diff --git a/cosmotech/coal/postgresql/runner.py b/cosmotech/coal/postgresql/runner.py index aee1aad5..392e00e7 100644 --- a/cosmotech/coal/postgresql/runner.py +++ b/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") diff --git a/cosmotech/coal/postgresql/store.py b/cosmotech/coal/postgresql/store.py index 95d6c422..d752208e 100644 --- a/cosmotech/coal/postgresql/store.py +++ b/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}", + ) + ) diff --git a/cosmotech/coal/postgresql/utils.py b/cosmotech/coal/postgresql/utils.py index a95e8d60..f5a76217 100644 --- a/cosmotech/coal/postgresql/utils.py +++ b/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: diff --git a/cosmotech/coal/store/output/postgres_channel.py b/cosmotech/coal/store/output/postgres_channel.py index 22894b80..e0edd83d 100644 --- a/cosmotech/coal/store/output/postgres_channel.py +++ b/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 diff --git a/cosmotech/translation/coal/en-US/coal/services/postgresql.yml b/cosmotech/translation/coal/en-US/coal/services/postgresql.yml index dc7a8a49..f6314d85 100644 --- a/cosmotech/translation/coal/en-US/coal/services/postgresql.yml +++ b/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" diff --git a/cosmotech/translation/coal/en-US/coal/store/output/postgres_channel.yml b/cosmotech/translation/coal/en-US/coal/store/output/postgres_channel.yml new file mode 100644 index 00000000..8c262539 --- /dev/null +++ b/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}" diff --git a/tests/unit/coal/test_postgresql/test_postgresql_runner.py b/tests/unit/coal/test_postgresql/test_postgresql_runner.py index 33c3f4a8..d529d9a3 100644 --- a/tests/unit/coal/test_postgresql/test_postgresql_runner.py +++ b/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") diff --git a/tests/unit/coal/test_store/test_output/test_postgres_channel.py b/tests/unit/coal/test_store/test_output/test_postgres_channel.py index 19046000..087f0fa0 100644 --- a/tests/unit/coal/test_store/test_output/test_postgres_channel.py +++ b/tests/unit/coal/test_store/test_output/test_postgres_channel.py @@ -21,7 +21,8 @@ def base_postgres_config(): "parameters_absolute_path": "/path/to/dataset", "organization_id": "org123", "workspace_id": "ws456", - "runner_id": "run789", + "runner_id": "r789", + "run_id": "run789", }, "postgres": { "host": "localhost", @@ -38,13 +39,23 @@ def base_postgres_config(): class TestPostgresChannel: """Tests for the PostgresChannel class.""" - def test_init_with_configuration(self, base_postgres_config): + def test_init(self, base_postgres_config): """Test PostgresChannel initialization with configuration.""" # Act channel = PostgresChannel(base_postgres_config) - # Assert - assert channel.configuration is not None + # Assert default value for setup_db is added + base_postgres_config.setup_db = True + assert channel.configuration == base_postgres_config + + def test_init_with_setup_db(self, base_postgres_config): + """Test PostgresChannel initialization with configuration.""" + # Act + base_postgres_config.setup_db = False # Set to False to test default behavior + channel = PostgresChannel(base_postgres_config) + + # Assert default value for setup_db is added + assert channel.configuration == base_postgres_config def test_required_keys(self): """Test that required_keys are properly defined.""" @@ -63,46 +74,165 @@ def test_required_keys(self): assert "user_name" in PostgresChannel.required_keys["postgres"] assert "user_password" in PostgresChannel.required_keys["postgres"] + @patch("cosmotech.coal.store.output.postgres_channel.create_metadata") + @patch("cosmotech.coal.store.output.postgres_channel.get_metadata_last_run") + @patch("cosmotech.coal.store.output.postgres_channel.dump_store_to_postgresql_from_conf") + @patch("cosmotech.coal.store.output.postgres_channel.send_runner_metadata_to_postgresql") + @patch("cosmotech.coal.store.output.postgres_channel.add_fk_constraints") + @patch("cosmotech.coal.store.output.postgres_channel.remove_run_metadata_from_postgresql") + def test_send_no_setup_db( + self, + mock_remove_run, + mock_add_fk, + mock_send_metadata, + mock_dump, + mock_get_metadata_last_run, + mock_create_metadata, + base_postgres_config, + ): + """Test sending data without table filter.""" + + base_postgres_config.setup_db = False # Set to False to test behavior when setup_db is False + channel = PostgresChannel(base_postgres_config) + mock_get_metadata_last_run.return_value = ["run789", "run456"] + + # Act + channel.send() + + # Assert + mock_create_metadata.assert_not_called() # This should be called only if setup_db is True + mock_get_metadata_last_run.Assert_called_once() + mock_send_metadata.assert_called_once() + mock_dump.assert_called_once() + mock_add_fk.assert_not_called() # This should be called only if setup_db is True + mock_remove_run.assert_called_once() + + @patch("cosmotech.coal.store.output.postgres_channel.create_metadata") + @patch("cosmotech.coal.store.output.postgres_channel.get_metadata_last_run") + @patch("cosmotech.coal.store.output.postgres_channel.dump_store_to_postgresql_from_conf") + @patch("cosmotech.coal.store.output.postgres_channel.send_runner_metadata_to_postgresql") + @patch("cosmotech.coal.store.output.postgres_channel.add_fk_constraints") + @patch("cosmotech.coal.store.output.postgres_channel.remove_run_metadata_from_postgresql") + def test_send_without_filter( + self, + mock_remove_run, + mock_add_fk, + mock_send_metadata, + mock_dump, + mock_get_metadata_last_run, + mock_create_metadata, + base_postgres_config, + ): + """Test sending data without table filter.""" + channel = PostgresChannel(base_postgres_config) + mock_get_metadata_last_run.return_value = ["run789", "run456"] + + # Act + channel.send() + + # Assert + base_postgres_config.setup_db = True + mock_create_metadata.assert_called_once() + mock_get_metadata_last_run.Assert_called_once() + mock_send_metadata.assert_called_once() + mock_dump.assert_called_once() + mock_add_fk.assert_called_once() + mock_remove_run.assert_called_once() + + # Check the arguments passed to dump_store_to_postgresql_from_conf + call_args = mock_dump.call_args + assert call_args.kwargs["configuration"] == base_postgres_config + assert call_args.kwargs["selected_tables"] is None + assert call_args.kwargs["fk_id"] == "run789" + + call_args_remove_run = mock_remove_run.call_args + assert call_args_remove_run.args[0] == base_postgres_config + assert call_args_remove_run.args[1] == "run456" + + @patch("cosmotech.coal.store.output.postgres_channel.create_metadata") + @patch("cosmotech.coal.store.output.postgres_channel.get_metadata_last_run") @patch("cosmotech.coal.store.output.postgres_channel.dump_store_to_postgresql_from_conf") @patch("cosmotech.coal.store.output.postgres_channel.send_runner_metadata_to_postgresql") - def test_send_without_filter(self, mock_send_metadata, mock_dump, base_postgres_config): + @patch("cosmotech.coal.store.output.postgres_channel.add_fk_constraints") + @patch("cosmotech.coal.store.output.postgres_channel.remove_run_metadata_from_postgresql") + def test_send_without_filter( + self, + mock_remove_run, + mock_add_fk, + mock_send_metadata, + mock_dump, + mock_get_metadata_last_run, + mock_create_metadata, + base_postgres_config, + ): """Test sending data without table filter.""" - mock_send_metadata.return_value = "run_id_123" channel = PostgresChannel(base_postgres_config) + mock_get_metadata_last_run.return_value = ["run789", "run456"] # Act channel.send() # Assert + base_postgres_config.setup_db = True + mock_create_metadata.assert_called_once() + mock_get_metadata_last_run.Assert_called_once() mock_send_metadata.assert_called_once() mock_dump.assert_called_once() + mock_add_fk.assert_called_once() + mock_remove_run.assert_called_once() # Check the arguments passed to dump_store_to_postgresql_from_conf call_args = mock_dump.call_args assert call_args.kwargs["configuration"] == base_postgres_config assert call_args.kwargs["selected_tables"] is None - assert call_args.kwargs["fk_id"] == "run_id_123" + assert call_args.kwargs["fk_id"] == "run789" + + call_args_remove_run = mock_remove_run.call_args + assert call_args_remove_run.args[0] == base_postgres_config + assert call_args_remove_run.args[1] == "run456" + @patch("cosmotech.coal.store.output.postgres_channel.create_metadata") + @patch("cosmotech.coal.store.output.postgres_channel.get_metadata_last_run") @patch("cosmotech.coal.store.output.postgres_channel.dump_store_to_postgresql_from_conf") @patch("cosmotech.coal.store.output.postgres_channel.send_runner_metadata_to_postgresql") - def test_send_with_filter(self, mock_send_metadata, mock_dump, base_postgres_config): + @patch("cosmotech.coal.store.output.postgres_channel.add_fk_constraints") + @patch("cosmotech.coal.store.output.postgres_channel.remove_run_metadata_from_postgresql") + def test_send_with_filter( + self, + mock_remove_run, + mock_add_fk, + mock_send_metadata, + mock_dump, + mock_get_metadata_last_run, + mock_create_metadata, + base_postgres_config, + ): """Test sending data with table filter.""" - mock_send_metadata.return_value = "run_id_456" channel = PostgresChannel(base_postgres_config) tables_filter = ["table1", "table2", "table3"] + mock_get_metadata_last_run.return_value = ["run789", "run456"] # Act channel.send(filter=tables_filter) # Assert + base_postgres_config.setup_db = True + mock_create_metadata.assert_called_once() + mock_get_metadata_last_run.Assert_called_once() mock_send_metadata.assert_called_once() mock_dump.assert_called_once() + mock_add_fk.assert_called_once() + mock_remove_run.assert_called_once() # Check the arguments passed to dump_store_to_postgresql_from_conf call_args = mock_dump.call_args - assert call_args[1]["configuration"] == base_postgres_config - assert call_args[1]["selected_tables"] == ["table1", "table2", "table3"] - assert call_args[1]["fk_id"] == "run_id_456" + assert call_args.kwargs["configuration"] == base_postgres_config + assert call_args.kwargs["selected_tables"] == ["table1", "table2", "table3"] + assert call_args.kwargs["fk_id"] == "run789" + + call_args_remove_run = mock_remove_run.call_args + assert call_args_remove_run.args[0] == base_postgres_config + assert call_args_remove_run.args[1] == "run456" @patch("cosmotech.coal.store.output.postgres_channel.remove_runner_metadata_from_postgresql") def test_delete(self, mock_remove_metadata, base_postgres_config): @@ -114,6 +244,7 @@ def test_delete(self, mock_remove_metadata, base_postgres_config): channel.delete() # Assert + base_postgres_config.setup_db = True mock_remove_metadata.assert_called_once() # Check that configuration was passed call_args = mock_remove_metadata.call_args