Add support for Migrate::*

This satisfies SDK tester's `schema_migrations_input_dml.json`.

Part 2: Make it work.

Author: amotl

src/cratedb_fivetran_destination/main.py

@@ -36,7 +36,6 @@ def __init__(self):
         self.metadata = sa.MetaData()
         self.engine: sa.Engine = None
         self.processor: Processor = None
-        self.migration_helper = SchemaMigrationHelper(CrateDBDestinationImpl.table_map)
 
     def ConfigurationForm(self, request, context):
         log_message(LOG_INFO, "Fetching configuration form")
@@ -246,9 +245,10 @@ def DescribeTable(self, request, context):
         """
         self._configure_database(request.configuration.get("url"))
         table_name = self.table_name(request)
+        schema_name = self.schema_name(request)
         table: common_pb2.Table = common_pb2.Table(name=table_name)
         try:
-            sa_table = self._reflect_table(schema=request.schema_name, table=table_name)
+            sa_table = self._reflect_table(schema=schema_name, table=table_name)
         except sa.exc.NoSuchTableError:
             msg = f"Table not found: {table_name}"
             log_message(LOG_WARNING, f"DescribeTable: {msg}")
@@ -270,7 +270,7 @@ def Migrate(self, request, context):
         """
         Example implementation of the new Migrate RPC introduced for schema migration support.
         This method inspects which migration operation (oneof) was requested and logs / handles it.
-        For demonstration, all recognized operations return success.
+        For demonstration, all recognized operations return `success`.
 
         :param request: The migration request contains details of the operation.
         :param context: gRPC context
