lsst-felis 27.2024.2500__py3-none-any.whl → 27.2024.2600__py3-none-any.whl

felis/db/utils.py

@@ -1,3 +1,5 @@
+"""Database utility functions and classes."""
+
 # This file is part of felis.
 #
 # Developed for the LSST Data Management System.
@@ -36,17 +38,44 @@ from sqlalchemy.types import TypeEngine
 
 from .dialects import get_dialect_module
 
+__all__ = ["string_to_typeengine", "SQLWriter", "ConnectionWrapper", "DatabaseContext"]
+
 logger = logging.getLogger("felis")
 
 _DATATYPE_REGEXP = re.compile(r"(\w+)(\((.*)\))?")
-"""Regular expression to match data types in the form "type(length)"""
+"""Regular expression to match data types with parameters in parentheses."""
 
 
 def string_to_typeengine(
     type_string: str, dialect: Dialect | None = None, length: int | None = None
 ) -> TypeEngine:
-    """Convert a string representation of a data type to a SQLAlchemy
-    TypeEngine.
+    """Convert a string representation of a datatype to a SQLAlchemy type.
+
+    Parameters
+    ----------
+    type_string
+        The string representation of the data type.
+    dialect
+        The SQLAlchemy dialect to use. If None, the default dialect will be
+        used.
+    length
+        The length of the data type. If the data type does not have a length
+        attribute, this parameter will be ignored.
+
+    Returns
+    -------
+    `sqlalchemy.types.TypeEngine`
+        The SQLAlchemy type engine object.
+
+    Raises
+    ------
+    ValueError
+        If the type string is invalid or the type is not supported.
+
+    Notes
+    -----
+    This function is used when converting type override strings defined in
+    fields such as ``mysql:datatype`` in the schema data.
     """
     match = _DATATYPE_REGEXP.search(type_string)
     if not match:
@@ -78,17 +107,17 @@ def string_to_typeengine(
 
 
 class SQLWriter:
-    """Writes SQL statements to stdout or a file."""
+    """Write SQL statements to stdout or a file.
 
-    def __init__(self, file: IO[str] | None = None) -> None:
-        """Initialize the SQL writer.
+    Parameters
+    ----------
+    file
+        The file to write the SQL statements to. If None, the statements
+        will be written to stdout.
+    """
 
-        Parameters
-        ----------
-        file : `io.TextIOBase` or `None`, optional
-            The file to write the SQL statements to. If None, the statements
-            will be written to stdout.
-        """
+    def __init__(self, file: IO[str] | None = None) -> None:
+        """Initialize the SQL writer."""
         self.file = file
         self.dialect: Dialect | None = None
 
@@ -100,12 +129,17 @@ class SQLWriter:
 
         Parameters
         ----------
-        sql : `typing.Any`
+        sql
             The SQL statement to write.
-        multiparams : `typing.Any`
+        *multiparams
             The multiparams to use for the SQL statement.
-        params : `typing.Any`
+        **params
             The params to use for the SQL statement.
+
+        Notes
+        -----
+        The functions arguments are typed very loosely because this method in
+        SQLAlchemy is untyped, amd we do not call it directly.
         """
         compiled = sql.compile(dialect=self.dialect)
         sql_str = str(compiled) + ";"
@@ -126,22 +160,37 @@ class SQLWriter:
 
 
 class ConnectionWrapper:
-    """A wrapper for a SQLAlchemy engine or mock connection which provides a
-    consistent interface for executing SQL statements.
+    """Wrap a SQLAlchemy engine or mock connection to provide a consistent
+    interface for executing SQL statements.
+
+    Parameters
+    ----------
+    engine
+        The SQLAlchemy engine or mock connection to wrap.
     """
 
     def __init__(self, engine: Engine | MockConnection):
-        """Initialize the connection wrapper.
+        """Initialize the connection wrapper."""
+        self.engine = engine
+
+    def execute(self, statement: Any) -> ResultProxy:
+        """Execute a SQL statement on the engine and return the result.
 
         Parameters
         ----------
-        engine : `sqlalchemy.Engine` or `sqlalchemy.MockConnection`
-            The SQLAlchemy engine or mock connection to wrap.
+        statement
+            The SQL statement to execute.
+
+        Returns
+        -------
+        ``sqlalchemy.engine.ResultProxy``
+            The result of the statement execution.
+
+        Notes
+        -----
+        The statement will be executed in a transaction block if not using
+        a mock connection.
         """
