Support for Starrocks

[maver1ck]

Hi,
I'm working on support for Starrocks dialect in Alembic.

The problem is that Starrocks support updates only on non-primary key columns in table that has different primary key.

As a result I want to create alembic_version table with following structure:

CREATE TABLE `alembic_version` (
  `id` bigint(20) NOT NULL AUTO_INCREMENT COMMENT "",
  `version_num` varchar(32) NULL COMMENT ""
) ENGINE=OLAP
PRIMARY KEY(`id`)
DISTRIBUTED BY HASH(`id`);

I'd like to create similar patch to Clickhouse dialect.
Source:
https://github.com/xzkostyan/clickhouse-sqlalchemy/blob/master/clickhouse_sqlalchemy/alembic/dialect.py#L39

My version:

def patch_alembic_version(context, **kwargs):
    migration_context = context._proxy._migration_context
    old_version = migration_context._version
    metadata = old_version.metadata
    metadata.remove(old_version)

    version = Table(
        "alembic_version",
        metadata,
        Column("id", types.BIGINT, autoincrement=True, primary_key=True),
        Column("version_num", types.VARCHAR(32), primary_key=False),
        starrocks_primary_key="id",
    )
    migration_context._version = version

The main problem is that alembic assumes that first column in alembic_version table is version_num.

As a result of this analysis I created PR to alembic repo that will select only version_num column even if alembic_version table has more than one column.
#1561
Unfortunally it was discarded.

How to proceed with this problem ?

[maver1ck]

Generally speaking Alembic should have plugable way of defining alembic_version table.
Currently multiple dialects (like clickhouse) are just patching private fields which should be avoided.

[zzzeek]

@maver1ck please start with this patch

diff --git a/alembic/ddl/impl.py b/alembic/ddl/impl.py
index 2574688..4618024 100644
--- a/alembic/ddl/impl.py
+++ b/alembic/ddl/impl.py
@@ -21,7 +21,12 @@ from typing import TYPE_CHECKING
 from typing import Union
 
 from sqlalchemy import cast
+from sqlalchemy import Column
+from sqlalchemy import MetaData
+from sqlalchemy import PrimaryKeyConstraint
 from sqlalchemy import schema
+from sqlalchemy import String
+from sqlalchemy import Table
 from sqlalchemy import text
 
 from . import _autogen
@@ -136,6 +141,30 @@ class DefaultImpl(metaclass=ImplMeta):
         self.output_buffer.write(text + "\n\n")
         self.output_buffer.flush()
 
+    def version_table_impl(self, version_table, version_table_schema, version_table_pk):
+        """create the Table object for the version_table.
+
+        Provided as part of impl so that third party dialects can override
+        this.
+
+        .. versionadded:: 1.13.4
+
+        """
+        vt = Table(
+            version_table,
+            MetaData(),
+            Column("version_num", String(32), nullable=False),
+            schema=version_table_schema,
+        )
+        if version_table_pk:
+            vt.append_constraint(
+                PrimaryKeyConstraint(
+                    "version_num", name=f"{version_table}_pkc"
+                )
+            )
+
+        return vt
+
     def requires_recreate_in_batch(
         self, batch_op: BatchOperationsImpl
     ) -> bool:
diff --git a/alembic/runtime/migration.py b/alembic/runtime/migration.py
index 6cfe5e2..f652808 100644
--- a/alembic/runtime/migration.py
+++ b/alembic/runtime/migration.py
@@ -24,10 +24,7 @@ from typing import Union
 
 from sqlalchemy import Column
 from sqlalchemy import literal_column
-from sqlalchemy import MetaData
-from sqlalchemy import PrimaryKeyConstraint
-from sqlalchemy import String
-from sqlalchemy import Table
+from sqlalchemy import select
 from sqlalchemy.engine import Engine
 from sqlalchemy.engine import url as sqla_url
 from sqlalchemy.engine.strategies import MockEngineStrategy
@@ -190,18 +187,6 @@ class MigrationContext:
         self.version_table_schema = version_table_schema = opts.get(
             "version_table_schema", None
         )
-        self._version = Table(
-            version_table,
-            MetaData(),
-            Column("version_num", String(32), nullable=False),
-            schema=version_table_schema,
-        )
-        if opts.get("version_table_pk", True):
-            self._version.append_constraint(
-                PrimaryKeyConstraint(
-                    "version_num", name="%s_pkc" % version_table
-                )
-            )
 
         self._start_from_rev: Optional[str] = opts.get("starting_rev")
         self.impl = ddl.DefaultImpl.get_by_dialect(dialect)(
@@ -212,6 +197,12 @@ class MigrationContext:
             self.output_buffer,
             opts,
         )
+
+        self._version = self.impl.version_table_impl(
+            version_table, version_table_schema,
+            opts.get("version_table_pk", True)
+        )
+
         log.info("Context impl %s.", self.impl.__class__.__name__)
         if self.as_sql:
             log.info("Generating static SQL")
@@ -540,7 +531,7 @@ class MigrationContext:
                 return ()
         assert self.connection is not None
         return tuple(
-            row[0] for row in self.connection.execute(self._version.select())
+            row[0] for row in self.connection.execute(select(self._version.c.version_num))
         )
 
     def _ensure_version_table(self, purge: bool = False) -> None:

tests for this need to be in tests/test_version_table.py. use a custom impl class which overrides the method. we can help you write the tests if you dont know how

[CaselIT]

ok, that use case is likely better. That should cover also the clickhouse monkey patch. I wonder why they didn't ask for this feature upstream...

[CaselIT]

Hi @xzkostyan sorry for the ping. Would this patch #1562 (reply in thread) work for you to avoid monkey patching the version table as done here? https://github.com/xzkostyan/clickhouse-sqlalchemy/blob/f52087844b86a9524c2295185ac29274b293d129/clickhouse_sqlalchemy/alembic/dialect.py#L39-L58

[maver1ck]

@zzzeek
Thanks. I will prepare PR tomorrow.

[maver1ck]

Please find my PR here:
#1563