Support for get_pk_constraint() method (#265)

* get_pk_constraint() implemented
xzkostyan · Oct 8, 2023 · 1b10fe0 · 1b10fe0
1 parent ddf9531
commit 1b10fe0
Show file tree

Hide file tree

Showing 3 changed files with 48 additions and 3 deletions.
diff --git a/clickhouse_sqlalchemy/drivers/base.py b/clickhouse_sqlalchemy/drivers/base.py
@@ -360,8 +360,26 @@ def get_foreign_keys(self, connection, table_name, schema=None, **kw):
 
     @reflection.cache
     def get_pk_constraint(self, connection, table_name, schema=None, **kw):
-        # No support for primary keys.
-        return []
+        if not self.supports_engine_reflection:
+            return {}
+
+        if schema:
+            query = (("SELECT primary_key FROM system.tables "
+                     "WHERE database='{}' AND name='{}'")
+                     .format(schema, table_name))
+        else:
+            query = (
+                "SELECT primary_key FROM system.tables WHERE name='{}'"
+            ).format(table_name)
+
+        rows = self._execute(connection, query)
+        for r in rows:
+            primary_keys = r.primary_key
+            if primary_keys:
+                return {
+                    "constrained_columns": tuple(primary_keys.split(", ")),
+                }
+        return {}
 
     @reflection.cache
     def get_indexes(self, connection, table_name, schema=None, **kw):

diff --git a/clickhouse_sqlalchemy/drivers/compilers/ddlcompiler.py b/clickhouse_sqlalchemy/drivers/compilers/ddlcompiler.py
@@ -129,6 +129,14 @@ def visit_merge_tree(self, engine):
                 )
             )
         if engine.primary_key:
+            if not engine.order_by:
+                text += ' ORDER BY {0}\n'.format(
+                    self._compile_param(
+                        engine.primary_key.get_expressions_or_columns(),
+                        opt_list=True
+                    )
+                )
+
             text += ' PRIMARY KEY {0}\n'.format(
                 self._compile_param(
                     engine.primary_key.get_expressions_or_columns(),

diff --git a/tests/drivers/test_clickhouse_dialect.py b/tests/drivers/test_clickhouse_dialect.py
@@ -2,6 +2,7 @@
 from sqlalchemy.ext.asyncio import create_async_engine
 
 from clickhouse_sqlalchemy import make_session, engines, types, Table
+from clickhouse_sqlalchemy.drivers.base import ClickHouseDialect
 from tests.testcase import BaseTestCase, BaseAsynchTestCase
 from tests.config import (
     system_http_uri, system_native_uri, system_asynch_uri,
@@ -16,7 +17,7 @@
 class ClickHouseDialectTestCase(BaseTestCase):
 
     @property
-    def dialect(self):
+    def dialect(self) -> ClickHouseDialect:
         return self.session.bind.dialect
 
     @property
@@ -31,6 +32,14 @@ def setUp(self):
             Column('x', types.Int32, primary_key=True),
             engines.Memory()
         )
+        self.table_with_primary = Table(
+            'test_primary_table',
+            self.metadata(),
+            Column('x', types.Int32, primary_key=True),
+            engines.MergeTree(
+                primary_key=['x']
+            )
+        )
         self.table.drop(self.session.bind, if_exists=True)
 
     def tearDown(self):
@@ -108,6 +117,16 @@ def test_columns_compilation(self):
         col = Column('x', types.Nullable(types.Int32))
         self.assertEqual(str(col.type), 'Nullable(Int32)')
 
+    @require_server_version(18, 16, 0)
+    def test_primary_keys(self):
+        self.table_with_primary.create(self.session.bind)
+        con = self.dialect.get_pk_constraint(self.session,
+                                             self.table_with_primary.name)
+        self.assertIsNotNone(con,
+                             "Table should contain primary key constrain")
+        self.assertTrue("constrained_columns" in con)
+        self.assertEqual(("x",), con["constrained_columns"])
+
     @require_server_version(19, 16, 2)
     def test_empty_set_expr(self):
         numbers = Table(