diff --git a/cosmotech/coal/postgresql/runner.py b/cosmotech/coal/postgresql/runner.py index 49d1638f..3634ccdf 100644 --- a/cosmotech/coal/postgresql/runner.py +++ b/cosmotech/coal/postgresql/runner.py @@ -44,7 +44,7 @@ def send_runner_metadata_to_postgresql( # Connect to PostgreSQL and update runner metadata with dbapi.connect(_psql.full_uri, autocommit=True) as conn: with conn.cursor() as curs: - schema_table = f"{str(_psql.db_schema)}.{str(_psql.table_prefix)}RunnerMetadata" + 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, @@ -102,7 +102,7 @@ def remove_runner_metadata_from_postgresql( # 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.table_prefix}RunnerMetadata" + schema_table = f"{_psql.db_schema}.{_psql.metadata_table_name}" last_run_id = runner.get("lastRunInfo").get("lastRunId") sql_delete_from_metatable = f""" DELETE FROM {schema_table} diff --git a/cosmotech/coal/postgresql/store.py b/cosmotech/coal/postgresql/store.py index 06425e96..8db8fb72 100644 --- a/cosmotech/coal/postgresql/store.py +++ b/cosmotech/coal/postgresql/store.py @@ -114,7 +114,7 @@ def dump_store_to_postgresql_from_conf( replace, ) if fk_id and _psql.is_metadata_exists(): - metadata_table = f"{_psql.table_prefix}RunnerMetadata" + metadata_table = f"{_psql.metadata_table_name}" _psql.add_fk_constraint(table_name, "csm_run_id", metadata_table, "last_csm_run_id") total_rows += rows diff --git a/cosmotech/coal/postgresql/utils.py b/cosmotech/coal/postgresql/utils.py index 62fed80c..4d863608 100644 --- a/cosmotech/coal/postgresql/utils.py +++ b/cosmotech/coal/postgresql/utils.py @@ -26,8 +26,8 @@ def __init__(self, configuration: Configuration): @property def table_prefix(self): if "table_prefix" in self._configuration: - return self._configuration.table_prefix - return "Cosmotech_" + return self._configuration.table_prefix.lower() + return "Cosmotech_".lower() @property def db_name(self): @@ -77,7 +77,7 @@ def full_uri(self) -> str: @property def metadata_table_name(self) -> str: - return f"{self.table_prefix}RunnerMetadata" + return f"{self.table_prefix}RunnerMetadata".lower() def get_postgresql_table_schema(self, target_table_name: str) -> Optional[pa.Schema]: """ diff --git a/tests/unit/coal/test_postgresql/test_postgresql_runner.py b/tests/unit/coal/test_postgresql/test_postgresql_runner.py index 711099af..870a1f89 100644 --- a/tests/unit/coal/test_postgresql/test_postgresql_runner.py +++ b/tests/unit/coal/test_postgresql/test_postgresql_runner.py @@ -46,7 +46,7 @@ def test_send_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ut mock_postgres_utils_instance = MagicMock() mock_postgres_utils_instance.full_uri = "postgresql://user:password@localhost:5432/testdb" mock_postgres_utils_instance.db_schema = "public" - mock_postgres_utils_instance.table_prefix = "Test_" + mock_postgres_utils_instance.metadata_table_name = "test_runnermetadata" mock_postgres_utils_class.return_value = mock_postgres_utils_instance # Mock PostgreSQL connection and cursor @@ -77,11 +77,11 @@ def test_send_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ut # 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] + assert "public.test_runnermetadata" in create_table_call[0][0] upsert_call = mock_cursor.execute.call_args_list[1] assert "INSERT INTO" in upsert_call[0][0] - assert "public.Test_RunnerMetadata" in upsert_call[0][0] + assert "public.test_runnermetadata" in upsert_call[0][0] assert upsert_call[0][1] == ( mock_runner["id"], mock_runner["name"], @@ -125,7 +125,7 @@ def test_remove_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ mock_postgres_utils_instance = MagicMock() mock_postgres_utils_instance.full_uri = "postgresql://user:password@localhost:5432/testdb" mock_postgres_utils_instance.db_schema = "public" - mock_postgres_utils_instance.table_prefix = "Test_" + mock_postgres_utils_instance.metadata_table_name = "test_runnermetadata" mock_postgres_utils_class.return_value = mock_postgres_utils_instance # Mock PostgreSQL connection and cursor @@ -156,7 +156,7 @@ def test_remove_runner_metadata_to_postgresql(self, mock_connect, mock_postgres_ # Verify the SQL statements (partially, since the exact SQL is complex) delete_call = mock_cursor.execute.call_args_list[0] assert "DELETE FROM" in delete_call[0][0] - assert "public.Test_RunnerMetadata" in delete_call[0][0] + assert "public.test_runnermetadata" in delete_call[0][0] # Check that commits were called assert mock_conn.commit.call_count == 1 diff --git a/tests/unit/coal/test_postgresql/test_postgresql_store.py b/tests/unit/coal/test_postgresql/test_postgresql_store.py index 6b548eec..f8bca0fc 100644 --- a/tests/unit/coal/test_postgresql/test_postgresql_store.py +++ b/tests/unit/coal/test_postgresql/test_postgresql_store.py @@ -42,7 +42,6 @@ def test_dump_store_to_postgresql_with_tables(self, mock_send_to_postgresql, moc postgres_schema = "public" postgres_user = "user" postgres_password = "password" - table_prefix = "Test_" replace = True _config = Configuration( @@ -58,7 +57,7 @@ def test_dump_store_to_postgresql_with_tables(self, mock_send_to_postgresql, moc "user_name": postgres_user, "user_password": postgres_password, "password_encoding": False, - "table_prefix": table_prefix, + "table_prefix": "Test_", }, } ) @@ -72,7 +71,7 @@ def test_dump_store_to_postgresql_with_tables(self, mock_send_to_postgresql, moc postgres_schema, postgres_user, postgres_password, - table_prefix, + "Test_", replace, ) @@ -93,12 +92,12 @@ def test_dump_store_to_postgresql_with_tables(self, mock_send_to_postgresql, moc [ call( table1_data, - f"{table_prefix}table1", + f"test_table1", replace, ), call( table2_data, - f"{table_prefix}table2", + f"test_table2", replace, ), ] @@ -266,6 +265,6 @@ def test_dump_store_to_postgresql_default_parameters(self, mock_send_to_postgres # Check that send_pyarrow_table_to_postgresql was called with default parameters mock_send_to_postgresql.assert_called_once_with( table_data, - "Cosmotech_table1", # Default table_prefix is "Cosmotech_" + "cosmotech_table1", # Default table_prefix is "Cosmotech_" but is sanitized to "cosmotech_" for psql True, # Default replace is True )