Improve SQLAlchemy's CrateDialect.get_pk_constraint for CrateDB>=5.1.0

Since CrateDB 3.0.0 already, it would have been correct to use the `information_schema.key_column_usage.table_schema` column instead of
`table_catalog` (mostly contains value `doc`, the default schema of
CrateDB). Now, with CrateDB >=5.1.0, `table_catalog` will always return
`crate`, thus the need for a fix.

Author: amotl

CHANGES.txt

@@ -5,6 +5,8 @@ Changes for crate
 Unreleased
 ==========
 
+- Improved SQLAlchemy's ``CrateDialect.get_pk_constraint`` to be compatible
+  with breaking changes in CrateDB >=5.1.0.
 
 2022/07/04 0.27.1
 =================

src/crate/client/sqlalchemy/dialect.py

@@ -264,14 +264,24 @@ def get_columns(self, connection, table_name, schema=None, **kw):
 
     @reflection.cache
     def get_pk_constraint(self, engine, table_name, schema=None, **kw):
-        if self.server_version_info >= (2, 3, 0):
+        if self.server_version_info >= (3, 0, 0):
+            query = """SELECT column_name
+                    FROM information_schema.key_column_usage
+                    WHERE table_name = ? AND table_schema = ?"""
+
+            def result_fun(result):
+                rows = result.fetchall()
+                return set(map(lambda el: el[0], rows))
+
+        elif self.server_version_info >= (2, 3, 0):
             query = """SELECT column_name
                     FROM information_schema.key_column_usage
                     WHERE table_name = ? AND table_catalog = ?"""
 
             def result_fun(result):
                 rows = result.fetchall()
                 return set(map(lambda el: el[0], rows))
+
         else:
             query = """SELECT constraint_name
                    FROM information_schema.table_constraints

src/crate/client/sqlalchemy/tests/dialect_test.py

@@ -67,7 +67,7 @@ class Character(self.base):
         self.character = Character
         self.session = Session()
 
-    def test_primary_keys(self):
+    def test_primary_keys_2_3_0(self):
         meta = self.character.metadata
         insp = inspect(meta.bind)
         self.engine.dialect.server_version_info = (2, 3, 0)
@@ -81,6 +81,23 @@ def test_primary_keys(self):
         eq_(insp.get_pk_constraint("characters")['constrained_columns'], {"id", "id2", "id3"})
         self.fake_cursor.fetchall.assert_called_once_with()
         in_("information_schema.key_column_usage", self.executed_statement)
+        in_("table_catalog = ?", self.executed_statement)
+
+    def test_primary_keys_3_0_0(self):
+        meta = self.character.metadata
+        insp = inspect(meta.bind)
+        self.engine.dialect.server_version_info = (3, 0, 0)
+
+        self.fake_cursor.rowcount = 3
+        self.fake_cursor.description = (
+            ('foo', None, None, None, None, None, None),
+        )
+        self.fake_cursor.fetchall = MagicMock(return_value=[["id"], ["id2"], ["id3"]])
+
+        eq_(insp.get_pk_constraint("characters")['constrained_columns'], {"id", "id2", "id3"})
+        self.fake_cursor.fetchall.assert_called_once_with()
+        in_("information_schema.key_column_usage", self.executed_statement)
+        in_("table_schema = ?", self.executed_statement)
 
     def test_get_table_names(self):
         self.fake_cursor.rowcount = 1