Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 2 additions & 2 deletions cosmotech/coal/postgresql/runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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,
Expand Down Expand Up @@ -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}
Expand Down
2 changes: 1 addition & 1 deletion cosmotech/coal/postgresql/store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
6 changes: 3 additions & 3 deletions cosmotech/coal/postgresql/utils.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down Expand Up @@ -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]:
"""
Expand Down
10 changes: 5 additions & 5 deletions tests/unit/coal/test_postgresql/test_postgresql_runner.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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"],
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand Down
11 changes: 5 additions & 6 deletions tests/unit/coal/test_postgresql/test_postgresql_store.py
Original file line number Diff line number Diff line change
Expand Up @@ -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(
Expand All @@ -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_",
},
}
)
Expand All @@ -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,
)

Expand All @@ -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,
),
]
Expand Down Expand Up @@ -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
)
Loading