lsst-felis 26.2024.1100__py3-none-any.whl → 26.2024.1400__py3-none-any.whl

felis/cli.py

@@ -26,7 +26,7 @@ import json
 import logging
 import sys
 from collections.abc import Iterable, Mapping, MutableMapping
-from typing import Any
+from typing import IO, Any
 
 import click
 import yaml
@@ -38,7 +38,7 @@ from sqlalchemy.engine.mock import MockConnection
 from . import DEFAULT_CONTEXT, DEFAULT_FRAME, __version__
 from .check import CheckingVisitor
 from .datamodel import Schema
-from .sql import SQLVisitor
+from .metadata import DatabaseContext, InsertDump, MetaDataBuilder
 from .tap import Tap11Base, TapLoadingVisitor, init_tables
 from .utils import ReorderingVisitor
 from .validation import get_schema
@@ -71,27 +71,70 @@ def cli(log_level: str, log_file: str | None) -> None:
         logging.basicConfig(level=log_level)
 
 
-@cli.command("create-all")
-@click.option("--engine-url", envvar="ENGINE_URL", help="SQLAlchemy Engine URL")
-@click.option("--schema-name", help="Alternate Schema Name for Felis File")
-@click.option("--dry-run", is_flag=True, help="Dry Run Only. Prints out the DDL that would be executed")
+@cli.command("create")
+@click.option("--engine-url", envvar="ENGINE_URL", help="SQLAlchemy Engine URL", default="sqlite://")
+@click.option("--schema-name", help="Alternate schema name to override Felis file")
+@click.option(
+    "--create-if-not-exists", is_flag=True, help="Create the schema in the database if it does not exist"
+)
+@click.option("--drop-if-exists", is_flag=True, help="Drop schema if it already exists in the database")
+@click.option("--echo", is_flag=True, help="Echo database commands as they are executed")
+@click.option("--dry-run", is_flag=True, help="Dry run only to print out commands instead of executing")
+@click.option(
+    "--output-file", "-o", type=click.File(mode="w"), help="Write SQL commands to a file instead of executing"
+)
 @click.argument("file", type=click.File())
-def create_all(engine_url: str, schema_name: str, dry_run: bool, file: io.TextIOBase) -> None:
-    """Create schema objects from the Felis FILE."""
-    schema_obj = yaml.load(file, Loader=yaml.SafeLoader)
-    visitor = SQLVisitor(schema_name=schema_name)
-    schema = visitor.visit_schema(schema_obj)
-
-    metadata = schema.metadata
+def create(
+    engine_url: str,
+    schema_name: str | None,
+    create_if_not_exists: bool,
+    drop_if_exists: bool,
+    echo: bool,
+    dry_run: bool,
+    output_file: IO[str] | None,
+    file: IO,
+) -> None:
+    """Create database objects from the Felis file."""
+    yaml_data = yaml.safe_load(file)
+    schema = Schema.model_validate(yaml_data)
+    url_obj = make_url(engine_url)
+    if schema_name:
+        logger.info(f"Overriding schema name with: {schema_name}")
+        schema.name = schema_name
+    elif url_obj.drivername == "sqlite":
+        logger.info("Overriding schema name for sqlite with: main")
+        schema.name = "main"
+    if not url_obj.host and not url_obj.drivername == "sqlite":
+        dry_run = True
+        logger.info("Forcing dry run for non-sqlite engine URL with no host")
+
+    builder = MetaDataBuilder(schema)
+    builder.build()
+    metadata = builder.metadata
+    logger.debug(f"Created metadata with schema name: {metadata.schema}")
 
     engine: Engine | MockConnection
-    if not dry_run:
-        engine = create_engine(engine_url)
+    if not dry_run and not output_file:
+        engine = create_engine(engine_url, echo=echo)
     else:
-        _insert_dump = InsertDump()
-        engine = create_mock_engine(make_url(engine_url), executor=_insert_dump.dump)
-        _insert_dump.dialect = engine.dialect
-    metadata.create_all(engine)
+        if dry_run:
+            logger.info("Dry run will be executed")
+        engine = DatabaseContext.create_mock_engine(url_obj, output_file)
+        if output_file:
+            logger.info("Writing SQL output to: " + output_file.name)
+
+    context = DatabaseContext(metadata, engine)
+
+    if drop_if_exists:
+        logger.debug("Dropping schema if it exists")
+        context.drop_if_exists()
+        create_if_not_exists = True  # If schema is dropped, it needs to be recreated.
+
+    if create_if_not_exists:
+        logger.debug("Creating schema if not exists")
+        context.create_if_not_exists()
+
+    context.create_all()
 
 
 @cli.command("init-tap")
@@ -404,30 +447,5 @@ def _normalize(schema_obj: Mapping[str, Any], embed: str = "@last") -> MutableMa
     return compacted
 
 
-class InsertDump:
-    """An Insert Dumper for SQL statements."""
-
-    dialect: Any = None
-
-    def dump(self, sql: Any, *multiparams: Any, **params: Any) -> None:
-        compiled = sql.compile(dialect=self.dialect)
-        sql_str = str(compiled) + ";"
-        params_list = [compiled.params]
-        for params in params_list:
-            if not params:
-                print(sql_str)
-                continue
-            new_params = {}
-            for key, value in params.items():
-                if isinstance(value, str):
-                    new_params[key] = f"'{value}'"
-                elif value is None:
-                    new_params[key] = "null"
-                else:
-                    new_params[key] = value
-
-            print(sql_str % new_params)
-
-
 if __name__ == "__main__":
     cli()

felis/datamodel.py

