Add option to set the schema index from the command line (DM-43683)

lsst · Apr 22, 2024 · 9d639fb · 9d639fb
1 parent 2340b21
commit 9d639fb
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 1 deletion.
diff --git a/python/felis/cli.py b/python/felis/cli.py
@@ -183,6 +183,7 @@ def init_tap(
 @click.option("--tap-columns-table", help="Alt Table Name for TAP_SCHEMA.columns")
 @click.option("--tap-keys-table", help="Alt Table Name for TAP_SCHEMA.keys")
 @click.option("--tap-key-columns-table", help="Alt Table Name for TAP_SCHEMA.key_columns")
+@click.option("--tap-schema-index", type=int, help="TAP_SCHEMA index of the schema")
 @click.argument("file", type=click.File())
 def load_tap(
     engine_url: str,
@@ -196,6 +197,7 @@ def load_tap(
     tap_columns_table: str,
     tap_keys_table: str,
     tap_key_columns_table: str,
+    tap_schema_index: int,
     file: io.TextIOBase,
 ) -> None:
     """Load TAP metadata from a Felis FILE.
@@ -228,6 +230,7 @@ def load_tap(
             catalog_name=catalog_name,
             schema_name=schema_name,
             tap_tables=tap_tables,
+            tap_schema_index=tap_schema_index,
         )
         tap_visitor.visit_schema(schema)
     else:
@@ -241,6 +244,7 @@ def load_tap(
             catalog_name=catalog_name,
             schema_name=schema_name,
             tap_tables=tap_tables,
+            tap_schema_index=tap_schema_index,
         )
         tap_visitor.visit_schema(schema)
 

diff --git a/python/felis/tap.py b/python/felis/tap.py
@@ -153,13 +153,15 @@ def __init__(
         catalog_name: str | None = None,
         schema_name: str | None = None,
         tap_tables: MutableMapping[str, Any] | None = None,
+        tap_schema_index: int | None = None,
     ):
         self.graph_index: MutableMapping[str, Any] = {}
         self.catalog_name = catalog_name
         self.schema_name = schema_name
         self.engine = engine
         self._mock_connection: MockConnection | None = None
         self.tables = tap_tables or init_tables()
+        self.tap_schema_index = tap_schema_index
 
     @classmethod
     def from_mock_connection(
@@ -168,9 +170,11 @@ def from_mock_connection(
         catalog_name: str | None = None,
         schema_name: str | None = None,
         tap_tables: MutableMapping[str, Any] | None = None,
+        tap_schema_index: int | None = None,
     ) -> TapLoadingVisitor:
         visitor = cls(engine=None, catalog_name=catalog_name, schema_name=schema_name, tap_tables=tap_tables)
         visitor._mock_connection = mock_connection
+        visitor.tap_schema_index = tap_schema_index
         return visitor
 
     def visit_schema(self, schema_obj: Schema) -> None:
@@ -181,7 +185,8 @@ def visit_schema(self, schema_obj: Schema) -> None:
         schema.schema_name = self._schema_name()
         schema.description = schema_obj.description
         schema.utype = schema_obj.votable_utype
-        schema.schema_index = schema_obj.tap_schema_index
+        schema.schema_index = self.tap_schema_index
+        logger.debug("Set TAP_SCHEMA index: {}".format(self.tap_schema_index))
 
         if self.engine is not None:
             session: Session = sessionmaker(self.engine)()