-        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):
@@ -153,19 +202,19 @@ class ConnectionWrapper:
 
 
 class DatabaseContext:
-    """A class for managing the schema and its database connection."""
+    """Manage the database connection and SQLAlchemy metadata.
 
-    def __init__(self, metadata: MetaData, engine: Engine | MockConnection):
-        """Initialize the database context.
+    Parameters
+    ----------
+    metadata
+        The SQLAlchemy metadata object.
 
-        Parameters
-        ----------
-        metadata : `sqlalchemy.MetaData`
-            The SQLAlchemy metadata object.
+    engine
+        The SQLAlchemy engine or mock connection object.
+    """
 
-        engine : `sqlalchemy.Engine` or `sqlalchemy.MockConnection`
-            The SQLAlchemy engine or mock connection object.
-        """
+    def __init__(self, metadata: MetaData, engine: Engine | MockConnection):
+        """Initialize the database context."""
         self.engine = engine
         self.dialect_name = engine.dialect.name
         self.metadata = metadata
@@ -174,16 +223,18 @@ class DatabaseContext:
     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
+        Raises
+        ------
+        ValueError
+            If the database is not supported.
+        sqlalchemy.exc.SQLAlchemyError
+            If there is an error creating the schema.
+
+        Notes
+        -----
+        In MySQL, this will create a new database and, 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.
         """
         schema_name = self.metadata.schema
         try:
@@ -202,15 +253,15 @@ class DatabaseContext:
     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.
+        Raises
+        ------
+        ValueError
+            If the database is not supported.
 
-        Parameters
-        ----------
-        engine: `sqlalchemy.Engine`
-            The SQLAlchemy engine object.
-        schema_name: `str`
-            The name of the schema (or database) to drop.
+        Notes
+        -----
+        In MySQL, this will drop a database. In PostgreSQL, it will drop a
+        schema. For other variants, this is an unsupported operation.
         """
         schema_name = self.metadata.schema
         try:
@@ -236,11 +287,16 @@ class DatabaseContext:
 
         Parameters
         ----------
-        engine_url : `sqlalchemy.engine.url.URL`
+        engine_url
             The SQLAlchemy engine URL.
-        output_file : `typing.IO` [ `str` ] or `None`, optional
+        output_file
             The file to write the SQL statements to. If None, the statements
             will be written to stdout.
+
+        Returns
+        -------
+        ``sqlalchemy.engine.mock.MockConnection``
+            The mock connection object.
         """
         writer = SQLWriter(output_file)
         engine = create_mock_engine(engine_url, executor=writer.write)

felis/db/variants.py

@@ -1,3 +1,5 @@
+"""Handle variant overrides for a Felis column."""
+
 # This file is part of felis.
 #
 # Developed for the LSST Data Management System.
@@ -19,7 +21,11 @@
 # 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 re
+from collections.abc import Mapping
+from types import MappingProxyType
 from typing import Any
 
 from sqlalchemy import types
@@ -28,26 +34,55 @@ from sqlalchemy.types import TypeEngine
 from ..datamodel import Column
 from .dialects import get_dialect_module, get_supported_dialects
 
+__all__ = ["make_variant_dict"]
+
 
 def _create_column_variant_overrides() -> dict[str, str]:
-    """Create a dictionary of column variant overrides."""
+    """Map column variant overrides to their dialect name.
+
+    Returns
+    -------
+    column_variant_overrides : `dict` [ `str`, `str` ]
+        A mapping of column variant overrides to their dialect name.
+
+    Notes
+    -----
+    This function is intended for internal use only.
+    """
     column_variant_overrides = {}
     for dialect_name in get_supported_dialects().keys():
         column_variant_overrides[f"{dialect_name}_datatype"] = dialect_name
     return column_variant_overrides
 
 
-_COLUMN_VARIANT_OVERRIDES = _create_column_variant_overrides()
+_COLUMN_VARIANT_OVERRIDES = MappingProxyType(_create_column_variant_overrides())
+"""Map of column variant overrides to their dialect name."""
+
 
+def _get_column_variant_overrides() -> Mapping[str, str]:
+    """Get a dictionary of column variant overrides.
 
