Skip to content

Commit

Permalink
Merge pull request #130 from lsst/tickets/DM-46957
Browse files Browse the repository at this point in the history
DM-46957: Add an optional table name postfix to the metadata builder
  • Loading branch information
JeremyMcCormick authored Jan 27, 2025
2 parents 065c647 + b9ed319 commit 8294325
Show file tree
Hide file tree
Showing 4 changed files with 30 additions and 5 deletions.
4 changes: 3 additions & 1 deletion python/felis/metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -134,6 +134,7 @@ def __init__(
schema: Schema,
apply_schema_to_metadata: bool = True,
ignore_constraints: bool = False,
table_name_postfix: str = "",
) -> None:
"""Initialize the metadata builder."""
self.schema = schema
Expand All @@ -142,6 +143,7 @@ def __init__(
self.metadata = MetaData(schema=schema.name if apply_schema_to_metadata else None)
self._objects: dict[str, Any] = {}
self.ignore_constraints = ignore_constraints
self.table_name_postfix = table_name_postfix

def build(self) -> MetaData:
"""Build the SQLAlchemy tables and constraints from the schema.
Expand Down Expand Up @@ -225,7 +227,7 @@ def build_table(self, table_obj: datamodel.Table) -> None:
description = table_obj.description
columns = [self.build_column(column) for column in table_obj.columns]
table = Table(
name,
name + self.table_name_postfix,
self.metadata,
*columns,
comment=description,
Expand Down
16 changes: 12 additions & 4 deletions python/felis/tap_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,9 +91,15 @@ def __init__(
self.table_name_postfix = table_name_postfix
self.apply_schema_to_metadata = apply_schema_to_metadata
self.schema_name = schema_name or TableManager._SCHEMA_NAME_STD
self.table_name_postfix = table_name_postfix

if is_valid_engine(engine):
assert isinstance(engine, Engine)
if table_name_postfix != "":
logger.warning(
"Table name postfix '%s' will be ignored when reflecting TAP_SCHEMA database",
table_name_postfix,
)
logger.debug(
"Reflecting TAP_SCHEMA database from existing database at %s",
engine.url._replace(password="***"),
Expand Down Expand Up @@ -131,7 +137,9 @@ def _load_yaml(self) -> None:
self.schema_name = self.schema.name

self._metadata = MetaDataBuilder(
self.schema, apply_schema_to_metadata=self.apply_schema_to_metadata
self.schema,
apply_schema_to_metadata=self.apply_schema_to_metadata,
table_name_postfix=self.table_name_postfix,
).build()

logger.debug("Loaded TAP_SCHEMA '%s' from YAML resource", self.schema_name)
Expand Down Expand Up @@ -424,7 +432,7 @@ def load(self) -> None:
# Execute the inserts if not in dry run mode.
self._execute_inserts()
else:
logger.info("Dry run: not loading data into database")
logger.info("Dry run - not loading data into database")

def _insert_schemas(self) -> None:
"""Insert the schema data into the schemas table."""
Expand Down Expand Up @@ -565,15 +573,15 @@ def _compiled_inserts(self) -> list[str]:
def _print_sql(self) -> None:
"""Print the generated inserts to stdout."""
for insert_str in self._compiled_inserts():
print(insert_str)
print(insert_str + ";")

def _write_sql_to_file(self) -> None:
"""Write the generated insert statements to a file."""
if not self.output_path:
raise ValueError("No output path specified")
with open(self.output_path, "w") as outfile:
for insert_str in self._compiled_inserts():
outfile.write(insert_str + "\n")
outfile.write(insert_str + ";" + "\n")

def _insert(self, table_name: str, record: list[Any] | dict[str, Any]) -> None:
"""Generate an insert statement for a record.
Expand Down
9 changes: 9 additions & 0 deletions tests/test_metadata.py
Original file line number Diff line number Diff line change
Expand Up @@ -231,6 +231,15 @@ def test_ignore_constraints(self) -> None:
msg=f"Table {table.name} has non-primary key constraints defined",
)

def test_table_name_postfix(self) -> None:
"""Test that table name postfixes are correctly applied."""
schema = Schema.model_validate(self.yaml_data)
schema.name = "main"
builder = MetaDataBuilder(schema, table_name_postfix="_test")
md = builder.build()
for table in md.tables.values():
self.assertTrue(table.name.endswith("_test"))


if __name__ == "__main__":
unittest.main()
6 changes: 6 additions & 0 deletions tests/test_tap_schema.py
Original file line number Diff line number Diff line change
Expand Up @@ -60,6 +60,12 @@ def test_create_table_manager(self) -> None:
# already been created.
mgr = TableManager()

def test_table_name_postfix(self) -> None:
"""Test the table name postfix."""
mgr = TableManager(table_name_postfix="_test")
for table_name in mgr.metadata.tables:
self.assertTrue(table_name.endswith("_test"))


class DataLoaderTestCase(unittest.TestCase):
"""Test the `DataLoader` class."""
Expand Down

0 comments on commit 8294325

Please sign in to comment.