Make index hinting work with CockroachDB (#178)

This includes support for asyncpg and testing.

Signed-off-by: Jonathan Dieter <[email protected]>
Co-authored-by: Gord Thompson <[email protected]>

Author: jdieter

CHANGES.md

@@ -2,9 +2,10 @@
 Unreleased
 
 - Added `include_hidden` option to `get_columns()` to enable reflection of columns like "rowid". (#173)
+- Added support for `.with_hint()` (patch courtesy of Jonathan Dieter)
 
 # Version 1.4.3
-Released Devember 10, 2021
+Released December 10, 2021
 
 - Added preliminary support for asyncpg. See instructions in the README.
 

sqlalchemy_cockroachdb/asyncpg.py

@@ -1,13 +1,15 @@
 from sqlalchemy.dialects.postgresql.asyncpg import PGDialect_asyncpg
 from .base import CockroachDBDialect
 from .ddl_compiler import CockroachDDLCompiler
+from .stmt_compiler import CockroachCompiler
 from .stmt_compiler import CockroachIdentifierPreparer
 
 
 class CockroachDBDialect_asyncpg(PGDialect_asyncpg, CockroachDBDialect):
     driver = "asyncpg"  # driver name
     preparer = CockroachIdentifierPreparer
     ddl_compiler = CockroachDDLCompiler
+    statement_compiler = CockroachCompiler
 
     supports_statement_cache = True
 

sqlalchemy_cockroachdb/psycopg2.py

@@ -1,12 +1,14 @@
 from sqlalchemy.dialects.postgresql.psycopg2 import PGDialect_psycopg2
 from .base import CockroachDBDialect
 from .ddl_compiler import CockroachDDLCompiler
+from .stmt_compiler import CockroachCompiler
 from .stmt_compiler import CockroachIdentifierPreparer
 
 
 class CockroachDBDialect_psycopg2(PGDialect_psycopg2, CockroachDBDialect):
     driver = "psycopg2"  # driver name
     preparer = CockroachIdentifierPreparer
     ddl_compiler = CockroachDDLCompiler
+    statement_compiler = CockroachCompiler
 
     supports_statement_cache = True

sqlalchemy_cockroachdb/stmt_compiler.py

@@ -1,5 +1,5 @@
+from sqlalchemy.dialects.postgresql.base import PGCompiler
 from sqlalchemy.dialects.postgresql.base import PGIdentifierPreparer
-from sqlalchemy.dialects.postgresql.psycopg2 import PGCompiler_psycopg2
 
 # This is extracted from CockroachDB's `sql.y`. Add keywords here if *NEW* reserved keywords
 # are added to sql.y. DO NOT DELETE keywords here, even if they are deleted from sql.y:
@@ -97,5 +97,6 @@ class CockroachIdentifierPreparer(PGIdentifierPreparer):
     reserved_words = CRDB_RESERVED_WORDS
 
 
-class CockroachCompiler(PGCompiler_psycopg2):
-    pass
+class CockroachCompiler(PGCompiler):
+    def format_from_hint_text(self, sqltext, table, hint, iscrud):
+        return f"{sqltext}@{hint}"

test/test_with_hint.py

@@ -0,0 +1,32 @@
+from sqlalchemy import Column
+from sqlalchemy import Index
+from sqlalchemy import Integer
+from sqlalchemy import select
+from sqlalchemy import String
+from sqlalchemy import Table
+from sqlalchemy.testing import AssertsCompiledSQL
+from sqlalchemy.testing import fixtures
+from sqlalchemy.testing import provide_metadata
+
+
+class WithHintTest(fixtures.TestBase, AssertsCompiledSQL):
+    # __requires__ = ("sync_driver",)
+
+    @provide_metadata
+    def test_with_hint(self):
+        meta = self.metadata
+        t = Table(
+            "t",
+            meta,
+            Column("id", Integer),
+            Column("txt", String(50)),
+            Index("ix_t_txt", "txt"),
+        )
+        self.assert_compile(
+            select(t).with_hint(t, "ix_t_txt"),
+            "SELECT t.id, t.txt FROM t@ix_t_txt",
+        )
+        self.assert_compile(
+            select(t).with_hint(t, "ix_t_txt").where(t.c.id < 3),
+            "SELECT t.id, t.txt FROM t@ix_t_txt WHERE t.id < %(id_1)s",
+        )