diff --git a/clickhouse_sqlalchemy/drivers/base.py b/clickhouse_sqlalchemy/drivers/base.py index b5d23fa2..7d9661cf 100644 --- 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 index 7da947be..8af8907a 100644 --- 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 index d310aabf..c8139bb4 100644 --- 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(