Coverage for python/felis/tap_schema.py: 23%
263 statements
« prev ^ index » next coverage.py v7.11.0, created at 2025-11-20 01:50 -0800
1"""Provides utilities for creating and populating the TAP_SCHEMA database."""
3# This file is part of felis.
4#
5# Developed for the LSST Data Management System.
6# This product includes software developed by the LSST Project
7# (https://www.lsst.org).
8# See the COPYRIGHT file at the top-level directory of this distribution
9# for details of code ownership.
10#
11# This program is free software: you can redistribute it and/or modify
12# it under the terms of the GNU General Public License as published by
13# the Free Software Foundation, either version 3 of the License, or
14# (at your option) any later version.
15#
16# This program is distributed in the hope that it will be useful,
17# but WITHOUT ANY WARRANTY; without even the implied warranty of
18# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
19# GNU General Public License for more details.
20#
21# You should have received a copy of the GNU General Public License
22# along with this program. If not, see <https://www.gnu.org/licenses/>.
24import csv
25import io
26import logging
27import os
28import re
29from typing import Any
31from lsst.resources import ResourcePath
32from sqlalchemy import MetaData, Table, select, text
33from sqlalchemy.engine import Connection, Engine
34from sqlalchemy.engine.mock import MockConnection
35from sqlalchemy.exc import SQLAlchemyError
36from sqlalchemy.schema import CreateSchema
37from sqlalchemy.sql.dml import Insert
39from felis import datamodel
40from felis.datamodel import Constraint, Schema
41from felis.db.utils import is_valid_engine
42from felis.metadata import MetaDataBuilder
44from .types import FelisType
46__all__ = ["DataLoader", "TableManager"]
48logger = logging.getLogger(__name__)
51class TableManager:
52 """Manage creation of TAP_SCHEMA tables.
54 Parameters
55 ----------
56 engine
57 The SQLAlchemy engine for reflecting the TAP_SCHEMA tables from an
58 existing database.
59 This can be a mock connection or None, in which case the internal
60 TAP_SCHEMA schema will be used by loading an internal YAML file.
61 schema_name
62 The name of the schema to use for the TAP_SCHEMA tables.
63 Leave as None to use the standard name of "TAP_SCHEMA".
64 apply_schema_to_metadata
65 If True, apply the schema to the metadata as well as the tables.
66 If False, these will be set to None, e.g., for sqlite.
67 table_name_postfix
68 A string to append to all the standard table names.
69 This needs to be used in a way such that the resultant table names
70 map to tables within the TAP_SCHEMA database.
71 extensions_path
72 The path to additional TAP_SCHEMA table definitions.
74 Notes
75 -----
76 The TAP_SCHEMA schema must either have been created already, in which case
77 the ``engine`` should be provided. Or the internal TAP_SCHEMA schema will
78 be used if ``engine`` is None or a ``MockConnection``.
79 """
81 _TABLE_NAMES_STD = ["schemas", "tables", "columns", "keys", "key_columns"]
82 """The standard table names for the TAP_SCHEMA tables."""
84 _SCHEMA_NAME_STD = "TAP_SCHEMA"
85 """The standard schema name for the TAP_SCHEMA tables."""
87 def __init__(
88 self,
89 engine: Engine | MockConnection | None = None,
90 schema_name: str | None = None,
91 apply_schema_to_metadata: bool = True,
92 table_name_postfix: str = "",
93 extensions_path: str | None = None,
94 ):
95 """Initialize the table manager."""
96 self.table_name_postfix = table_name_postfix
97 self.apply_schema_to_metadata = apply_schema_to_metadata
98 self.schema_name = schema_name or TableManager._SCHEMA_NAME_STD
99 self.extensions_path = extensions_path
101 if is_valid_engine(engine):
102 assert isinstance(engine, Engine)
103 if table_name_postfix != "":
104 logger.warning(
105 "Table name postfix '%s' will be ignored when reflecting TAP_SCHEMA database",
106 table_name_postfix,
107 )
108 logger.debug(
109 "Reflecting TAP_SCHEMA database from existing database at %s",
110 engine.url._replace(password="***"),
111 )
112 self._reflect(engine)
113 else:
114 self._load_yaml()
115 self._apply_extensions()
116 self._build_metadata()
118 self._create_table_map()
119 self._check_tables()
121 def _apply_extensions(self) -> None:
122 """Apply extensions from a YAML file to the TAP_SCHEMA schema.
124 This method loads extension column definitions from a YAML file and
125 adds them to the appropriate TAP_SCHEMA tables.
126 """
127 if not self.extensions_path:
128 return
130 logger.info("Loading TAP_SCHEMA extensions from: %s", self.extensions_path)
131 extensions_schema = Schema.from_uri(self.extensions_path, context={"id_generation": True})
133 if not extensions_schema.tables:
134 logger.warning("Extensions schema does not contain any tables, no extensions applied")
135 return
137 extension_count = 0
138 extension_tables = {table.name: table.columns for table in extensions_schema.tables}
140 for table in self.schema.tables:
141 extension_columns = extension_tables.get(table.name)
142 if extension_columns:
143 table.columns = list(table.columns) + list(extension_columns)
144 extension_count += len(extension_columns)
145 logger.debug("Added %d extension columns to table '%s'", len(extension_columns), table.name)
147 logger.info("Applied %d extension columns to TAP_SCHEMA", extension_count)
149 def _reflect(self, engine: Engine) -> None:
150 """Reflect the TAP_SCHEMA database tables into the metadata.
152 Parameters
153 ----------
154 engine
155 The SQLAlchemy engine to use to reflect the tables.
156 """
157 self._metadata = MetaData(schema=self.schema_name if self.apply_schema_to_metadata else None)
158 try:
159 self.metadata.reflect(bind=engine)
160 except SQLAlchemyError as e:
161 logger.error("Error reflecting TAP_SCHEMA database: %s", e)
162 raise
164 def _build_metadata(self) -> None:
165 """Build SQLAlchemy metadata from the Felis schema."""
166 self._metadata = MetaDataBuilder(
167 self.schema,
168 apply_schema_to_metadata=self.apply_schema_to_metadata,
169 table_name_postfix=self.table_name_postfix,
170 ).build()
172 def _load_yaml(self) -> None:
173 """Load the standard TAP_SCHEMA schema from a Felis package
174 resource.
175 """
176 self._load_schema()
177 if self.schema_name != TableManager._SCHEMA_NAME_STD:
178 self.schema.name = self.schema_name
179 else:
180 self.schema_name = self.schema.name
182 logger.debug("Loaded TAP_SCHEMA '%s' from YAML resource", self.schema_name)
184 def __getitem__(self, table_name: str) -> Table:
185 """Get one of the TAP_SCHEMA tables by its standard TAP_SCHEMA name.
187 Parameters
188 ----------
189 table_name
190 The name of the table to get.
192 Returns
193 -------
194 Table
195 The table with the given name.
197 Notes
198 -----
199 This implements array semantics for the table manager, allowing
200 tables to be accessed by their standard TAP_SCHEMA names.
201 """
202 if table_name not in self._table_map:
203 raise KeyError(f"Table '{table_name}' not found in TAP_SCHEMA")
204 return self.metadata.tables[self._table_map[table_name]]
206 @property
207 def schema(self) -> Schema:
208 """Get the TAP_SCHEMA schema.
210 Returns
211 -------
212 Schema
213 The TAP_SCHEMA schema.
215 Notes
216 -----
217 This will only be set if the TAP_SCHEMA schema was loaded from a
218 Felis package resource. In the case where the TAP_SCHEMA schema was
219 reflected from an existing database, this will be None.
220 """
221 return self._schema
223 @property
224 def metadata(self) -> MetaData:
225 """Get the metadata for the TAP_SCHEMA tables.
227 Returns
228 -------
229 `~sqlalchemy.sql.schema.MetaData`
230 The metadata for the TAP_SCHEMA tables.
232 Notes
233 -----
234 This will either be the metadata that was reflected from an existing
235 database or the metadata that was loaded from a Felis package resource.
236 """
237 return self._metadata
239 @classmethod
240 def get_tap_schema_std_path(cls) -> str:
241 """Get the path to the standard TAP_SCHEMA schema resource.
243 Returns
244 -------
245 str
246 The path to the standard TAP_SCHEMA schema resource.
247 """
248 return os.path.join(os.path.dirname(__file__), "config", "tap_schema", "tap_schema_std.yaml")
250 @classmethod
251 def get_tap_schema_std_resource(cls) -> ResourcePath:
252 """Get the standard TAP_SCHEMA schema resource.
254 Returns
255 -------
256 `~lsst.resources.ResourcePath`
257 The standard TAP_SCHEMA schema resource.
258 """
259 return ResourcePath("resource://felis/config/tap_schema/tap_schema_std.yaml")
261 @classmethod
262 def get_table_names_std(cls) -> list[str]:
263 """Get the standard column names for the TAP_SCHEMA tables.
265 Returns
266 -------
267 list
268 The standard table names for the TAP_SCHEMA tables.
269 """
270 return cls._TABLE_NAMES_STD
272 @classmethod
273 def get_schema_name_std(cls) -> str:
274 """Get the standard schema name for the TAP_SCHEMA tables.
276 Returns
277 -------
278 str
279 The standard schema name for the TAP_SCHEMA tables.
280 """
281 return cls._SCHEMA_NAME_STD
283 @classmethod
284 def load_schema_resource(cls) -> Schema:
285 """Load the standard TAP_SCHEMA schema from a Felis package
286 resource into a Felis `~felis.datamodel.Schema`.
288 Returns
289 -------
290 Schema
291 The TAP_SCHEMA schema.
292 """
293 rp = cls.get_tap_schema_std_resource()
294 return Schema.from_uri(rp, context={"id_generation": True})
296 def _load_schema(self) -> None:
297 """Load the TAP_SCHEMA schema from a Felis package resource."""
298 self._schema = self.load_schema_resource()
300 def _create_table_map(self) -> None:
301 """Create a mapping of standard table names to the table names modified
302 with a postfix, as well as the prepended schema name if it is set.
304 Returns
305 -------
306 dict
307 A dictionary mapping the standard table names to the modified
308 table names.
310 Notes
311 -----
312 This is a private method that is called during initialization, allowing
313 us to use table names like ``schemas11`` such as those used by the CADC
314 TAP library instead of the standard table names. It also maps between
315 the standard table names and those with the schema name prepended like
316 SQLAlchemy uses.
317 """
318 self._table_map = {
319 table_name: (
320 f"{self.schema_name + '.' if self.apply_schema_to_metadata else ''}"
321 f"{table_name}{self.table_name_postfix}"
322 )
323 for table_name in TableManager.get_table_names_std()
324 }
325 logger.debug(f"Created TAP_SCHEMA table map: {self._table_map}")
327 def _check_tables(self) -> None:
328 """Check that there is a valid mapping to each standard table.
330 Raises
331 ------
332 KeyError
333 If a table is missing from the table map.
334 """
335 for table_name in TableManager.get_table_names_std():
336 self[table_name]
338 def _create_schema(self, engine: Engine) -> None:
339 """Create the database schema for TAP_SCHEMA if it does not already
340 exist.
342 Parameters
343 ----------
344 engine
345 The SQLAlchemy engine to use to create the schema.
347 Notes
348 -----
349 This method only creates the schema in the database. It does not create
350 the tables.
351 """
352 create_schema_functions = {
353 "postgresql": self._create_schema_postgresql,
354 "mysql": self._create_schema_mysql,
355 }
357 dialect_name = engine.dialect.name
358 if dialect_name == "sqlite":
359 # SQLite doesn't have schemas.
360 return
362 create_function = create_schema_functions.get(dialect_name)
364 if create_function:
365 with engine.begin() as connection:
366 create_function(connection)
367 else:
368 # Some other database engine we don't currently know how to handle.
369 raise NotImplementedError(
370 f"Database engine '{engine.dialect.name}' is not supported for schema creation"
371 )
373 def _create_schema_postgresql(self, connection: Connection) -> None:
374 """Create the schema in a PostgreSQL database.
376 Parameters
377 ----------
378 connection
379 The SQLAlchemy connection to use to create the schema.
380 """
381 connection.execute(CreateSchema(self.schema_name, if_not_exists=True))
383 def _create_schema_mysql(self, connection: Connection) -> None:
384 """Create the schema in a MySQL database.
386 Parameters
387 ----------
388 connection
389 The SQLAlchemy connection to use to create the schema.
390 """
391 connection.execute(text(f"CREATE DATABASE IF NOT EXISTS {self.schema_name}"))
393 def initialize_database(self, engine: Engine) -> None:
394 """Initialize a database with the TAP_SCHEMA tables.
396 Parameters
397 ----------
398 engine
399 The SQLAlchemy engine to use to create the tables.
400 """
401 logger.info("Creating TAP_SCHEMA database '%s'", self.schema_name)
402 self._create_schema(engine)
403 self.metadata.create_all(engine)
405 def select(self, engine: Engine, table_name: str, filter_condition: str = "") -> list[dict[str, Any]]:
406 """Select all rows from a TAP_SCHEMA table with an optional filter
407 condition.
409 Parameters
410 ----------
411 engine
412 The SQLAlchemy engine to use to connect to the database.
413 table_name
414 The name of the table to select from.
415 filter_condition
416 The filter condition as a string. If empty, no filter will be
417 applied.
419 Returns
420 -------
421 list
422 A list of dictionaries containing the rows from the table.
423 """
424 table = self[table_name]
425 query = select(table)
426 if filter_condition:
427 query = query.where(text(filter_condition))
428 with engine.connect() as connection:
429 result = connection.execute(query)
430 rows = [dict(row._mapping) for row in result]
431 return rows
434class DataLoader:
435 """Load data into the TAP_SCHEMA tables.
437 Parameters
438 ----------
439 schema
440 The Felis ``Schema`` to load into the TAP_SCHEMA tables.
441 mgr
442 The table manager that contains the TAP_SCHEMA tables.
443 engine
444 The SQLAlchemy engine to use to connect to the database.
445 tap_schema_index
446 The index of the schema in the TAP_SCHEMA database.
447 output_path
448 The file to write the SQL statements to. If None, printing will be
449 suppressed.
450 print_sql
451 If True, print the SQL statements that will be executed.
452 dry_run
453 If True, the data will not be loaded into the database.
454 unique_keys
455 If True, prepend the schema name to the key name to make it unique
456 when loading data into the keys and key_columns tables.
457 """
459 def __init__(
460 self,
461 schema: Schema,
462 mgr: TableManager,
463 engine: Engine | MockConnection,
464 tap_schema_index: int = 0,
465 output_path: str | None = None,
466 print_sql: bool = False,
467 dry_run: bool = False,
468 unique_keys: bool = False,
469 ):
470 self.schema = schema
471 self.mgr = mgr
472 self.engine = engine
473 self.tap_schema_index = tap_schema_index
474 self.inserts: list[Insert] = []
475 self.output_path = output_path
476 self.print_sql = print_sql
477 self.dry_run = dry_run
478 self.unique_keys = unique_keys
480 def load(self) -> None:
481 """Load the schema data into the TAP_SCHEMA tables.
483 Notes
484 -----
485 This will generate inserts for the data, print the SQL statements if
486 requested, save the SQL statements to a file if requested, and load the
487 data into the database if not in dry run mode. These are done as
488 sequential operations rather than for each insert. The logic is that
489 the user may still want the complete SQL output to be printed or saved
490 to a file even if loading into the database causes errors. If there are
491 errors when inserting into the database, the SQLAlchemy error message
492 should indicate which SQL statement caused the error.
493 """
494 self._generate_all_inserts()
495 if self.print_sql:
496 # Print to stdout.
497 self._print_sql()
498 if self.output_path:
499 # Print to an output file.
500 self._write_sql_to_file()
501 if not self.dry_run:
502 # Execute the inserts if not in dry run mode.
503 self._execute_inserts()
504 else:
505 logger.info("Dry run - not loading data into database")
507 def _insert_schemas(self) -> None:
508 """Insert the schema data into the ``schemas`` table."""
509 schema_record = {
510 "schema_name": self.schema.name,
511 "utype": self.schema.votable_utype,
512 "description": self.schema.description,
513 "schema_index": self.tap_schema_index,
514 }
515 self._insert("schemas", schema_record)
517 def _get_table_name(self, table: datamodel.Table) -> str:
518 """Get the name of the table with the schema name prepended.
520 Parameters
521 ----------
522 table
523 The table to get the name for.
525 Returns
526 -------
527 str
528 The name of the table with the schema name prepended.
529 """
530 return f"{self.schema.name}.{table.name}"
532 def _insert_tables(self) -> None:
533 """Insert the table data into the ``tables`` table."""
534 for table in self.schema.tables:
535 table_record = {
536 "schema_name": self.schema.name,
537 "table_name": self._get_table_name(table),
538 "table_type": "table",
539 "utype": table.votable_utype,
540 "description": table.description,
541 "table_index": 0 if table.tap_table_index is None else table.tap_table_index,
542 }
543 self._insert("tables", table_record)
545 def _insert_columns(self) -> None:
546 """Insert the column data into the ``columns`` table."""
547 for table in self.schema.tables:
548 for column in table.columns:
549 felis_type = FelisType.felis_type(column.datatype.value)
550 arraysize = str(column.votable_arraysize) if column.votable_arraysize else None
551 size = DataLoader._get_size(column)
552 indexed = DataLoader._is_indexed(column, table)
553 tap_column_index = column.tap_column_index
554 unit = column.ivoa_unit or column.fits_tunit
556 column_record = {
557 "table_name": self._get_table_name(table),
558 "column_name": column.name,
559 "datatype": felis_type.votable_name,
560 "arraysize": arraysize,
561 "size": size,
562 "xtype": column.votable_xtype,
563 "description": column.description,
564 "utype": column.votable_utype,
565 "unit": unit,
566 "ucd": column.ivoa_ucd,
567 "indexed": indexed,
568 "principal": column.tap_principal,
569 "std": column.tap_std,
570 "column_index": tap_column_index,
571 }
572 self._insert("columns", column_record)
574 def _get_key(self, constraint: Constraint) -> str:
575 """Get the key name for a constraint.
577 Parameters
578 ----------
579 constraint
580 The constraint to get the key name for.
582 Returns
583 -------
584 str
585 The key name for the constraint.
587 Notes
588 -----
589 This will prepend the name of the schema to the key name if the
590 `unique_keys` attribute is set to True. Otherwise, it will just return
591 the name of the constraint.
592 """
593 if self.unique_keys:
594 key_id = f"{self.schema.name}_{constraint.name}"
595 logger.debug("Generated unique key_id: %s -> %s", constraint.name, key_id)
596 else:
597 key_id = constraint.name
598 return key_id
600 def _insert_keys(self) -> None:
601 """Insert the foreign keys into the ``keys`` and ``key_columns``
602 tables.
603 """
604 for table in self.schema.tables:
605 for constraint in table.constraints:
606 if isinstance(constraint, datamodel.ForeignKeyConstraint):
607 ###########################################################
608 # Handle keys table
609 ###########################################################
610 referenced_column = self.schema.find_object_by_id(
611 constraint.referenced_columns[0], datamodel.Column
612 )
613 referenced_table = self.schema.get_table_by_column(referenced_column)
614 key_id = self._get_key(constraint)
615 key_record = {
616 "key_id": key_id,
617 "from_table": self._get_table_name(table),
618 "target_table": self._get_table_name(referenced_table),
619 "description": constraint.description,
620 "utype": constraint.votable_utype,
621 }
622 self._insert("keys", key_record)
624 ###########################################################
625 # Handle key_columns table
626 ###########################################################
627 # Loop over the corresponding columns and referenced
628 # columns and insert a record for each pair. This is
629 # necessary for proper handling of composite keys.
630 for from_column_id, target_column_id in zip(
631 constraint.columns, constraint.referenced_columns
632 ):
633 from_column = self.schema.find_object_by_id(from_column_id, datamodel.Column)
634 target_column = self.schema.find_object_by_id(target_column_id, datamodel.Column)
635 key_columns_record = {
636 "key_id": key_id,
637 "from_column": from_column.name,
638 "target_column": target_column.name,
639 }
640 self._insert("key_columns", key_columns_record)
642 def _generate_all_inserts(self) -> None:
643 """Generate the inserts for all the data."""
644 self.inserts.clear()
645 self._insert_schemas()
646 self._insert_tables()
647 self._insert_columns()
648 self._insert_keys()
649 logger.debug("Generated %d insert statements", len(self.inserts))
651 def _execute_inserts(self) -> None:
652 """Load the `~felis.datamodel.Schema` data into the TAP_SCHEMA
653 tables.
654 """
655 if isinstance(self.engine, Engine):
656 with self.engine.connect() as connection:
657 transaction = connection.begin()
658 try:
659 for insert in self.inserts:
660 connection.execute(insert)
661 transaction.commit()
662 except Exception as e:
663 logger.error("Error loading data into database: %s", e)
664 transaction.rollback()
665 raise
667 def _compiled_inserts(self) -> list[str]:
668 """Compile the inserts to SQL.
670 Returns
671 -------
672 list
673 A list of the compiled insert statements.
674 """
675 return [
676 str(insert.compile(self.engine, compile_kwargs={"literal_binds": True}))
677 for insert in self.inserts
678 ]
680 def _print_sql(self) -> None:
681 """Print the generated inserts to stdout."""
682 for insert_str in self._compiled_inserts():
683 print(insert_str + ";")
685 def _write_sql_to_file(self) -> None:
686 """Write the generated insert statements to a file."""
687 if not self.output_path:
688 raise ValueError("No output path specified")
689 with open(self.output_path, "w") as outfile:
690 for insert_str in self._compiled_inserts():
691 outfile.write(insert_str + ";" + "\n")
693 def _insert(self, table_name: str, record: list[Any] | dict[str, Any]) -> None:
694 """Generate an insert statement for a record.
696 Parameters
697 ----------
698 table_name
699 The name of the table to insert the record into.
700 record
701 The record to insert into the table.
702 """
703 table = self.mgr[table_name]
704 insert_statement = table.insert().values(record)
705 self.inserts.append(insert_statement)
707 @staticmethod
708 def _get_size(column: datamodel.Column) -> int | None:
709 """Get the size of the column.
711 Parameters
712 ----------
713 column
714 The column to get the size for.
716 Returns
717 -------
718 int or None
719 The size of the column or None if not applicable.
720 """
721 arraysize = column.votable_arraysize
723 if not arraysize:
724 return None
726 arraysize_str = str(arraysize)
727 if arraysize_str.isdigit():
728 return int(arraysize_str)
730 match = re.match(r"^([0-9]+)\*$", arraysize_str)
731 if match and match.group(1) is not None:
732 return int(match.group(1))
734 return None
736 @staticmethod
737 def _is_indexed(column: datamodel.Column, table: datamodel.Table) -> int:
738 """Check if the column is indexed in the table.
740 Parameters
741 ----------
742 column
743 The column to check.
744 table
745 The table to check.
747 Returns
748 -------
749 int
750 1 if the column is indexed, 0 otherwise.
751 """
752 if isinstance(table.primary_key, str) and table.primary_key == column.id:
753 return 1
754 for index in table.indexes:
755 if index.columns and len(index.columns) == 1 and index.columns[0] == column.id:
756 return 1
757 return 0
760class MetadataInserter:
761 """Insert TAP_SCHEMA self-description rows into the database.
763 Parameters
764 ----------
765 mgr
766 The table manager that contains the TAP_SCHEMA tables.
767 engine
768 The engine for connecting to the TAP_SCHEMA database.
769 """
771 def __init__(self, mgr: TableManager, engine: Engine):
772 """Initialize the metadata inserter.
774 Parameters
775 ----------
776 mgr
777 The table manager representing the TAP_SCHEMA tables.
778 engine
779 The SQLAlchemy engine for connecting to the database.
780 """
781 self._mgr = mgr
782 self._engine = engine
784 def insert_metadata(self) -> None:
785 """Insert the TAP_SCHEMA metadata into the database."""
786 for table_name in self._mgr.get_table_names_std():
787 table = self._mgr[table_name]
788 csv_bytes = ResourcePath(f"resource://felis/config/tap_schema/{table_name}.csv").read()
789 text_stream = io.TextIOWrapper(io.BytesIO(csv_bytes), encoding="utf-8")
790 reader = csv.reader(text_stream)
791 headers = next(reader)
792 rows = [
793 {key: None if value == "\\N" else value for key, value in zip(headers, row)} for row in reader
794 ]
795 logger.debug(
796 "Inserting %d rows into table '%s' with headers: %s",
797 len(rows),
798 table_name,
799 headers,
800 )
801 with self._engine.begin() as conn:
802 conn.execute(table.insert(), rows)