Lal/fix deadlock on rollout sdcosmo 3601

cosmotech/coal/postgresql/runner.py

@@ -21,9 +21,39 @@
 from cosmotech.coal.utils.logger import LOGGER
 
 
+def create_metadata(configuration: Configuration) -> None:
+    _psql = PostgresUtils(configuration)
+    with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
+        with conn.cursor() as curs:
+            schema_table = f"{str(_psql.db_schema)}.{str(_psql.metadata_table_name)}"
+            sql_create_table = f"""
+                CREATE TABLE IF NOT EXISTS {schema_table}  (
+                  id varchar(32),
+                  name varchar(256),
+                  last_csm_run_id varchar(32) PRIMARY KEY,
+                  run_template_id varchar(32)
+                );
+            """
+            LOGGER.info(T("coal.services.postgresql.creating_table").format(schema_table=schema_table))
+            curs.execute(sql_create_table)
+            conn.commit()
+
+
+def get_metadata_last_run(configuration: Configuration) -> list:
+    _psql = PostgresUtils(configuration)
+    with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
+        with conn.cursor() as curs:
+            metadata_table = f"{str(_psql.db_schema)}.{str(_psql.metadata_table_name)}"
+            sql_query_run = f"""
+                SELECT last_csm_run_id FROM {metadata_table} WHERE id = $1
+            """
+            curs.execute(sql_query_run, (configuration.cosmotech.runner_id,))
+            return [row[0] for row in curs.fetchall()]
+
+
 def send_runner_metadata_to_postgresql(
     configuration: Configuration,
-) -> str:
+) -> None:
     """
     Send runner metadata to a PostgreSQL database.
 
@@ -45,25 +75,6 @@ def send_runner_metadata_to_postgresql(
     with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
         with conn.cursor() as curs:
             schema_table = f"{str(_psql.db_schema)}.{str(_psql.metadata_table_name)}"
-            sql_create_table = f"""
-                CREATE TABLE IF NOT EXISTS {schema_table}  (
-                  id varchar(32) PRIMARY KEY,
-                  name varchar(256),
-                  last_csm_run_id varchar(32) UNIQUE,
-                  run_template_id varchar(32)
-                );
-            """
-            LOGGER.info(T("coal.services.postgresql.creating_table").format(schema_table=schema_table))
-            curs.execute(sql_create_table)
-            conn.commit()
-
-            runner_id = runner.get("id")
-            sql_delete_from_metatable = f"""
-                DELETE FROM {schema_table}
-                WHERE id= $1;
-            """
-            curs.execute(sql_delete_from_metatable, (runner_id,))
-            conn.commit()
 
             sql_upsert = f"""
                 INSERT INTO {schema_table} (id, name, last_csm_run_id, run_template_id)