@@ -279,34 +279,37 @@ def Migrate(self, request, context):
               rather different migration methods are just manipulating table_map to simulate
               the migration operations.
         """
+        self._configure_database(request.configuration.get("url"))
+        migration_helper = SchemaMigrationHelper(self.engine, CrateDBDestinationImpl.table_map)
+
         details = request.details
         schema = details.schema
         table = details.table
 
         operation_case = details.WhichOneof("operation")
         log_message(LOG_INFO, f"[Migrate] schema={schema} table={table} operation={operation_case}")
 
-        response = None
+        table_obj = self.DescribeTable(request, context).table
 
         if operation_case == "drop":
-            response = self.migration_helper.handle_drop(details.drop, schema, table)
+            response = migration_helper.handle_drop(details.drop, schema, table)
 
         elif operation_case == "copy":
-            response = self.migration_helper.handle_copy(details.copy, schema, table)
+            response = migration_helper.handle_copy(details.copy, schema, table, table_obj)
 
         elif operation_case == "rename":
-            response = self.migration_helper.handle_rename(details.rename, schema, table)
+            response = migration_helper.handle_rename(details.rename, schema, table)
 
         elif operation_case == "add":
-            response = self.migration_helper.handle_add(details.add, schema, table)
+            response = migration_helper.handle_add(details.add, schema, table, table_obj)
 
         elif operation_case == "update_column_value":
-            response = self.migration_helper.handle_update_column_value(
+            response = migration_helper.handle_update_column_value(
                 details.update_column_value, schema, table
             )
 
         elif operation_case == "table_sync_mode_migration":
-            response = self.migration_helper.handle_table_sync_mode_migration(
+            response = migration_helper.handle_table_sync_mode_migration(
                 details.table_sync_mode_migration, schema, table
             )
 
@@ -368,11 +371,22 @@ def _table_info_from_fivetran(self, table: common_pb2.Table) -> TableInfo:  # pr
         primary_keys = [column.name for column in table.columns if column.primary_key]
         return TableInfo(fullname=table.name, primary_keys=primary_keys)
 
+    @staticmethod
+    def schema_name(request):
+        """
+        Return schema name from request object.
+        """
+        if hasattr(request, "details"):
+            return request.details.schema
+        return request.schema_name
+
     @staticmethod
     def table_name(request):
         """
         Return table name from request object.
         """
+        if hasattr(request, "details"):
+            return request.details.table
         if hasattr(request, "table"):
             return request.table.name
         return request.table_name

src/cratedb_fivetran_destination/schema_migration_helper.py

@@ -1,7 +1,11 @@
-# ruff: noqa: E501
+# ruff: noqa: E501,S608
 # https://github.com/fivetran/fivetran_partner_sdk/blob/main/examples/destination_connector/python/schema_migration_helper.py
 import sys
 
+import sqlalchemy as sa
+
+from cratedb_fivetran_destination.model import SqlBag, TypeMap
+
 sys.path.append("sdk_pb2")
 
 from fivetran_sdk import common_pb2, destination_sdk_pb2
@@ -18,16 +22,18 @@
 class SchemaMigrationHelper:
     """Helper class for handling migration operations"""
 
-    def __init__(self, table_map):
+    def __init__(self, engine: sa.Engine, table_map):
+        self.engine = engine
         self.table_map = table_map
 
     def handle_drop(self, drop_op, schema, table):
         """Handles drop operations (drop table, drop column in history mode)."""
         entity_case = drop_op.WhichOneof("entity")
 
         if entity_case == "drop_table":
-            # table-map manipulation to simulate drop, replace with actual logic.
-            self.table_map.pop(table, None)
+            sql = f'DROP TABLE "{schema}"."{table}"'
+            with self.engine.connect() as conn:
+                conn.execute(sa.text(sql))
 
             log_message(INFO, f"[Migrate:Drop] Dropping table {schema}.{table}")
             return destination_sdk_pb2.MigrateResponse(success=True)
@@ -53,15 +59,18 @@ def handle_drop(self, drop_op, schema, table):
         log_message(WARNING, "[Migrate:Drop] No drop entity specified")
         return destination_sdk_pb2.MigrateResponse(unsupported=True)
 
-    def handle_copy(self, copy_op, schema, table):
+    def handle_copy(self, copy_op, schema, table, table_obj: common_pb2.Table):
         """Handles copy operations (copy table, copy column, copy table to history mode)."""
         entity_case = copy_op.WhichOneof("entity")
 
         if entity_case == "copy_table":
-            # table-map manipulation to simulate copy, replace with actual logic.
             copy_table = copy_op.copy_table
-            if copy_table.from_table in self.table_map:
-                self.table_map[copy_table.to_table] = self.table_map[copy_table.from_table]
+            sql = (
+                f'CREATE TABLE "{schema}"."{copy_table.to_table}" '
+                f'AS SELECT * FROM "{schema}"."{copy_table.from_table}"'
+            )
+            with self.engine.connect() as conn:
+                conn.execute(sa.text(sql))
 
             log_message(
                 INFO,
@@ -70,17 +79,25 @@ def handle_copy(self, copy_op, schema, table):
             return destination_sdk_pb2.MigrateResponse(success=True)
 
         if entity_case == "copy_column":
-            # table-map manipulation to simulate copy column, replace with actual logic.
+            sql_bag = SqlBag()
             copy_column = copy_op.copy_column
-            table_obj = self.table_map.get(table)
-            if table_obj:
-                for col in table_obj.columns:
-                    if col.name == copy_column.from_column:
-                        new_col = type(col)()
-                        new_col.CopyFrom(col)
-                        new_col.name = copy_column.to_column
-                        table_obj.columns.add().CopyFrom(new_col)
-                        break
+            for col in table_obj.columns:
+                if col.name == copy_column.from_column:
+                    new_col = type(col)()
+                    new_col.CopyFrom(col)
+                    new_col.name = copy_column.to_column
+                    table_obj.columns.add().CopyFrom(new_col)
+                    type_ = TypeMap.to_cratedb(new_col.type, new_col.params)
+                    sql_bag.add(
+                        f'ALTER TABLE "{schema}"."{table}" ADD COLUMN "{new_col.name}" {type_};'
+                    )
+                    sql_bag.add(f'UPDATE "{schema}"."{table}" SET "{new_col.name}"="{col.name}";')
+                    break
+
+            if sql_bag:
+                with self.engine.connect() as conn:
+                    for command in sql_bag.statements:
+                        conn.execute(sa.text(command))
 
             log_message(
                 INFO,
@@ -116,31 +133,21 @@ def handle_rename(self, rename_op, schema, table):
         entity_case = rename_op.WhichOneof("entity")
 
         if entity_case == "rename_table":
-            # table-map manipulation to simulate rename, replace with actual logic.
             rt = rename_op.rename_table
-            if rt.from_table in self.table_map:
-                tbl = self.table_map.pop(rt.from_table)
-                # Adjust name inside the Table metadata if needed
-                tbl_copy = tbl.__class__.FromString(tbl.SerializeToString())
-                if hasattr(tbl_copy, "name"):
-                    tbl_copy.name = rt.to_table
-                self.table_map[rt.to_table] = tbl_copy
+            sql = f'ALTER TABLE "{schema}"."{rt.from_table}" RENAME TO "{rt.to_table}";'
+            with self.engine.connect() as conn:
+                conn.execute(sa.text(sql))
 
             log_message(
                 INFO, f"[Migrate:RenameTable] from={rt.from_table} to={rt.to_table} schema={schema}"
             )
             return destination_sdk_pb2.MigrateResponse(success=True)
 
         if entity_case == "rename_column":
-            # table-map manipulation to simulate rename column, replace with actual logic.
             rename_column = rename_op.rename_column
-            table_obj = self.table_map.get(table)
-            if table_obj:
-                # Rename the column
-                for col in table_obj.columns:
-                    if col.name == rename_column.from_column:
-                        col.name = rename_column.to_column
-                        break
+            sql = f'ALTER TABLE "{schema}"."{table}" RENAME "{rename_column.from_column}" TO "{rename_column.to_column}";'
+            with self.engine.connect() as conn:
+                conn.execute(sa.text(sql))
 
             log_message(
                 INFO,
@@ -151,7 +158,7 @@ def handle_rename(self, rename_op, schema, table):
         log_message(WARNING, "[Migrate:Rename] No rename entity specified")
         return destination_sdk_pb2.MigrateResponse(unsupported=True)
 
-    def handle_add(self, add_op, schema, table):
+    def handle_add(self, add_op, schema, table, table_obj: common_pb2.Table):
         """Handles add operations (add column in history mode, add column with default value)."""
         entity_case = add_op.WhichOneof("entity")
 
@@ -171,13 +178,15 @@ def handle_add(self, add_op, schema, table):
             return destination_sdk_pb2.MigrateResponse(success=True)
 
         if entity_case == "add_column_with_default_value":
-            # table-map manipulation to simulate add column with default value, replace with actual logic.
             add_col_default_with_value = add_op.add_column_with_default_value
-            table_obj = self.table_map.get(table)
             if table_obj:
                 new_col = table_obj.columns.add()
                 new_col.name = add_col_default_with_value.column
                 new_col.type = add_col_default_with_value.column_type
+                type_ = TypeMap.to_cratedb(new_col.type, new_col.params)
+                sql = f'ALTER TABLE "{schema}"."{table}" ADD COLUMN "{new_col.name}" {type_};'
+                with self.engine.connect() as conn:
+                    conn.execute(sa.text(sql))
 
             log_message(
                 INFO,
@@ -190,7 +199,12 @@ def handle_add(self, add_op, schema, table):
 
     def handle_update_column_value(self, upd, schema, table):
         """Handles update column value operation."""
-        # Placeholder: Update all existing rows' column value.
+        with self.engine.connect() as conn:
+            conn.execute(
+                sa.text(f'UPDATE "{schema}"."{table}" SET "{upd.column}"=:value;'),  # noqa: S608
+                parameters={"value": upd.value},
+            )
+            conn.execute(sa.text(f'REFRESH TABLE "{schema}"."{table}";'))
 
         log_message(
             INFO,

tests/conftest.py

@@ -8,7 +8,8 @@ def engine():
     with engine.connect() as conn:
         # Clean up stale "read-only" mode states.
         try:
-            conn.execute(sa.text('ALTER TABLE testdrive.foo RESET ("blocks.write");'))
+            conn.execute(sa.text('ALTER TABLE testdrive.foo SET ("blocks.write"=false);'))
+            conn.execute(sa.text('ALTER TABLE tester.transaction SET ("blocks.write"=false);'))
         except Exception:  # noqa: S110
             pass
         conn.execute(sa.text("DROP TABLE IF EXISTS testdrive.foo"))

tests/data/fivetran_migrations_dml/configuration.json

@@ -0,0 +1,11 @@
+{
+  "writerType": "Database",
+  "url": "crate://",
+  "database": "doc",
+  "table": "testdrive",
+  "host": "localhost",
+  "port": "4200",
+  "user": "crate",
+  "password": "",
+  "enableEncryption": "false"
+}

tests/data/fivetran_migrations_dml/schema_migrations_input_dml.json

@@ -0,0 +1,100 @@
+{
+   "create_table" : {
+      "transaction": {
+         "columns": {
+            "id": "INT",
+            "amount" : "DOUBLE",
+            "desc": "STRING"
+         },
+         "primary_key": ["id"]
+      },
+      "transaction_drop": {
+         "columns": {
+            "id": "INT",
+            "amount" : "DOUBLE",
+            "desc": "STRING"
+         },
+         "primary_key": ["id"]
+      }
+   },
+   "ops" : [
+      {
+         "upsert": {
+            "transaction": [
+               {"id":1, "amount": 100.45, "desc":null},
+               {"id":2, "amount": 150.33, "desc": "two"},
+               {"id":3, "amount": 150.33, "desc": "two"},
+               {"id":4, "amount": 150.33, "desc": "two"},
+               {"id":10, "amount": 200, "desc": "three"},
+               {"id":20, "amount": 50, "desc": "money"}
+            ],
+            "transaction_drop": [
+               {"id":1, "amount": 100.45, "desc":null},
+               {"id":2, "amount": 150.33, "desc": "two"}
+            ]
+         }
+      }
+   ],
+   "schema_migration" : [
+      {
+         "copy_column": [
+            {
+               "table": "transaction",
+               "from_column": "desc",
+               "to_column": "desc_detailed"
+            }
+         ],
+         "update_column_value": [
+            {
+               "table": "transaction",
+               "column": "amount",
+               "value": "202.57"
+            }
+         ],
+         "add_column_with_default_value": [
+            {
+               "table": "transaction",
+               "column": "operation_time",
+               "data_type": "UTC_DATETIME",
+               "default_value": "2005-05-23T20:57:00Z"
+            }
+         ],
+         "set_column_to_null": [
+            {
+               "table": "transaction",
+               "column": "desc"
+            }
+         ],
+         "copy_table": [
+            {
+               "from_table": "transaction",
+               "to_table": "transaction_new"
+            }
+         ],
+         "rename_column": [
+            {
+               "table": "transaction",
+               "from_column": "amount",
+               "to_column": "amount_renamed"
+            }
+         ],
+         "rename_table": [
+            {
+               "from_table": "transaction_new",
+               "to_table": "transaction_renamed"
+            }
+         ],
+         "drop_table": [
+            {
+               "table": "transaction_drop"
+            }
+         ]
+      }
+   ],
+   "describe_table" : [
+      "transaction",
+      "transaction_drop",
+      "transaction_new",
+      "transaction_renamed"
+   ]
+}