@@ -31,7 +31,6 @@ from astropy.io.votable import ucd  # type: ignore
 from pydantic import BaseModel, ConfigDict, Field, field_validator, model_validator
 
 logger = logging.getLogger(__name__)
-# logger.setLevel(logging.DEBUG)
 
 __all__ = (
     "BaseObject",
@@ -50,7 +49,6 @@ __all__ = (
 CONFIG = ConfigDict(
     populate_by_name=True,  # Populate attributes by name.
     extra="forbid",  # Do not allow extra fields.
-    use_enum_values=True,  # Use enum values instead of names.
     validate_assignment=True,  # Validate assignments after model is created.
     str_strip_whitespace=True,  # Strip whitespace from string fields.
 )
@@ -130,8 +128,13 @@ class Column(BaseObject):
     length: int | None = None
     """The length of the column."""
 
-    nullable: bool = True
-    """Whether the column can be `NULL`."""
+    nullable: bool | None = None
+    """Whether the column can be ``NULL``.
+
+    If `None`, this value was not set explicitly in the YAML data. In this
+    case, it will be set to `False` for columns with numeric types and `True`
+    otherwise.
+    """
 
     value: Any = None
     """The default value of the column."""
@@ -142,6 +145,9 @@ class Column(BaseObject):
     mysql_datatype: str | None = Field(None, alias="mysql:datatype")
     """The MySQL datatype of the column."""
 
+    postgresql_datatype: str | None = Field(None, alias="postgresql:datatype")
+    """The PostgreSQL datatype of the column."""
+
     ivoa_ucd: str | None = Field(None, alias="ivoa:ucd")
     """The IVOA UCD of the column."""
 
@@ -280,7 +286,7 @@ class Table(BaseObject):
     indexes: list[Index] = Field(default_factory=list)
     """The indexes on the table."""
 
-    primaryKey: str | list[str] | None = None
+    primary_key: str | list[str] | None = Field(None, alias="primaryKey")
     """The primary key of the table."""
 
     tap_table_index: int | None = Field(None, alias="tap:table_index")
@@ -427,6 +433,9 @@ class Schema(BaseObject):
     @model_validator(mode="after")
     def create_id_map(self: Schema) -> Schema:
         """Create a map of IDs to objects."""
+        if len(self.id_map):
+            logger.debug("ID map was already populated")
+            return self
         visitor: SchemaIdVisitor = SchemaIdVisitor()
         visitor.visit_schema(self)
         logger.debug(f"ID map contains {len(self.id_map.keys())} objects")

felis/db/_variants.py

@@ -0,0 +1,94 @@
+# This file is part of felis.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+import re
+from typing import Any
+
+from sqlalchemy import types
+from sqlalchemy.dialects import mysql, oracle, postgresql, sqlite
+from sqlalchemy.types import TypeEngine
+
+from ..datamodel import Column
+
+MYSQL = "mysql"
+ORACLE = "oracle"
+POSTGRES = "postgresql"
+SQLITE = "sqlite"
+
+TABLE_OPTS = {
+    "mysql:engine": "mysql_engine",
+    "mysql:charset": "mysql_charset",
+    "oracle:compress": "oracle_compress",
+}
+
+COLUMN_VARIANT_OVERRIDE = {
+    "mysql:datatype": "mysql",
+    "oracle:datatype": "oracle",
+    "postgresql:datatype": "postgresql",
+    "sqlite:datatype": "sqlite",
+}
+
+DIALECT_MODULES = {MYSQL: mysql, ORACLE: oracle, SQLITE: sqlite, POSTGRES: postgresql}
+
+_length_regex = re.compile(r"\((\d+)\)")
+"""A regular expression that is looking for numbers within parentheses."""
+
+
+def process_variant_override(dialect_name: str, variant_override_str: str) -> types.TypeEngine:
+    """Return variant type for given dialect."""
+    dialect = DIALECT_MODULES[dialect_name]
+    variant_type_name = variant_override_str.split("(")[0]
+
+    # Process Variant Type
+    if variant_type_name not in dir(dialect):
+        raise ValueError(f"Type {variant_type_name} not found in dialect {dialect_name}")
+    variant_type = getattr(dialect, variant_type_name)
+    length_params = []
+    if match := _length_regex.search(variant_override_str):
+        length_params.extend([int(i) for i in match.group(1).split(",")])
+    return variant_type(*length_params)
+
+
+def make_variant_dict(column_obj: Column) -> dict[str, TypeEngine[Any]]:
+    """Handle variant overrides for a `felis.datamodel.Column`.
+
+    This function will return a dictionary of `str` to
+    `sqlalchemy.types.TypeEngine` containing variant datatype information
+    (e.g., for mysql, postgresql, etc).
+
+    Parameters
+    ----------
+    column_obj : `felis.datamodel.Column`
+        The column object from which to build the variant dictionary.
+
+    Returns
+    -------
+    variant_dict : `dict`
+        The dictionary of `str` to `sqlalchemy.types.TypeEngine` containing
+        variant datatype information (e.g., for mysql, postgresql, etc).
+    """
+    variant_dict = {}
+    for field_name, value in iter(column_obj):
+        if field_name in COLUMN_VARIANT_OVERRIDE:
+            dialect = COLUMN_VARIANT_OVERRIDE[field_name]
+            variant: TypeEngine = process_variant_override(dialect, value)
+            variant_dict[dialect] = variant
+    return variant_dict

felis/metadata.py

@@ -0,0 +1,504 @@
+# This file is part of felis.
+#
+# Developed for the LSST Data Management System.
+# This product includes software developed by the LSST Project
+# (https://www.lsst.org).
+# See the COPYRIGHT file at the top-level directory of this distribution
+# for details of code ownership.
+#
+# This program is free software: you can redistribute it and/or modify
+# it under the terms of the GNU General Public License as published by
+# the Free Software Foundation, either version 3 of the License, or
+# (at your option) any later version.
+#
+# This program is distributed in the hope that it will be useful,
+# but WITHOUT ANY WARRANTY; without even the implied warranty of
+# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
+# GNU General Public License for more details.
+#
+# You should have received a copy of the GNU General Public License
+# along with this program.  If not, see <https://www.gnu.org/licenses/>.
+
+from __future__ import annotations
+
+import logging
+from typing import IO, Any, Literal
+
+import sqlalchemy.schema as sqa_schema
+from lsst.utils.iteration import ensure_iterable
+from sqlalchemy import (
+    CheckConstraint,
+    Column,
+    Constraint,
+    Engine,
+    ForeignKeyConstraint,
+    Index,
+    MetaData,
+    Numeric,
+    PrimaryKeyConstraint,
+    ResultProxy,
+    Table,
+    UniqueConstraint,
+    create_mock_engine,
+    make_url,
+    text,
+)
+from sqlalchemy.engine.interfaces import Dialect
+from sqlalchemy.engine.mock import MockConnection
+from sqlalchemy.engine.url import URL
+from sqlalchemy.exc import SQLAlchemyError
+from sqlalchemy.types import TypeEngine
+
+from felis.datamodel import Schema
+from felis.db._variants import make_variant_dict
+
+from . import datamodel
+from .db import sqltypes
+from .types import FelisType
+
+logger = logging.getLogger(__name__)
+
+
+class InsertDump:
+    """An Insert Dumper for SQL statements which supports writing messages
+    to stdout or a file.
+    """
+
+    def __init__(self, file: IO[str] | None = None) -> None:
+        """Initialize the insert dumper.
+
+        Parameters
+        ----------
+        file : `io.TextIOBase` or `None`, optional
+            The file to write the SQL statements to. If None, the statements
+            will be written to stdout.
+        """
+        self.file = file
+        self.dialect: Dialect | None = None
+
+    def dump(self, sql: Any, *multiparams: Any, **params: Any) -> None:
+        """Dump the SQL statement to a file or stdout.
+
+        Statements with parameters will be formatted with the values
+        inserted into the resultant SQL output.
+
+        Parameters
+        ----------
+        sql : `typing.Any`
+            The SQL statement to dump.
+        multiparams : `typing.Any`
+            The multiparams to use for the SQL statement.
+        params : `typing.Any`
+            The params to use for the SQL statement.
+        """
+        compiled = sql.compile(dialect=self.dialect)
+        sql_str = str(compiled) + ";"
+        params_list = [compiled.params]
+        for params in params_list:
+            if not params:
+                print(sql_str, file=self.file)
+                continue
+            new_params = {}
+            for key, value in params.items():
+                if isinstance(value, str):
+                    new_params[key] = f"'{value}'"
+                elif value is None:
+                    new_params[key] = "null"
+                else:
+                    new_params[key] = value
+            print(sql_str % new_params, file=self.file)
+
+
+def get_datatype_with_variants(column_obj: datamodel.Column) -> TypeEngine:
+    """Use the Felis type system to get a SQLAlchemy datatype with variant
+    overrides from the information in a `Column` object.
+
+    Parameters
+    ----------
+    column_obj : `felis.datamodel.Column`
+        The column object from which to get the datatype.
+
+    Raises
+    ------
+    ValueError
+        If the column has a sized type but no length.
+    """
+    variant_dict = make_variant_dict(column_obj)
+    felis_type = FelisType.felis_type(column_obj.datatype.value)
+    datatype_fun = getattr(sqltypes, column_obj.datatype.value)
+    if felis_type.is_sized:
+        if not column_obj.length:
+            raise ValueError(f"Column {column_obj.name} has sized type '{column_obj.datatype}' but no length")
+        datatype = datatype_fun(column_obj.length, **variant_dict)
+    else:
+        datatype = datatype_fun(**variant_dict)
+    return datatype
+
+
+class MetaDataBuilder:
+    """A class for building a `MetaData` object from a Felis `Schema`."""
+
+    def __init__(
+        self, schema: Schema, apply_schema_to_metadata: bool = True, apply_schema_to_tables: bool = True
+    ) -> None:
+        """Initialize the metadata builder.
+
+        Parameters
+        ----------
+        schema : `felis.datamodel.Schema`
+            The schema object from which to build the SQLAlchemy metadata.
+        apply_schema_to_metadata : `bool`, optional
+            Whether to apply the schema name to the metadata object.
+        apply_schema_to_tables : `bool`, optional
+            Whether to apply the schema name to the tables.
+        """
+        self.schema = schema
+        if not apply_schema_to_metadata:
+            logger.debug("Schema name will not be applied to metadata")
+        if not apply_schema_to_tables:
+            logger.debug("Schema name will not be applied to tables")
+        self.metadata = MetaData(schema=schema.name if apply_schema_to_metadata else None)
+        self._objects: dict[str, Any] = {}
+        self.apply_schema_to_tables = apply_schema_to_tables
+
+    def build(self) -> MetaData:
+        """Build the SQLAlchemy tables and constraints from the schema."""
+        self.build_tables()
+        self.build_constraints()
+        return self.metadata
+
+    def build_tables(self) -> None:
+        """Build the SQLAlchemy tables from the schema.
+
+        Notes
+        -----
+        This function builds all the tables by calling ``build_table`` on
+        each Pydantic object. It also calls ``build_primary_key`` to create the
+        primary key constraints.
+        """
+        for table in self.schema.tables:
+            self.build_table(table)
+            if table.primary_key:
+                primary_key = self.build_primary_key(table.primary_key)
+                self._objects[table.id].append_constraint(primary_key)
+
+    def build_primary_key(self, primary_key_columns: str | list[str]) -> PrimaryKeyConstraint:
+        """Build a SQLAlchemy `PrimaryKeyConstraint` from a single column ID
+        or a list.
+
+        The `primary_key_columns` are strings or a list of strings representing
+        IDs pointing to columns that will be looked up in the internal object
+        dictionary.
+
+        Parameters
+        ----------
+        primary_key_columns : `str` or `list` of `str`
+            The column ID or list of column IDs from which to build the primary
+            key.
+
+        Returns
+        -------
+        primary_key: `sqlalchemy.PrimaryKeyConstraint`
+            The SQLAlchemy primary key constraint object.
+        """
+        return PrimaryKeyConstraint(
+            *[self._objects[column_id] for column_id in ensure_iterable(primary_key_columns)]
+        )
+
+    def build_table(self, table_obj: datamodel.Table) -> None:
+        """Build a `sqlalchemy.Table` from a `felis.datamodel.Table` and add
+        it to the `sqlalchemy.MetaData` object.
+
+        Several MySQL table options are handled by annotations on the table,
+        including the engine and charset. This is not needed for Postgres,
+        which does not have equivalent options.
+
+        Parameters
+        ----------
+        table_obj : `felis.datamodel.Table`
+            The table object to build the SQLAlchemy table from.
+        """
+        # Process mysql table options.
+        optargs = {}
+        if table_obj.mysql_engine:
+            optargs["mysql_engine"] = table_obj.mysql_engine
+        if table_obj.mysql_charset:
+            optargs["mysql_charset"] = table_obj.mysql_charset
+
+        # Create the SQLAlchemy table object and its columns.
+        name = table_obj.name
+        id = table_obj.id
+        description = table_obj.description
+        columns = [self.build_column(column) for column in table_obj.columns]
+        table = Table(
+            name,
+            self.metadata,
+            *columns,
+            comment=description,
+            schema=self.schema.name if self.apply_schema_to_tables else None,
+            **optargs,  # type: ignore[arg-type]
+        )
+
+        # Create the indexes and add them to the table.
+        indexes = [self.build_index(index) for index in table_obj.indexes]
+        for index in indexes:
+            index._set_parent(table)
+            table.indexes.add(index)
+
+        self._objects[id] = table
+
+    def build_column(self, column_obj: datamodel.Column) -> Column:
+        """Build a SQLAlchemy column from a `felis.datamodel.Column` object.
+
+        Parameters
+        ----------
+        column_obj : `felis.datamodel.Column`
+            The column object from which to build the SQLAlchemy column.
+
+        Returns
+        -------
+        column: `sqlalchemy.Column`
+            The SQLAlchemy column object.
+        """
+        # Get basic column attributes.
+        name = column_obj.name
+        id = column_obj.id
+        description = column_obj.description
+        default = column_obj.value
+
+        # Handle variant overrides for the column (e.g., "mysql:datatype").
+        datatype = get_datatype_with_variants(column_obj)
+
+        # Set default value of nullable based on column type and then whether
+        # it was explicitly provided in the schema data.
+        nullable = column_obj.nullable
+        if nullable is None:
+            nullable = False if isinstance(datatype, Numeric) else True
+
+        # Set autoincrement depending on if it was provided explicitly.
+        autoincrement: Literal["auto"] | bool = (
+            column_obj.autoincrement if column_obj.autoincrement is not None else "auto"
+        )
+
+        column: Column = Column(
+            name,
+            datatype,
+            comment=description,
+            autoincrement=autoincrement,
+            nullable=nullable,
+            server_default=default,
+        )
+
+        self._objects[id] = column
+
+        return column
+
+    def build_constraints(self) -> None:
+        """Build the SQLAlchemy constraints in the Felis schema and append them
+        to the associated `Table`.
+
+        Notes
+        -----
+        This is performed as a separate step after building the tables so that
+        all the referenced objects in the constraints will be present and can
+        be looked up by their ID.
+        """
+        for table_obj in self.schema.tables:
+            table = self._objects[table_obj.id]
+            for constraint_obj in table_obj.constraints:
+                constraint = self.build_constraint(constraint_obj)
+                table.append_constraint(constraint)
+
+    def build_constraint(self, constraint_obj: datamodel.Constraint) -> Constraint:
+        """Build a SQLAlchemy `Constraint` from a `felis.datamodel.Constraint`
+        object.
+
+        Parameters
+        ----------
+        constraint_obj : `felis.datamodel.Constraint`
+            The constraint object from which to build the SQLAlchemy
+            constraint.
+
+        Returns
+        -------
+        constraint: `sqlalchemy.Constraint`
+            The SQLAlchemy constraint object.
+
+        Raises
+        ------
+        ValueError
+            If the constraint type is not recognized.
+        TypeError
+            If the constraint object is not the expected type.
+        """
+        args: dict[str, Any] = {
+            "name": constraint_obj.name or None,
+            "info": constraint_obj.description or None,
+            "deferrable": constraint_obj.deferrable or None,
+            "initially": constraint_obj.initially or None,
+        }
+        constraint: Constraint
+        constraint_type = constraint_obj.type
+
+        if isinstance(constraint_obj, datamodel.ForeignKeyConstraint):
+            fk_obj: datamodel.ForeignKeyConstraint = constraint_obj
+            columns = [self._objects[column_id] for column_id in fk_obj.columns]
+            refcolumns = [self._objects[column_id] for column_id in fk_obj.referenced_columns]
+            constraint = ForeignKeyConstraint(columns, refcolumns, **args)
+        elif isinstance(constraint_obj, datamodel.CheckConstraint):
+            check_obj: datamodel.CheckConstraint = constraint_obj
+            expression = check_obj.expression
+            constraint = CheckConstraint(expression, **args)
+        elif isinstance(constraint_obj, datamodel.UniqueConstraint):
+            uniq_obj: datamodel.UniqueConstraint = constraint_obj
+            columns = [self._objects[column_id] for column_id in uniq_obj.columns]
+            constraint = UniqueConstraint(*columns, **args)
+        else:
+            raise ValueError(f"Unknown constraint type: {constraint_type}")
+
+        self._objects[constraint_obj.id] = constraint
+
+        return constraint
+
+    def build_index(self, index_obj: datamodel.Index) -> Index:
+        """Build a SQLAlchemy `Index` from a `felis.datamodel.Index` object.
+
+        Parameters
+        ----------
+        index_obj : `felis.datamodel.Index`
+            The index object from which to build the SQLAlchemy index.
+
+        Returns
+        -------
+        index: `sqlalchemy.Index`
+            The SQLAlchemy index object.
+        """
+        columns = [self._objects[c_id] for c_id in (index_obj.columns if index_obj.columns else [])]
+        expressions = index_obj.expressions if index_obj.expressions else []
+        index = Index(index_obj.name, *columns, *expressions)
+        self._objects[index_obj.id] = index
+        return index
+
+
+class ConnectionWrapper:
+    """A wrapper for a SQLAlchemy engine or mock connection which provides a
+    consistent interface for executing SQL statements.
+    """
+
+    def __init__(self, engine: Engine | MockConnection):
+        """Initialize the connection wrapper.
+
+        Parameters
+        ----------
+        engine : `sqlalchemy.Engine` or `sqlalchemy.MockConnection`
+            The SQLAlchemy engine or mock connection to wrap.
+        """
+        self.engine = engine
+
+    def execute(self, statement: Any) -> ResultProxy:
+        """Execute a SQL statement on the engine and return the result."""
+        if isinstance(statement, str):
+            statement = text(statement)
+        if isinstance(self.engine, MockConnection):
+            return self.engine.connect().execute(statement)
+        else:
+            with self.engine.begin() as connection:
+                result = connection.execute(statement)
+                return result
+
+
+class DatabaseContext:
+    """A class for managing the schema and its database connection."""
+
+    def __init__(self, metadata: MetaData, engine: Engine | MockConnection):
+        """Initialize the database context.
+
+        Parameters
+        ----------
+        metadata : `sqlalchemy.MetaData`
+            The SQLAlchemy metadata object.
+
+        engine : `sqlalchemy.Engine` or `sqlalchemy.MockConnection`
+            The SQLAlchemy engine or mock connection object.
+        """
+        self.engine = engine
+        self.metadata = metadata
+        self.connection = ConnectionWrapper(engine)
+
+    def create_if_not_exists(self) -> None:
+        """Create the schema in the database if it does not exist.
+
+        In MySQL, this will create a new database. In PostgreSQL, it will
+        create a new schema. For other variants, this is an unsupported
+        operation.
+
+        Parameters
+        ----------
+        engine: `sqlalchemy.Engine`
+            The SQLAlchemy engine object.
+        schema_name: `str`
+            The name of the schema (or database) to create.
+        """
+        db_type = self.engine.dialect.name
+        schema_name = self.metadata.schema
+        try:
+            if db_type == "mysql":
+                logger.info(f"Creating MySQL database: {schema_name}")
+                self.connection.execute(text(f"CREATE DATABASE IF NOT EXISTS {schema_name}"))
+            elif db_type == "postgresql":
+                logger.info(f"Creating PG schema: {schema_name}")
+                self.connection.execute(sqa_schema.CreateSchema(schema_name, if_not_exists=True))
+            else:
+                raise ValueError("Unsupported database type:" + db_type)
+        except SQLAlchemyError as e:
+            logger.error(f"Error creating schema: {e}")
+            raise
+
+    def drop_if_exists(self) -> None:
+        """Drop the schema in the database if it exists.
+
+        In MySQL, this will drop a database. In PostgreSQL, it will drop a
+        schema. For other variants, this is unsupported for now.
+
+        Parameters
+        ----------
+        engine: `sqlalchemy.Engine`
+            The SQLAlchemy engine object.
+        schema_name: `str`
+            The name of the schema (or database) to drop.
+        """
+        db_type = self.engine.dialect.name
+        schema_name = self.metadata.schema
+        try:
+            if db_type == "mysql":
+                logger.info(f"Dropping MySQL database if exists: {schema_name}")
+                self.connection.execute(text(f"DROP DATABASE IF EXISTS {schema_name}"))
+            elif db_type == "postgresql":
+                logger.info(f"Dropping PostgreSQL schema if exists: {schema_name}")
+                self.connection.execute(sqa_schema.DropSchema(schema_name, if_exists=True))
+            else:
+                raise ValueError(f"Unsupported database type: {db_type}")
+        except SQLAlchemyError as e:
+            logger.error(f"Error dropping schema: {e}")
+            raise
+
+    def create_all(self) -> None:
+        """Create all tables in the schema using the metadata object."""
+        self.metadata.create_all(self.engine)
+
+    @staticmethod
+    def create_mock_engine(engine_url: URL, output_file: IO[str] | None = None) -> MockConnection:
+        """Create a mock engine for testing or dumping DDL statements.
+
+        Parameters
+        ----------
+        engine_url : `sqlalchemy.engine.url.URL`
+            The SQLAlchemy engine URL.
+        output_file : `typing.IO` [ `str` ] or `None`, optional
+            The file to write the SQL statements to. If None, the statements
+            will be written to stdout.
+        """
+        dumper = InsertDump(output_file)
+        engine = create_mock_engine(make_url(engine_url), executor=dumper.dump)
+        dumper.dialect = engine.dialect
+        return engine

felis/version.py

@@ -1,2 +1,2 @@
 __all__ = ["__version__"]
-__version__ = "26.2024.1100"
+__version__ = "26.2024.1400"

{lsst_felis-26.2024.1100.dist-info → lsst_felis-26.2024.1400.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lsst-felis
-Version: 26.2024.1100
+Version: 26.2024.1400
 Summary: A vocabulary for describing catalogs and acting on those descriptions
 Author-email: Rubin Observatory Data Management <dm-admin@lists.lsst.org>
 License: GNU General Public License v3 or later (GPLv3+)
@@ -23,6 +23,7 @@ Requires-Dist: click >=7
 Requires-Dist: pyyaml >=6
 Requires-Dist: pyld >=2
 Requires-Dist: pydantic <3,>=2
+Requires-Dist: lsst-utils
 Provides-Extra: test
 Requires-Dist: pytest >=3.2 ; extra == 'test'
 

{lsst_felis-26.2024.1100.dist-info → lsst_felis-26.2024.1400.dist-info}/RECORD

@@ -1,23 +1,24 @@
 felis/__init__.py,sha256=_Pw-QKMYj0WRgE8fW2N2pBXJUj-Pjv8dSKJBzykjyZU,1842
 felis/check.py,sha256=RBxXq7XwPGIucrs1PPgPtgk8MrWAJlOmoxCNySEz9-I,13892
-felis/cli.py,sha256=7lxVM46uaHdBhlNmIGz9PKleI0jrx6VswVQdAv2v-mM,15785
-felis/datamodel.py,sha256=6i_BtmU1c5RaIZxeA-XbnHkikJb-fnxSWrsX0zdVwwU,15648
+felis/cli.py,sha256=YeGSiA3ywPVMMdB1YxH1_Gdac1kl4oPJvJtajfCs5VU,16637
+felis/datamodel.py,sha256=ooNSg68OuNk89EVu1MtxupLUWgSyzmb50wjza1joDO4,16002
+felis/metadata.py,sha256=5DE2YMnu6YuhwntBSe-OheCD7C2-vA4yb64BpjTC68A,18542
 felis/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 felis/simple.py,sha256=yzv_aoZrZhfakd1Xm7gLDeVKyJjCDZ7wAyYYp-l_Sxs,14414
-felis/sql.py,sha256=Kq4Y7XxXDUTrzz8eIR5VLQFq0n899uZlcjIyorZqZ70,10293
 felis/tap.py,sha256=RBwEKyU3S0oXSNIMoI2WRAuC9WB5eai9BdQQUYN5Qdc,17704
 felis/types.py,sha256=1GL6IkHcIsIydiyw1eX98REh-lWCVIRO-9qjmaZfqvw,4440
 felis/utils.py,sha256=tYxr0xFdPN4gDHibeAD9d5DFgU8hKlSZVKmZoDzi4e8,4164
 felis/validation.py,sha256=f9VKvp7q-cnim2D5voTKwCdt0NRsYBpTwom1Z_3OKkc,3469
-felis/version.py,sha256=xrLeMJU99wLKp0EMd5_cq5lPeP_4FfIT_M5Mei_Lydc,55
+felis/version.py,sha256=Bkf7vQojVYVWgjqmEFJRzUA7IlOUtWz7QH7SQ12F5YA,55
 felis/visitor.py,sha256=EazU4nYbkKBj3mCZYvsTCBTNmh0qRaUNZIzCcM3dqOQ,6439
 felis/db/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+felis/db/_variants.py,sha256=aW0Q7R4KEtxLR7VMashQjDLWdzDNrMVAH521MSvMey0,3346
 felis/db/sqltypes.py,sha256=0HOEqvL0OailGP-j6Jj5tnOSu_Pt7Hi29PPof4Q5d2c,5787
-lsst_felis-26.2024.1100.dist-info/COPYRIGHT,sha256=bUmNy19uUxqITMpjeHFe69q3IzQpjxvvBw6oV7kR7ho,129
-lsst_felis-26.2024.1100.dist-info/LICENSE,sha256=jOtLnuWt7d5Hsx6XXB2QxzrSe2sWWh3NgMfFRetluQM,35147
-lsst_felis-26.2024.1100.dist-info/METADATA,sha256=T6zHQB8sB6wuj48oVXRsJBn4LBgU-cKoj7sDb7eZPI4,1075
-lsst_felis-26.2024.1100.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-lsst_felis-26.2024.1100.dist-info/entry_points.txt,sha256=Gk2XFujA_Gp52VBk45g5kim8TDoMDJFPctsMqiq72EM,40
-lsst_felis-26.2024.1100.dist-info/top_level.txt,sha256=F4SvPip3iZRVyISi50CHhwTIAokAhSxjWiVcn4IVWRI,6
-lsst_felis-26.2024.1100.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
-lsst_felis-26.2024.1100.dist-info/RECORD,,
+lsst_felis-26.2024.1400.dist-info/COPYRIGHT,sha256=bUmNy19uUxqITMpjeHFe69q3IzQpjxvvBw6oV7kR7ho,129
+lsst_felis-26.2024.1400.dist-info/LICENSE,sha256=jOtLnuWt7d5Hsx6XXB2QxzrSe2sWWh3NgMfFRetluQM,35147
+lsst_felis-26.2024.1400.dist-info/METADATA,sha256=y4cVRA28EHR2jiWXFDDzSRCm16TH4YnJrept5w3fkMM,1101
+lsst_felis-26.2024.1400.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+lsst_felis-26.2024.1400.dist-info/entry_points.txt,sha256=Gk2XFujA_Gp52VBk45g5kim8TDoMDJFPctsMqiq72EM,40
+lsst_felis-26.2024.1400.dist-info/top_level.txt,sha256=F4SvPip3iZRVyISi50CHhwTIAokAhSxjWiVcn4IVWRI,6
+lsst_felis-26.2024.1400.dist-info/zip-safe,sha256=AbpHGcgLb-kRsJGnwFEktk7uzpZOCcBY74-YBdrKVGs,1
+lsst_felis-26.2024.1400.dist-info/RECORD,,

felis/sql.py

@@ -1,275 +0,0 @@
-# This file is part of felis.
-#
-# Developed for the LSST Data Management System.
-# This product includes software developed by the LSST Project
-# (https://www.lsst.org).
-# See the COPYRIGHT file at the top-level directory of this distribution
-# for details of code ownership.
-#
-# This program is free software: you can redistribute it and/or modify
-# it under the terms of the GNU General Public License as published by
-# the Free Software Foundation, either version 3 of the License, or
-# (at your option) any later version.
-#
-# This program is distributed in the hope that it will be useful,
-# but WITHOUT ANY WARRANTY; without even the implied warranty of
-# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
-# GNU General Public License for more details.
-#
-# You should have received a copy of the GNU General Public License
-# along with this program.  If not, see <https://www.gnu.org/licenses/>.
-
-from __future__ import annotations
-
-__all__ = ["SQLVisitor"]
-
-import logging
-import re
-from collections.abc import Iterable, Mapping, MutableMapping
-from typing import Any, NamedTuple
-
-from sqlalchemy import (
-    CheckConstraint,
-    Column,
-    Constraint,
-    ForeignKeyConstraint,
-    Index,
-    MetaData,
-    Numeric,
-    PrimaryKeyConstraint,
-    UniqueConstraint,
-    types,
-)
-from sqlalchemy.dialects import mysql, oracle, postgresql, sqlite
-from sqlalchemy.schema import Table
-
-from .check import FelisValidator
-from .db import sqltypes
-from .types import FelisType
-from .visitor import Visitor
-
-_Mapping = Mapping[str, Any]
-_MutableMapping = MutableMapping[str, Any]
-
-logger = logging.getLogger("felis")
-
-MYSQL = "mysql"
-ORACLE = "oracle"
-POSTGRES = "postgresql"
-SQLITE = "sqlite"
-
-TABLE_OPTS = {
-    "mysql:engine": "mysql_engine",
-    "mysql:charset": "mysql_charset",
-    "oracle:compress": "oracle_compress",
-}
-
-COLUMN_VARIANT_OVERRIDE = {
-    "mysql:datatype": "mysql",
-    "oracle:datatype": "oracle",
-    "postgresql:datatype": "postgresql",
-    "sqlite:datatype": "sqlite",
-}
-
-DIALECT_MODULES = {MYSQL: mysql, ORACLE: oracle, SQLITE: sqlite, POSTGRES: postgresql}
-
-length_regex = re.compile(r"\((.+)\)")
-
-
-class Schema(NamedTuple):
-    name: str | None
-    tables: list[Table]
-    metadata: MetaData
-    graph_index: Mapping[str, Any]
-
-
-class SQLVisitor(Visitor[Schema, Table, Column, PrimaryKeyConstraint | None, Constraint, Index, None]):
-    """A Felis Visitor which populates a SQLAlchemy metadata object.
-
-    Parameters
-    ----------
-    schema_name : `str`, optional
-        Override for the schema name.
-    """
-
-    def __init__(self, schema_name: str | None = None):
-        self.metadata = MetaData()
-        self.schema_name = schema_name
-        self.checker = FelisValidator()
-        self.graph_index: MutableMapping[str, Any] = {}
-
-    def visit_schema(self, schema_obj: _Mapping) -> Schema:
-        # Docstring is inherited.
-        self.checker.check_schema(schema_obj)
-        if (version_obj := schema_obj.get("version")) is not None:
-            self.visit_schema_version(version_obj, schema_obj)
-
-        # Create tables but don't add constraints yet.
-        tables = [self.visit_table(t, schema_obj) for t in schema_obj["tables"]]
-
-        # Process constraints after the tables are created so that all
-        # referenced columns are available.
-        for table_obj in schema_obj["tables"]:
-            constraints = [
-                self.visit_constraint(constraint, table_obj)
-                for constraint in table_obj.get("constraints", [])
-            ]
-            table = self.graph_index[table_obj["@id"]]
-            for constraint in constraints:
-                table.append_constraint(constraint)
-
-        schema = Schema(
-            name=self.schema_name or schema_obj["name"],
-            tables=tables,
-            metadata=self.metadata,
-            graph_index=self.graph_index,
-        )
-        return schema
-
-    def visit_schema_version(
-        self, version_obj: str | Mapping[str, Any], schema_obj: Mapping[str, Any]
-    ) -> None:
-        # Docstring is inherited.
-
-        # For now we ignore schema versioning completely, still do some checks.
-        self.checker.check_schema_version(version_obj, schema_obj)
-
-    def visit_table(self, table_obj: _Mapping, schema_obj: _Mapping) -> Table:
-        # Docstring is inherited.
-        self.checker.check_table(table_obj, schema_obj)
-        columns = [self.visit_column(c, table_obj) for c in table_obj["columns"]]
-
-        name = table_obj["name"]
-        table_id = table_obj["@id"]
-        description = table_obj.get("description")
-        schema_name = self.schema_name or schema_obj["name"]
-
-        table = Table(name, self.metadata, *columns, schema=schema_name, comment=description)
-
-        primary_key = self.visit_primary_key(table_obj.get("primaryKey", []), table_obj)
-        if primary_key:
-            table.append_constraint(primary_key)
-
-        indexes = [self.visit_index(i, table_obj) for i in table_obj.get("indexes", [])]
-        for index in indexes:
-            # FIXME: Hack because there's no table.add_index
-            index._set_parent(table)
-            table.indexes.add(index)
-        self.graph_index[table_id] = table
-        return table
-
-    def visit_column(self, column_obj: _Mapping, table_obj: _Mapping) -> Column:
-        # Docstring is inherited.
-        self.checker.check_column(column_obj, table_obj)
-        column_name = column_obj["name"]
-        column_id = column_obj["@id"]
-        datatype_name = column_obj["datatype"]
-        column_description = column_obj.get("description")
-        column_default = column_obj.get("value")
-        column_length = column_obj.get("length")
-
-        kwargs = {}
-        for column_opt in column_obj.keys():
-            if column_opt in COLUMN_VARIANT_OVERRIDE:
-                dialect = COLUMN_VARIANT_OVERRIDE[column_opt]
-                variant = _process_variant_override(dialect, column_obj[column_opt])
-                kwargs[dialect] = variant
-
-        felis_type = FelisType.felis_type(datatype_name)
-        datatype_fun = getattr(sqltypes, datatype_name)
-
-        if felis_type.is_sized:
-            datatype = datatype_fun(column_length, **kwargs)
-        else:
-            datatype = datatype_fun(**kwargs)
-
-        nullable_default = True
-        if isinstance(datatype, Numeric):
-            nullable_default = False
-
-        column_nullable = column_obj.get("nullable", nullable_default)
-        column_autoincrement = column_obj.get("autoincrement", "auto")
-
-        column: Column = Column(
-            column_name,
-            datatype,
-            comment=column_description,
-            autoincrement=column_autoincrement,
-            nullable=column_nullable,
-            server_default=column_default,
-        )
-        if column_id in self.graph_index:
-            logger.warning(f"Duplication of @id {column_id}")
-        self.graph_index[column_id] = column
-        return column
-
-    def visit_primary_key(
-        self, primary_key_obj: str | Iterable[str], table_obj: _Mapping
-    ) -> PrimaryKeyConstraint | None:
-        # Docstring is inherited.
-        self.checker.check_primary_key(primary_key_obj, table_obj)
-        if primary_key_obj:
-            if isinstance(primary_key_obj, str):
-                primary_key_obj = [primary_key_obj]
-            columns = [self.graph_index[c_id] for c_id in primary_key_obj]
-            return PrimaryKeyConstraint(*columns)
-        return None
-
-    def visit_constraint(self, constraint_obj: _Mapping, table_obj: _Mapping) -> Constraint:
-        # Docstring is inherited.
-        self.checker.check_constraint(constraint_obj, table_obj)
-        constraint_type = constraint_obj["@type"]
-        constraint_id = constraint_obj["@id"]
-
-        constraint_args: _MutableMapping = {}
-        # The following are not used on every constraint
-        _set_if("name", constraint_obj.get("name"), constraint_args)
-        _set_if("info", constraint_obj.get("description"), constraint_args)
-        _set_if("expression", constraint_obj.get("expression"), constraint_args)
-        _set_if("deferrable", constraint_obj.get("deferrable"), constraint_args)
-        _set_if("initially", constraint_obj.get("initially"), constraint_args)
-
-        columns = [self.graph_index[c_id] for c_id in constraint_obj.get("columns", [])]
-        constraint: Constraint
-        if constraint_type == "ForeignKey":
-            refcolumns = [self.graph_index[c_id] for c_id in constraint_obj.get("referencedColumns", [])]
-            constraint = ForeignKeyConstraint(columns, refcolumns, **constraint_args)
-        elif constraint_type == "Check":
-            expression = constraint_obj["expression"]
-            constraint = CheckConstraint(expression, **constraint_args)
-        elif constraint_type == "Unique":
-            constraint = UniqueConstraint(*columns, **constraint_args)
-        else:
-            raise ValueError(f"Unexpected constraint type: {constraint_type}")
-        self.graph_index[constraint_id] = constraint
-        return constraint
-
-    def visit_index(self, index_obj: _Mapping, table_obj: _Mapping) -> Index:
-        # Docstring is inherited.
-        self.checker.check_index(index_obj, table_obj)
-        name = index_obj["name"]
-        description = index_obj.get("description")
-        columns = [self.graph_index[c_id] for c_id in index_obj.get("columns", [])]
-        expressions = index_obj.get("expressions", [])
-        return Index(name, *columns, *expressions, info=description)
-
-
-def _set_if(key: str, value: Any, mapping: _MutableMapping) -> None:
-    if value is not None:
-        mapping[key] = value
-
-
-def _process_variant_override(dialect_name: str, variant_override_str: str) -> types.TypeEngine:
-    """Return variant type for given dialect."""
-    match = length_regex.search(variant_override_str)
-    dialect = DIALECT_MODULES[dialect_name]
-    variant_type_name = variant_override_str.split("(")[0]
-
-    # Process Variant Type
-    if variant_type_name not in dir(dialect):
-        raise ValueError(f"Type {variant_type_name} not found in dialect {dialect_name}")
-    variant_type = getattr(dialect, variant_type_name)
-    length_params = []
-    if match:
-        length_params.extend([int(i) for i in match.group(1).split(",")])
-    return variant_type(*length_params)