@@ -73,15 +84,33 @@ def send_runner_metadata_to_postgresql(
             curs.execute(
                 sql_upsert,
                 (
-                    runner.get("id"),
+                    configuration.cosmotech.runner_id,
                     runner.get("name"),
-                    runner.get("lastRunInfo").get("lastRunId"),
-                    runner.get("runTemplateId"),
+                    configuration.cosmotech.run_id,
+                    configuration.cosmotech.run_template_id,
                 ),
             )
             conn.commit()
             LOGGER.info(T("coal.services.postgresql.metadata_updated"))
-    return runner.get("lastRunInfo").get("lastRunId")
+
+
+def remove_run_metadata_from_postgresql(
+    configuration: Configuration,
+    run_id: str,
+) -> str:
+    _psql = PostgresUtils(configuration)
+
+    # Connect to PostgreSQL and remove runner metadata row
+    with dbapi.connect(_psql.full_uri, autocommit=True) as conn:
+        with conn.cursor() as curs:
+            schema_table = f"{_psql.db_schema}.{_psql.metadata_table_name}"
+            sql_delete_from_metatable = f"""
+                DELETE FROM {schema_table}
+                WHERE last_csm_run_id = $1;
+            """
+            curs.execute(sql_delete_from_metatable, (run_id,))
+            conn.commit()
+    LOGGER.info(T("coal.services.postgresql.metadata_removed").format(id=run_id))
 
 
 def remove_runner_metadata_from_postgresql(
@@ -115,5 +144,5 @@ def remove_runner_metadata_from_postgresql(
             """
             curs.execute(sql_delete_from_metatable, (runner_id,))
             conn.commit()
-    LOGGER.info(T("coal.services.postgresql.metadata_removed").format(runner_id=runner_id))
+    LOGGER.info(T("coal.services.postgresql.metadata_removed").format(id=runner_id))
     return runner.get("lastRunInfo").get("lastRunId")

cosmotech/coal/postgresql/store.py

@@ -72,6 +72,23 @@ def dump_store_to_postgresql(
     dump_store_to_postgresql_from_conf(configuration=_c, replace=replace, selected_tables=selected_tables, fk_id=fk_id)
 
 
+def add_fk_constraints(configuration: Configuration) -> None:
+    """
+    Add constraints on the column 'csm_run_id' on every PSQL tables that are in the current store
+
+    Args:
+        configuration: coal Configuration
+    """
+    _s = Store(configuration=configuration)
+    _psql = PostgresUtils(configuration)
+
+    tables = list(_s.list_tables())
+    for table_name in tables:
+        target_table_name = f"{_psql.table_prefix}{table_name}"
+        metadata_table = f"{_psql.metadata_table_name}"
+        _psql.add_fk_constraint(target_table_name, "csm_run_id", metadata_table, "last_csm_run_id")
+
+
 def dump_store_to_postgresql_from_conf(
     configuration: Configuration,
     replace: bool = True,
@@ -90,53 +107,54 @@ def dump_store_to_postgresql_from_conf(
     _psql = PostgresUtils(configuration)
     _s = Store(configuration=configuration)
 
+    # Apply table filter
     tables = list(_s.list_tables())
     if selected_tables:
         tables = [t for t in tables if t in selected_tables]
-    if len(tables):
-        LOGGER.info(T("coal.services.database.sending_data").format(table=f"{_psql.db_name}.{_psql.db_schema}"))
-        total_rows = 0
-        _process_start = perf_counter()
-        for table_name in tables:
-            _s_time = perf_counter()
-            target_table_name = f"{_psql.table_prefix}{table_name}"
-            LOGGER.info(T("coal.services.database.table_entry").format(table=target_table_name))
-            data = _s.get_table(table_name)
-            if not len(data):
-                LOGGER.info(T("coal.services.database.no_rows"))
-                continue
-            if fk_id:
-                data = data.append_column("csm_run_id", [[fk_id] * data.num_rows])
-            _dl_time = perf_counter()
-            rows = _psql.send_pyarrow_table_to_postgresql(
-                data,
-                target_table_name,
-                replace,
-            )
-            if fk_id and _psql.is_metadata_exists():
-                metadata_table = f"{_psql.metadata_table_name}"
-                _psql.add_fk_constraint(target_table_name, "csm_run_id", metadata_table, "last_csm_run_id")
-
-            total_rows += rows
-            _up_time = perf_counter()
-            LOGGER.info(T("coal.services.database.row_count").format(count=rows))
-            LOGGER.debug(
-                T("coal.common.timing.operation_completed").format(
-                    operation="Load from datastore", time=f"{_dl_time - _s_time:0.3}"
-                )
-            )
-            LOGGER.debug(
-                T("coal.common.timing.operation_completed").format(
-                    operation="Send to postgresql", time=f"{_up_time - _dl_time:0.3}"
-                )
+
+    if not tables:
+        LOGGER.info(T("coal.services.database.store_empty"))
+
+    LOGGER.info(T("coal.services.database.sending_data").format(table=f"{_psql.db_name}.{_psql.db_schema}"))
+    total_rows = 0
+    _process_start = perf_counter()
+    for table_name in tables:
+        _s_time = perf_counter()
+        target_table_name = f"{_psql.table_prefix}{table_name}"
+        LOGGER.info(T("coal.services.database.table_entry").format(table=target_table_name))
+
+        data = _s.get_table(table_name)
+        if not len(data):
+            LOGGER.info(T("coal.services.database.no_rows"))
+            continue
+        if fk_id:
+            data = data.append_column("csm_run_id", [[fk_id] * data.num_rows])
+        _dl_time = perf_counter()
+
+        rows = _psql.send_pyarrow_table_to_postgresql(
+            data,
+            target_table_name,
+            replace,
+        )
+
+        total_rows += rows
+        _up_time = perf_counter()
+        LOGGER.info(T("coal.services.database.row_count").format(count=rows))
+        LOGGER.debug(
+            T("coal.common.timing.operation_completed").format(
+                operation="Load from datastore", time=f"{_dl_time - _s_time:0.3}"
             )
-        _process_end = perf_counter()
-        LOGGER.info(
-            T("coal.services.database.rows_fetched").format(
-                table="all tables",
-                count=total_rows,
-                time=f"{_process_end - _process_start:0.3}",
+        )
+        LOGGER.debug(
+            T("coal.common.timing.operation_completed").format(
+                operation="Send to postgresql", time=f"{_up_time - _dl_time:0.3}"
             )
         )
-    else:
-        LOGGER.info(T("coal.services.database.store_empty"))
+    _process_end = perf_counter()
+    LOGGER.info(
+        T("coal.services.database.rows_fetched").format(
+            table="all tables",
+            count=total_rows,
+            time=f"{_process_end - _process_start:0.3}",
+        )
+    )

cosmotech/coal/postgresql/utils.py

@@ -168,8 +168,13 @@ def add_fk_constraint(
                     REFERENCES {self.db_schema}.{to_table}({to_col})
                     ON DELETE CASCADE;
                 """
+                sql_add_fk_index = f"""
+                    CREATE INDEX CONCURRENTLY IF NOT EXISTS {from_table}_index
+                    on {self.db_schema}.{from_table}(csm_run_id)
+                """
                 curs.execute(sql_drop_fk)
                 curs.execute(sql_add_fk)
+                curs.execute(sql_add_fk_index)
                 conn.commit()
 
     def is_metadata_exists(self) -> None:

cosmotech/coal/store/output/postgres_channel.py

@@ -1,11 +1,22 @@
 from typing import Optional
 
+import adbc_driver_manager
+from cosmotech.orchestrator.utils.translate import T
+
 from cosmotech.coal.postgresql.runner import (
+    create_metadata,
+    get_metadata_last_run,
+    remove_run_metadata_from_postgresql,
     remove_runner_metadata_from_postgresql,
     send_runner_metadata_to_postgresql,
 )
-from cosmotech.coal.postgresql.store import dump_store_to_postgresql_from_conf
+from cosmotech.coal.postgresql.store import (
+    add_fk_constraints,
+    dump_store_to_postgresql_from_conf,
+)
 from cosmotech.coal.store.output.channel_interface import ChannelInterface
+from cosmotech.coal.utils.configuration import Dotdict
+from cosmotech.coal.utils.logger import LOGGER
 
 
 class PostgresChannel(ChannelInterface):
@@ -23,14 +34,43 @@ class PostgresChannel(ChannelInterface):
     }
     requirement_string = required_keys
 
+    def __init__(self, dct: Dotdict = None):
+        super().__init__(dct)
+        # set setup_db to default to True for compatibility
+        self.configuration.setup_db = self.configuration.safe_get("setup_db", True)
+
     def send(self, filter: Optional[list[str]] = None) -> bool:
-        run_id = send_runner_metadata_to_postgresql(self.configuration)
-        dump_store_to_postgresql_from_conf(
-            configuration=self.configuration,
-            selected_tables=filter,
-            fk_id=run_id,
-            replace=False,
-        )
+        if self.configuration.setup_db:
+            create_metadata(self.configuration)
+
+        # get old run id (present in metadata table)
+        last_run_ids = get_metadata_last_run(self.configuration)
+        rollout_run_id_list = [run_id for run_id in last_run_ids if run_id != self.configuration.cosmotech.run_id]
+
+        # add new run id in metadata table
+        new_run_id = self.configuration.cosmotech.run_id
+        send_runner_metadata_to_postgresql(self.configuration)
+
+        try:
+            # Send store's tables to PSQL
+            dump_store_to_postgresql_from_conf(
+                configuration=self.configuration,
+                selected_tables=filter,
+                fk_id=new_run_id,
+                replace=False,
+            )
+
+            if self.configuration.setup_db:
+                add_fk_constraints(self.configuration)
+        except adbc_driver_manager.ProgrammingError as e:
+            # remove newly created id
+            LOGGER.error(T("coal.store.output.postgres_channel.data_send_fail"))
+            remove_run_metadata_from_postgresql(self.configuration, new_run_id)
+            raise e
+        # rollout old run id
+        LOGGER.info(T("coal.store.output.postgres_channel.output_rollout").format(run_id=rollout_run_id_list))
+        for last_run_id in rollout_run_id_list:
+            remove_run_metadata_from_postgresql(self.configuration, last_run_id)
 
     def delete(self):
         # removing metadata will trigger cascade delete on real data

cosmotech/translation/coal/en-US/coal/services/postgresql.yml

@@ -25,4 +25,4 @@ ingesting_data: "Ingesting data with mode: {mode}"
 ingestion_success: "Successfully ingested {rows} rows"
 creating_table: "Creating table {schema_table}"
 metadata_updated: "Metadata updated"
-metadata_removed: "{runner_id} has been removed from Metadata table"
+metadata_removed: "{id} has been removed from Metadata table"

cosmotech/translation/coal/en-US/coal/store/output/postgres_channel.yml

@@ -0,0 +1,2 @@
+data_send_fail: "Data dump to postgres failed. Rolling back data."
+output_rollout: "Rolling out data: removing {run_id}"

tests/unit/coal/test_postgresql/test_postgresql_runner.py

@@ -28,6 +28,8 @@ def test_send_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ut
         mock_configuration.cosmotech.organization_id = "test-org"
         mock_configuration.cosmotech.workspace_id = "test-workspace"
         mock_configuration.cosmotech.runner_id = "test-runner-id"
+        mock_configuration.cosmotech.run_id = "test-run-id"
+        mock_configuration.cosmotech.run_template_id = "test-template-id"
 
         # Mock runner metadata
         mock_runner = {
@@ -56,7 +58,7 @@ def test_send_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ut
         mock_connect.return_value.__enter__.return_value = mock_conn
 
         # Act
-        result = send_runner_metadata_to_postgresql(mock_configuration)
+        send_runner_metadata_to_postgresql(mock_configuration)
 
         # Assert
         # Verify PostgresUtils was instantiated with configuration
@@ -72,19 +74,14 @@ def test_send_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ut
         mock_connect.assert_called_once_with("postgresql://user:password@localhost:5432/testdb", autocommit=True)
 
         # Check that SQL statements were executed
-        assert mock_cursor.execute.call_count == 3
+        assert mock_cursor.execute.call_count == 1
 
         # Verify the SQL statements (partially, since the exact SQL is complex)
-        create_table_call = mock_cursor.execute.call_args_list[0]
-        assert "CREATE TABLE IF NOT EXISTS" in create_table_call[0][0]
-        assert "public.test_runnermetadata" in create_table_call[0][0]
-
-        delete_call = mock_cursor.execute.call_args_list[1]
-        assert "DELETE FROM" in delete_call[0][0]
-        assert "public.test_runnermetadata" in delete_call[0][0]
-        assert delete_call[0][1] == ("test-runner-id",)
+        # create_table_call = mock_cursor.execute.call_args_list[0]
+        # assert "CREATE TABLE IF NOT EXISTS" in create_table_call[0][0]
+        # assert "public.test_runnermetadata" in create_table_call[0][0]
 
-        upsert_call = mock_cursor.execute.call_args_list[2]
+        upsert_call = mock_cursor.execute.call_args_list[0]
         assert "INSERT INTO" in upsert_call[0][0]
         assert "public.test_runnermetadata" in upsert_call[0][0]
         assert upsert_call[0][1] == (
@@ -95,10 +92,7 @@ def test_send_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ut
         )
 
         # Check that commits were called
-        assert mock_conn.commit.call_count == 3
-
-        # Verify the function returns the lastRunId
-        assert result == "test-run-id"
+        assert mock_conn.commit.call_count == 1
 
     @patch("cosmotech.coal.postgresql.runner.RunnerApi")
     @patch("cosmotech.coal.postgresql.runner.PostgresUtils")