apdbSqlReplica.py

Go to the documentation of this file.

# This file is part of dax_apdb.
#
# Developed for the LSST Data Management System.
# This product includes software developed by the LSST Project
# (http://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 <http://www.gnu.org/licenses/>.
 
"""Module defining Apdb class and related methods."""
 
from __future__ import annotations
 
__all__ = ["ApdbSqlReplica"]
 
import logging
from collections.abc import Collection, Iterable, Mapping, Sequence
from typing import TYPE_CHECKING, cast
 
import astropy.time
import felis.datamodel
import sqlalchemy
from sqlalchemy import sql
 
from ..apdbReplica import ApdbReplica, ApdbTableData, ReplicaChunk
from ..apdbSchema import ApdbTables
from ..monitor import MonAgent
from ..schema_model import ExtraDataTypes
from ..timer import Timer
from ..versionTuple import VersionTuple
from .apdbSqlSchema import ExtraTables
 
if TYPE_CHECKING:
    from .apdbSqlSchema import ApdbSqlSchema
 
 
_LOG = logging.getLogger(__name__)
 
_MON = MonAgent(__name__)
 
VERSION = VersionTuple(1, 0, 0)
"""Version for the code controlling replication tables. This needs to be
updated following compatibility rules when schema produced by this code
changes.
"""
 
 
class ApdbSqlTableData(ApdbTableData):
    """Implementation of ApdbTableData that wraps sqlalchemy Result."""
 
    def __init__(self, result: sqlalchemy.engine.Result, column_types: dict[str, felis.datamodel.DataType]):
        self._column_defs = tuple((column, column_types[column]) for column in result.keys())
        self._rows: list[tuple] = cast(list[tuple], list(result.fetchall()))
 
    def column_names(self) -> Sequence[str]:
        return tuple(column_def[0] for column_def in self._column_defs)
 
    def column_defs(self) -> Sequence[tuple[str, felis.datamodel.DataType]]:
        return self._column_defs
 
    def rows(self) -> Collection[tuple]:
        return self._rows
 
 
class ApdbSqlReplica(ApdbReplica):
    """Implementation of `ApdbReplica` for SQL backend.
 
    Parameters
    ----------
    schema : `ApdbSqlSchema`
        Instance of `ApdbSqlSchema` class for APDB database.
    engine : `sqlalchemy.engine.Engine`
        Engine for database access.
    db_schema_version : `VersionTuple`
        Version of the database schema.
    timer : `bool`, optional
        If `True` then log timing information.
    """
 
    def __init__(
        self,
        schema: ApdbSqlSchema,
        engine: sqlalchemy.engine.Engine,
        db_schema_version: VersionTuple,
        timer: bool = False,
    ):
        self._schema = schema
        self._engine = engine
        self._db_schema_version = db_schema_version
 
        self._timer_args: list[MonAgent | logging.Logger] = [_MON]
        if timer:
            self._timer_args.append(_LOG)
 
    def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
        """Create `Timer` instance given its name."""
        return Timer(name, *self._timer_args, tags=tags)
 
    def schemaVersion(self) -> VersionTuple:
        # Docstring inherited from base class.
        return self._db_schema_version
 
    @classmethod
    def apdbReplicaImplementationVersion(cls) -> VersionTuple:
        # Docstring inherited from base class.
        return VERSION
 
    def getReplicaChunks(self) -> list[ReplicaChunk] | None:
        # docstring is inherited from a base class
        if not self._schema.replication_enabled:
            return None
 
        table = self._schema.get_table(ExtraTables.ApdbReplicaChunks)
        assert table is not None, "replication_enabled=True means it must be defined"
        query = sql.select(
            table.columns["apdb_replica_chunk"], table.columns["last_update_time"], table.columns["unique_id"]
        ).order_by(table.columns["last_update_time"])
        with self._timer("chunks_select_time") as timer:
            with self._engine.connect() as conn:
                result = conn.execution_options(stream_results=True, max_row_buffer=10000).execute(query)
                ids = []
                for row in result:
                    last_update_time = astropy.time.Time(row[1].timestamp(), format="unix_tai")
                    ids.append(ReplicaChunk(id=row[0], last_update_time=last_update_time, unique_id=row[2]))
                timer.add_values(row_count=len(ids))
                return ids
 
    def deleteReplicaChunks(self, chunks: Iterable[int]) -> None:
        # docstring is inherited from a base class
        if not self._schema.replication_enabled:
            raise ValueError("APDB is not configured for replication")
 
        table = self._schema.get_table(ExtraTables.ApdbReplicaChunks)
        chunk_list = list(chunks)
        where_clause = table.columns["apdb_replica_chunk"].in_(chunk_list)
        stmt = table.delete().where(where_clause)
        with self._timer("chunks_delete_time") as timer:
            with self._engine.begin() as conn:
                conn.execute(stmt)
            timer.add_values(row_count=len(chunk_list))
 
    def getTableDataChunks(self, table: ApdbTables, chunks: Iterable[int]) -> ApdbTableData:
        # docstring is inherited from a base class
        for chunk_table, table_enum in ExtraTables.replica_chunk_tables().items():
            if table is table_enum:
                return self._get_chunks(chunks, table, chunk_table)
        raise ValueError(f"Table {table} does not support replica chunks.")
 
    def _get_chunks(
        self,
        chunks: Iterable[int],
        table_enum: ApdbTables,
        chunk_table_enum: ExtraTables,
    ) -> ApdbTableData:
        """Return catalog of records for given insert identifiers, common
        implementation for all DIA tables.
        """
        if not self._schema.replication_enabled:
            raise ValueError("APDB is not configured for replication")
 
        table = self._schema.get_table(table_enum)
        chunk_table = self._schema.get_table(chunk_table_enum)
 
        join = table.join(chunk_table)
        chunk_id_column = chunk_table.columns["apdb_replica_chunk"]
        apdb_columns = self._schema.get_apdb_columns(table_enum)
        where_clause = chunk_id_column.in_(chunks)
        query = sql.select(chunk_id_column, *apdb_columns).select_from(join).where(where_clause)
 
        table_schema = self._schema.tableSchemas[table_enum]
        # Regular tables should never have columns of ExtraDataTypes, this is
        # just to make mypy happy.
        column_types = {
            column.name: column.datatype
            for column in table_schema.columns
            if not isinstance(column.datatype, ExtraDataTypes)
        }
        column_types["apdb_replica_chunk"] = felis.datamodel.DataType.long
 
        # execute select
        with self._timer("table_chunk_select_time", tags={"table": table.name}) as timer:
            with self._engine.begin() as conn:
                result = conn.execution_options(stream_results=True, max_row_buffer=10000).execute(query)
                table_data = ApdbSqlTableData(result, column_types)
                timer.add_values(row_count=len(table_data.rows()))
                return table_data
 
    def getTableUpdateChunks(self, table: ApdbTables, chunks: Iterable[int]) -> ApdbTableData:
        # docstring is inherited from a base class
        raise NotImplementedError()