Better handling of non-existent tables

[albireox]

The catalogdb models are defined for operations, where all the tables exist. The copy in pipelines has only a subset for the tables. Currently the import of the models for catalogdb works but many of the models are orphan and fail when a query is attempted, as there is no associated table. In SQLA this causes an error that requires a rollback.

We should find a better way to handle this. Ideally we would avoid creating the models, or remove them once we verify that the table does not exist. But this requires an active connection, at which point the model has already been defined. Removing the model from the module is not trivial since it depends on how the module or class were imported.

[albireox]

As an alternative, I tried to create the missing tables as temporary tables. This does not fully work yet (it creates the tables, but not as temporary) but may be promising. The code is for sqlalchemy/sdss5db/__init__.py

# !usr/bin/env python
# -*- coding: utf-8 -*-
#
from attr import has

from sqlalchemy import inspect
from sqlalchemy.ext.declarative import DeferredReflection
from sqlalchemy.orm import DeclarativeBase, Mapped, declared_attr

from sdssdb.connection import SQLADatabaseConnection
from sdssdb.sqlalchemy import BaseModel


# we need a shared common Base when joining across multiple schema
class SDSS5dbBase(BaseModel, DeclarativeBase):
    pass


def recursive_subclasses(cls):
    """Recursively get all subclasses of a class."""

    subclasses = set()

    for subclass in cls.__subclasses__():
        if issubclass(subclass, BaseModel):
            subclasses.add(subclass)
            subclasses.update(recursive_subclasses(subclass))

    return subclasses


def get_schema(model: type[SDSS5dbBase]) -> str | None:
    """Returns the schema of a model, if any."""

    if hasattr(model, "__table_args__") and model.__table_args__:
        table_args = model.__table_args__
        if isinstance(table_args, dict) and "schema" in table_args:
            return table_args["schema"]
        elif isinstance(table_args, tuple):
            for arg in table_args:
                if isinstance(arg, dict) and "schema" in arg:
                    return arg["schema"]

    return None


class SDSS5dbDatabaseConnection(SQLADatabaseConnection):
    dbname = "sdss5db"
    base = SDSS5dbBase

    def post_connect(self):
        """Called after a successful connection to the database."""

        if self.engine is None:
            return

        models = recursive_subclasses(SDSS5dbBase)

        schemas = set()
        for model in models:
            schema = get_schema(model)
            if schema:
                schemas.add(schema)

        inspector = inspect(self.engine)

        for schema in schemas:
            table_names = inspector.get_table_names(schema=schema)

            for model in models:
                schema = get_schema(model)
                if schema != schema or not hasattr(model, "__tablename__"):
                    continue

                table_name = model.__tablename__
                if table_name is not None and table_name not in table_names:
                    if isinstance(model.__table_args__, dict):
                        model.__table_args__ = declared_attr(
                            lambda cls: {"schema": schema, "prefixes": ["TEMPORARY"]}
                        )
                        model.metadata.create_all(self.engine)


database = SDSS5dbDatabaseConnection()