From 5cb7bf8137becbc40d12ffd4dcaf49d23913a926 Mon Sep 17 00:00:00 2001 From: Tobias Pfeiffer Date: Mon, 28 Nov 2022 07:52:31 -0500 Subject: [PATCH] add partial index predicate to SQLiteDialect.get_indexes() result Added support for reflection of expression-oriented WHERE criteria included in indexes on the SQLite dialect, in a manner similar to that of the PostgreSQL dialect. Pull request courtesy Tobias Pfeiffer. Fixes: #8804 Closes: #8806 Pull-request: https://github.com/sqlalchemy/sqlalchemy/pull/8806 Pull-request-sha: 539dfcb372360911b69aed2a804698bb1a2220b1 Change-Id: I0e34d47dbe2b9c1da6fce531363084843e5127a3 (cherry picked from commit ed39e846cd8ae2714c47fc3d563582f72483df0c) --- doc/build/changelog/unreleased_14/8804.rst | 7 +++ lib/sqlalchemy/dialects/sqlite/base.py | 51 ++++++++++++++++++- .../testing/suite/test_reflection.py | 13 +++-- test/dialect/test_sqlite.py | 51 +++++++++++++++++++ 4 files changed, 116 insertions(+), 6 deletions(-) create mode 100644 doc/build/changelog/unreleased_14/8804.rst diff --git a/doc/build/changelog/unreleased_14/8804.rst b/doc/build/changelog/unreleased_14/8804.rst new file mode 100644 index 0000000000..c3f91a16d2 --- /dev/null +++ b/doc/build/changelog/unreleased_14/8804.rst @@ -0,0 +1,7 @@ +.. change:: + :tags: usecase, sqlite + :tickets: 8804 + + Added support for reflection of expression-oriented WHERE criteria included + in indexes on the SQLite dialect, in a manner similar to that of the + PostgreSQL dialect. Pull request courtesy Tobias Pfeiffer. diff --git a/lib/sqlalchemy/dialects/sqlite/base.py b/lib/sqlalchemy/dialects/sqlite/base.py index 612d8f9063..24166717a4 100644 --- a/lib/sqlalchemy/dialects/sqlite/base.py +++ b/lib/sqlalchemy/dialects/sqlite/base.py @@ -821,6 +821,7 @@ from ... import exc from ... import processors from ... import schema as sa_schema from ... import sql +from ... import text from ... import types as sqltypes from ... import util from ...engine import default @@ -2474,6 +2475,21 @@ class SQLiteDialect(default.DefaultDialect): ) indexes = [] + # regular expression to extract the filter predicate of a partial + # index. this could fail to extract the predicate correctly on + # indexes created like + # CREATE INDEX i ON t (col || ') where') WHERE col <> '' + # but as this function does not support expression-based indexes + # this case does not occur. + partial_pred_re = re.compile(r"\)\s+where\s+(.+)", re.IGNORECASE) + + if schema: + schema_expr = "%s." % self.identifier_preparer.quote_identifier( + schema + ) + else: + schema_expr = "" + include_auto_indexes = kw.pop("include_auto_indexes", False) for row in pragma_indexes: # ignore implicit primary key index. @@ -2482,7 +2498,38 @@ class SQLiteDialect(default.DefaultDialect): "sqlite_autoindex" ): continue - indexes.append(dict(name=row[1], column_names=[], unique=row[2])) + indexes.append( + dict( + name=row[1], + column_names=[], + unique=row[2], + dialect_options={}, + ) + ) + + # check partial indexes + if row[4]: + s = ( + "SELECT sql FROM %(schema)ssqlite_master " + "WHERE name = ? " + "AND type = 'index'" % {"schema": schema_expr} + ) + rs = connection.exec_driver_sql(s, (row[1],)) + index_sql = rs.scalar() + predicate_match = partial_pred_re.search(index_sql) + if predicate_match is None: + # unless the regex is broken this case shouldn't happen + # because we know this is a partial index, so the + # definition sql should match the regex + util.warn( + "Failed to look up filter predicate of " + "partial index %s" % row[1] + ) + else: + predicate = predicate_match.group(1) + indexes[-1]["dialect_options"]["sqlite_where"] = text( + predicate + ) # loop thru unique indexes to get the column names. for idx in list(indexes): @@ -2500,6 +2547,8 @@ class SQLiteDialect(default.DefaultDialect): break else: idx["column_names"].append(row[2]) + + indexes.sort(key=lambda d: d["name"] or "~") # sort None as last return indexes @reflection.cache diff --git a/lib/sqlalchemy/testing/suite/test_reflection.py b/lib/sqlalchemy/testing/suite/test_reflection.py index 12949fe02b..4e575046d3 100644 --- a/lib/sqlalchemy/testing/suite/test_reflection.py +++ b/lib/sqlalchemy/testing/suite/test_reflection.py @@ -1300,8 +1300,14 @@ class ComponentReflectionTestExtra(fixtures.TestBase): insp = inspect(connection) expected = [ - {"name": "t_idx_2", "column_names": ["x"], "unique": False} + { + "name": "t_idx_2", + "column_names": ["x"], + "unique": False, + "dialect_options": {}, + } ] + if testing.requires.index_reflects_included_columns.enabled: expected[0]["include_columns"] = [] expected[0]["dialect_options"] = { @@ -1311,10 +1317,7 @@ class ComponentReflectionTestExtra(fixtures.TestBase): with expect_warnings( "Skipped unsupported reflection of expression-based index t_idx" ): - eq_( - insp.get_indexes("t"), - expected, - ) + eq_(insp.get_indexes("t"), expected) @testing.requires.index_reflects_included_columns def test_reflect_covering_index(self, metadata, connection): diff --git a/test/dialect/test_sqlite.py b/test/dialect/test_sqlite.py index 1f7a06dffb..3da4d6574b 100644 --- a/test/dialect/test_sqlite.py +++ b/test/dialect/test_sqlite.py @@ -2312,6 +2312,7 @@ class ConstraintReflectionTest(fixtures.TestBase): "unique": 1, "name": "sqlite_autoindex_o_1", "column_names": ["foo"], + "dialect_options": {}, } ], ) @@ -2327,10 +2328,60 @@ class ConstraintReflectionTest(fixtures.TestBase): "unique": 0, "name": u"ix_main_l_bar", "column_names": [u"bar"], + "dialect_options": {}, } ], ) + def test_reflect_partial_indexes(self, connection): + connection.exec_driver_sql( + "create table foo_with_partial_index (x integer, y integer)" + ) + connection.exec_driver_sql( + "create unique index ix_partial on " + "foo_with_partial_index (x) where y > 10" + ) + connection.exec_driver_sql( + "create unique index ix_no_partial on " + "foo_with_partial_index (x)" + ) + connection.exec_driver_sql( + "create unique index ix_partial2 on " + "foo_with_partial_index (x, y) where " + "y = 10 or abs(x) < 5" + ) + + inspector = inspect(connection) + indexes = inspector.get_indexes("foo_with_partial_index") + eq_( + indexes, + [ + { + "unique": 1, + "name": "ix_no_partial", + "column_names": ["x"], + "dialect_options": {}, + }, + { + "unique": 1, + "name": "ix_partial", + "column_names": ["x"], + "dialect_options": {"sqlite_where": mock.ANY}, + }, + { + "unique": 1, + "name": "ix_partial2", + "column_names": ["x", "y"], + "dialect_options": {"sqlite_where": mock.ANY}, + }, + ], + ) + eq_(indexes[1]["dialect_options"]["sqlite_where"].text, "y > 10") + eq_( + indexes[2]["dialect_options"]["sqlite_where"].text, + "y = 10 or abs(x) < 5", + ) + def test_unique_constraint_named(self): inspector = inspect(testing.db) eq_( -- 2.47.2