From 5475d87039a715d8d7c6c4e05f9841451f6222a8 Mon Sep 17 00:00:00 2001 From: johnfolly Date: Wed, 8 Apr 2026 17:50:56 +0200 Subject: [PATCH 1/4] fix: update foreign key constraints and ensure uniqueness in runner metadata --- cosmotech/coal/postgresql/runner.py | 2 +- cosmotech/coal/postgresql/store.py | 2 +- cosmotech/coal/postgresql/utils.py | 17 ++++++++++++++--- 3 files changed, 16 insertions(+), 5 deletions(-) diff --git a/cosmotech/coal/postgresql/runner.py b/cosmotech/coal/postgresql/runner.py index 34497bcd..503a24dd 100644 --- a/cosmotech/coal/postgresql/runner.py +++ b/cosmotech/coal/postgresql/runner.py @@ -49,7 +49,7 @@ def send_runner_metadata_to_postgresql( CREATE TABLE IF NOT EXISTS {schema_table} ( id varchar(32) PRIMARY KEY, name varchar(256), - last_csm_run_id varchar(32), + last_csm_run_id varchar(32) UNIQUE, run_template_id varchar(32) ); """ diff --git a/cosmotech/coal/postgresql/store.py b/cosmotech/coal/postgresql/store.py index 8db8fb72..95d6c422 100644 --- a/cosmotech/coal/postgresql/store.py +++ b/cosmotech/coal/postgresql/store.py @@ -115,7 +115,7 @@ def dump_store_to_postgresql_from_conf( ) if fk_id and _psql.is_metadata_exists(): metadata_table = f"{_psql.metadata_table_name}" - _psql.add_fk_constraint(table_name, "csm_run_id", metadata_table, "last_csm_run_id") + _psql.add_fk_constraint(target_table_name, "csm_run_id", metadata_table, "last_csm_run_id") total_rows += rows _up_time = perf_counter() diff --git a/cosmotech/coal/postgresql/utils.py b/cosmotech/coal/postgresql/utils.py index 4d863608..80b27c86 100644 --- a/cosmotech/coal/postgresql/utils.py +++ b/cosmotech/coal/postgresql/utils.py @@ -155,12 +155,23 @@ def add_fk_constraint( to_table: str, to_col: str, ) -> None: - # Connect to PostgreSQL and remove runner metadata row + # Connect to PostgreSQL and add a foreign key constraint with dbapi.connect(self.full_uri, autocommit=True) as conn: with conn.cursor() as curs: sql_add_fk = f""" - ALTER TABLE {self.db_schema}.{from_table} - CONSTRAINT metadata FOREIGN KEY ({from_col}) REFERENCES {to_table}({to_col}) + DO $$ + BEGIN + IF NOT EXISTS ( + SELECT 1 + FROM pg_constraint + WHERE conname = 'metadata' + AND conrelid = '{self.db_schema}.{from_table}'::regclass + ) THEN + ALTER TABLE {self.db_schema}.{from_table} + ADD CONSTRAINT metadata FOREIGN KEY ({from_col}) + REFERENCES {self.db_schema}.{to_table}({to_col}); + END IF; + END $$; """ curs.execute(sql_add_fk) conn.commit() From 57591eae31264d261ba6ec6432d8c3ee371ba10b Mon Sep 17 00:00:00 2001 From: johnfolly Date: Thu, 9 Apr 2026 09:54:15 +0200 Subject: [PATCH 2/4] fix: add metadata cleanup trigger to prevent foreign key violations --- cosmotech/coal/postgresql/runner.py | 10 ++++++++++ cosmotech/coal/postgresql/utils.py | 3 ++- .../coal/test_postgresql/test_postgresql_runner.py | 11 ++++++++--- 3 files changed, 20 insertions(+), 4 deletions(-) diff --git a/cosmotech/coal/postgresql/runner.py b/cosmotech/coal/postgresql/runner.py index 503a24dd..3649beb8 100644 --- a/cosmotech/coal/postgresql/runner.py +++ b/cosmotech/coal/postgresql/runner.py @@ -56,6 +56,16 @@ def send_runner_metadata_to_postgresql( LOGGER.info(T("coal.services.postgresql.creating_table").format(schema_table=schema_table)) curs.execute(sql_create_table) conn.commit() + + last_run_id = runner.get("lastRunInfo").get("lastRunId") + if last_run_id: + sql_delete_from_metatable = f""" + DELETE FROM {schema_table} + WHERE last_csm_run_id= $1; + """ + curs.execute(sql_delete_from_metatable, (last_run_id,)) + conn.commit() + sql_upsert = f""" INSERT INTO {schema_table} (id, name, last_csm_run_id, run_template_id) VALUES ($1, $2, $3, $4) diff --git a/cosmotech/coal/postgresql/utils.py b/cosmotech/coal/postgresql/utils.py index 80b27c86..4ac078ea 100644 --- a/cosmotech/coal/postgresql/utils.py +++ b/cosmotech/coal/postgresql/utils.py @@ -169,7 +169,8 @@ def add_fk_constraint( ) THEN ALTER TABLE {self.db_schema}.{from_table} ADD CONSTRAINT metadata FOREIGN KEY ({from_col}) - REFERENCES {self.db_schema}.{to_table}({to_col}); + REFERENCES {self.db_schema}.{to_table}({to_col}) + ON DELETE CASCADE; END IF; END $$; """ diff --git a/tests/unit/coal/test_postgresql/test_postgresql_runner.py b/tests/unit/coal/test_postgresql/test_postgresql_runner.py index 870a1f89..83789a07 100644 --- a/tests/unit/coal/test_postgresql/test_postgresql_runner.py +++ b/tests/unit/coal/test_postgresql/test_postgresql_runner.py @@ -72,14 +72,19 @@ 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 == 2 + assert mock_cursor.execute.call_count == 3 # 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] - upsert_call = mock_cursor.execute.call_args_list[1] + 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-run-id",) + + upsert_call = mock_cursor.execute.call_args_list[2] assert "INSERT INTO" in upsert_call[0][0] assert "public.test_runnermetadata" in upsert_call[0][0] assert upsert_call[0][1] == ( @@ -90,7 +95,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 == 2 + assert mock_conn.commit.call_count == 3 # Verify the function returns the lastRunId assert result == "test-run-id" From 85a3c38b35f47b40488e9dd4eaeefdcf04179aa0 Mon Sep 17 00:00:00 2001 From: johnfolly Date: Thu, 9 Apr 2026 12:32:04 +0200 Subject: [PATCH 3/4] fix: update foreign key reference in runner metadata deletion logic --- cosmotech/coal/postgresql/runner.py | 15 +++++++-------- .../test_postgresql/test_postgresql_runner.py | 2 +- 2 files changed, 8 insertions(+), 9 deletions(-) diff --git a/cosmotech/coal/postgresql/runner.py b/cosmotech/coal/postgresql/runner.py index 3649beb8..21c454e5 100644 --- a/cosmotech/coal/postgresql/runner.py +++ b/cosmotech/coal/postgresql/runner.py @@ -57,14 +57,13 @@ def send_runner_metadata_to_postgresql( curs.execute(sql_create_table) conn.commit() - last_run_id = runner.get("lastRunInfo").get("lastRunId") - if last_run_id: - sql_delete_from_metatable = f""" - DELETE FROM {schema_table} - WHERE last_csm_run_id= $1; - """ - curs.execute(sql_delete_from_metatable, (last_run_id,)) - 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) diff --git a/tests/unit/coal/test_postgresql/test_postgresql_runner.py b/tests/unit/coal/test_postgresql/test_postgresql_runner.py index 83789a07..33c3f4a8 100644 --- a/tests/unit/coal/test_postgresql/test_postgresql_runner.py +++ b/tests/unit/coal/test_postgresql/test_postgresql_runner.py @@ -82,7 +82,7 @@ def test_send_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ut 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-run-id",) + assert delete_call[0][1] == ("test-runner-id",) upsert_call = mock_cursor.execute.call_args_list[2] assert "INSERT INTO" in upsert_call[0][0] From 7a13c1f2fb5c4bfab8bcc8e87448959382976d8e Mon Sep 17 00:00:00 2001 From: =?UTF-8?q?Laurent=20Al=C3=A9p=C3=A9e?= Date: Thu, 9 Apr 2026 16:31:23 +0200 Subject: [PATCH 4/4] chore: remove unecessary code check on conflict is not necessary anymore as we remove the row beforehand --- cosmotech/coal/postgresql/runner.py | 5 +---- 1 file changed, 1 insertion(+), 4 deletions(-) diff --git a/cosmotech/coal/postgresql/runner.py b/cosmotech/coal/postgresql/runner.py index 21c454e5..9236f3e6 100644 --- a/cosmotech/coal/postgresql/runner.py +++ b/cosmotech/coal/postgresql/runner.py @@ -67,10 +67,7 @@ def send_runner_metadata_to_postgresql( sql_upsert = f""" INSERT INTO {schema_table} (id, name, last_csm_run_id, run_template_id) - VALUES ($1, $2, $3, $4) - ON CONFLICT (id) - DO - UPDATE SET name = EXCLUDED.name, last_csm_run_id = EXCLUDED.last_csm_run_id; + VALUES ($1, $2, $3, $4) """ LOGGER.debug(runner) curs.execute(