-def _get_column_variant_overrides() -> dict[str, str]:
-    """Return a dictionary of column variant overrides."""
+    Returns
+    -------
+    column_variant_overrides : `dict` [ `str`, `str` ]
+        A mapping of column variant overrides to their dialect name.
+    """
     return _COLUMN_VARIANT_OVERRIDES
 
 
 def _get_column_variant_override(field_name: str) -> str:
-    """Return the dialect name from an override field name on the column like
+    """Get the dialect name from an override field name on the column like
     ``mysql_datatype``.
+
+    Returns
+    -------
+    dialect_name : `str`
+        The name of the dialect.
+
+    Raises
+    ------
+    ValueError
+        If the field name is not found in the column variant overrides.
     """
     if field_name not in _COLUMN_VARIANT_OVERRIDES:
         raise ValueError(f"Field name {field_name} not found in column variant overrides")
@@ -59,7 +94,30 @@ _length_regex = re.compile(r"\((\d+)\)")
 
 
 def _process_variant_override(dialect_name: str, variant_override_str: str) -> types.TypeEngine:
-    """Return variant type for given dialect."""
+    """Get the variant type for the given dialect.
+
+    Parameters
+    ----------
+    dialect_name
+        The name of the dialect to create.
+    variant_override_str
+        The string representation of the variant override.
+
+    Returns
+    -------
+    variant_type : `~sqlalchemy.types.TypeEngine`
+        The variant type for the given dialect.
+
+    Raises
+    ------
+    ValueError
+        If the type is not found in the dialect.
+
+    Notes
+    -----
+    This function converts a string representation of a variant override
+    into a `sqlalchemy.types.TypeEngine` object.
+    """
     dialect = get_dialect_module(dialect_name)
     variant_type_name = variant_override_str.split("(")[0]
 
@@ -82,12 +140,12 @@ def make_variant_dict(column_obj: Column) -> dict[str, TypeEngine[Any]]:
 
     Parameters
     ----------
-    column_obj : `felis.datamodel.Column`
+    column_obj
         The column object from which to build the variant dictionary.
 
     Returns
     -------
-    variant_dict : `dict`
+    `dict` [ `str`, `~sqlalchemy.types.TypeEngine` ]
         The dictionary of `str` to `sqlalchemy.types.TypeEngine` containing
         variant datatype information (e.g., for mysql, postgresql, etc).
     """

felis/metadata.py

@@ -1,3 +1,5 @@
+"""Build SQLAlchemy metadata from a Felis schema."""
+
 # This file is part of felis.
 #
 # Developed for the LSST Data Management System.
@@ -47,18 +49,25 @@ from . import datamodel
 from .db import sqltypes
 from .types import FelisType
 
+__all__ = ("MetaDataBuilder", "get_datatype_with_variants")
+
 logger = logging.getLogger(__name__)
 
 
 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.
+    overrides from the information in a Felis column object.
 
     Parameters
     ----------
-    column_obj : `felis.datamodel.Column`
+    column_obj
         The column object from which to get the datatype.
 
+    Returns
+    -------
+    `~sqlalchemy.types.TypeEngine`
+        The SQLAlchemy datatype object.
+
     Raises
     ------
     ValueError
@@ -80,22 +89,22 @@ _VALID_SERVER_DEFAULTS = ("CURRENT_TIMESTAMP", "NOW()", "LOCALTIMESTAMP", "NULL"
 
 
 class MetaDataBuilder:
-    """A class for building a `MetaData` object from a Felis `Schema`."""
+    """Build a SQLAlchemy metadata object from a Felis schema.
+
+    Parameters
+    ----------
+    schema
+        The schema object from which to build the SQLAlchemy metadata.
+    apply_schema_to_metadata
+        Whether to apply the schema name to the metadata object.
+    apply_schema_to_tables
+        Whether to apply the schema name to the tables.
+    """
 
     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.
-        """
+        """Initialize the metadata builder."""
         self.schema = schema
         if not apply_schema_to_metadata:
             logger.debug("Schema name will not be applied to metadata")
@@ -106,20 +115,25 @@ class MetaDataBuilder:
         self.apply_schema_to_tables = apply_schema_to_tables
 
     def build(self) -> MetaData:
-        """Build the SQLAlchemy tables and constraints from the schema."""
+        """Build the SQLAlchemy tables and constraints from the schema.
+
+        Notes
+        -----
+        This first builds the tables and then makes a second pass to build the
+        constraints. This is necessary because the constraints may reference
+        objects that are not yet created when the tables are built.
+
+        Returns
+        -------
+        `~sqlalchemy.sql.schema.MetaData`
+            The SQLAlchemy metadata object.
+        """
         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.
