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

diff --git a/doc/build/changelog/unreleased_14/8804.rst b/doc/build/changelog/unreleased_14/8804.rst
new file mode 100644 (file)
index 0000000..c3f91a1
--- /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 4e5808f623cd38b455c1b35af2ff381da151be05..11554fcc0cd15d93a458867f6fca50360823e4a2 100644 (file)
--- a/lib/sqlalchemy/dialects/sqlite/base.py
+++ b/lib/sqlalchemy/dialects/sqlite/base.py
@@ -893,6 +893,7 @@ from .json import JSONPathType
 from ... import exc
 from ... import schema as sa_schema
 from ... import sql
+from ... import text
 from ... import types as sqltypes
 from ... import util
 from ...engine import default
@@ -2626,6 +2627,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.
@@ -2634,7 +2650,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):
@@ -2652,6 +2699,7 @@ class SQLiteDialect(default.DefaultDialect):
                     break
                 else:
                     idx["column_names"].append(row[2])
+
         indexes.sort(key=lambda d: d["name"] or "~")  # sort None as last
         if indexes:
             return indexes

diff --git a/lib/sqlalchemy/testing/suite/test_reflection.py b/lib/sqlalchemy/testing/suite/test_reflection.py
index 563f129abad15fa4f136f95c0b9ccdf194cc4625..60f8fabf17b051faff42cabce84b06d3d6d4fbc0 100644 (file)
--- a/lib/sqlalchemy/testing/suite/test_reflection.py
+++ b/lib/sqlalchemy/testing/suite/test_reflection.py
@@ -2393,7 +2393,12 @@ class ComponentReflectionTestExtra(ComparesIndexes, 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": {},
+            }
         ]
 
         def completeIndex(entry):
@@ -2402,6 +2407,8 @@ class ComponentReflectionTestExtra(ComparesIndexes, fixtures.TestBase):
                 entry["dialect_options"] = {
                     f"{connection.engine.name}_include": []
                 }
+            else:
+                entry.setdefault("dialect_options", {})
 
         completeIndex(expected[0])
 

diff --git a/test/dialect/test_sqlite.py b/test/dialect/test_sqlite.py
index 5bda6577fe9dfd6254c28789cf3798cdaab3f7df..d460ef64e05e26f52764d7c4467cc2f59033d51e 100644 (file)
--- a/test/dialect/test_sqlite.py
+++ b/test/dialect/test_sqlite.py
@@ -2280,6 +2280,7 @@ class ConstraintReflectionTest(fixtures.TestBase):
                     "unique": 1,
                     "name": "sqlite_autoindex_o_1",
                     "column_names": ["foo"],
+                    "dialect_options": {},
                 }
             ],
         )
@@ -2295,10 +2296,60 @@ class ConstraintReflectionTest(fixtures.TestBase):
                     "unique": 0,
                     "name": "ix_main_l_bar",
                     "column_names": ["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_(