-        """
+        """Build the SQLAlchemy tables from the schema."""
         for table in self.schema.tables:
             self.build_table(table)
             if table.primary_key:
@@ -127,40 +141,44 @@ class MetaDataBuilder:
                 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.
+        """Build a SQAlchemy ``PrimaryKeyConstraint`` from a single column ID
+        or a list of them.
 
         Parameters
         ----------
-        primary_key_columns : `str` or `list` of `str`
+        primary_key_columns
             The column ID or list of column IDs from which to build the primary
             key.
 
         Returns
         -------
-        primary_key: `sqlalchemy.PrimaryKeyConstraint`
+        `~sqlalchemy.sql.schema.PrimaryKeyConstraint`
             The SQLAlchemy primary key constraint object.
+
+        Notes
+        -----
+        The ``primary_key_columns`` is a string or a list of strings
+        representing IDs which will be used to find the columnn objects in the
+        builder's internal ID map.
         """
         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.
+        """Build a SQLAlchemy ``Table`` from a Felis table and add it to the
+        metadata.
 
         Parameters
         ----------
-        table_obj : `felis.datamodel.Table`
-            The table object to build the SQLAlchemy table from.
+        table_obj
+            The Felis table object from which to build the SQLAlchemy table.
+
+        Notes
+        -----
+        Several MySQL table options, including the engine and charset, are
+        handled by adding annotations to the table. This is not needed for
+        Postgres, as Felis does not support any table options for this dialect.
         """
         # Process mysql table options.
         optargs = {}
@@ -192,16 +210,16 @@ class MetaDataBuilder:
         self._objects[id] = table
 
     def build_column(self, column_obj: datamodel.Column) -> Column:
-        """Build a SQLAlchemy column from a `felis.datamodel.Column` object.
+        """Build a SQLAlchemy ``Column`` from a Felis column object.
 
         Parameters
         ----------
-        column_obj : `felis.datamodel.Column`
+        column_obj
             The column object from which to build the SQLAlchemy column.
 
         Returns
         -------
-        column: `sqlalchemy.Column`
+        `~sqlalchemy.sql.schema.Column`
             The SQLAlchemy column object.
         """
         # Get basic column attributes.
@@ -244,8 +262,8 @@ class MetaDataBuilder:
         return column
 
     def build_constraints(self) -> None:
-        """Build the SQLAlchemy constraints in the Felis schema and append them
-        to the associated `Table`.
+        """Build the SQLAlchemy constraints from the Felis schema and append
+        them to the associated table in the metadata.
 
         Notes
         -----
@@ -260,18 +278,16 @@ class MetaDataBuilder:
                 table.append_constraint(constraint)
 
     def build_constraint(self, constraint_obj: datamodel.Constraint) -> Constraint:
-        """Build a SQLAlchemy `Constraint` from a `felis.datamodel.Constraint`
-        object.
+        """Build a SQLAlchemy ``Constraint`` from a  Felis constraint.
 
         Parameters
         ----------
-        constraint_obj : `felis.datamodel.Constraint`
-            The constraint object from which to build the SQLAlchemy
-            constraint.
+        constraint_obj
+            The Felis object from which to build the constraint.
 
         Returns
         -------
-        constraint: `sqlalchemy.Constraint`
+        `~sqlalchemy.sql.schema.Constraint`
             The SQLAlchemy constraint object.
 
         Raises
@@ -311,16 +327,16 @@ class MetaDataBuilder:
         return constraint
 
     def build_index(self, index_obj: datamodel.Index) -> Index:
-        """Build a SQLAlchemy `Index` from a `felis.datamodel.Index` object.
+        """Build a SQLAlchemy ``Index`` from a Felis `~felis.datamodel.Index`.
 
         Parameters
         ----------
-        index_obj : `felis.datamodel.Index`
-            The index object from which to build the SQLAlchemy index.
+        index_obj
+            The Felis object from which to build the SQLAlchemy index.
 
         Returns
         -------
-        index: `sqlalchemy.Index`
+        `~sqlalchemy.sql.schema.Index`
             The SQLAlchemy index object.
         """
         columns = [self._objects[c_id] for c_id in (index_obj.columns if index_obj.columns else [])]