Coverage for python/lsst/dax/apdb/sql/apdbSqlReplica.py: 35%
75 statements
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-26 09:55 +0000
« prev ^ index » next coverage.py v7.5.0, created at 2024-04-26 09:55 +0000
1# This file is part of dax_apdb.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22"""Module defining Apdb class and related methods.
23"""
25from __future__ import annotations
27__all__ = ["ApdbSqlReplica"]
29import logging
30from collections.abc import Collection, Iterable, Sequence
31from typing import TYPE_CHECKING, cast
33import astropy.time
34import sqlalchemy
35from sqlalchemy import sql
37from ..apdbReplica import ApdbReplica, ApdbTableData, ReplicaChunk
38from ..apdbSchema import ApdbTables
39from ..timer import Timer
40from ..versionTuple import VersionTuple
41from .apdbSqlSchema import ExtraTables
43if TYPE_CHECKING:
44 from .apdbSqlSchema import ApdbSqlSchema
47_LOG = logging.getLogger(__name__)
49VERSION = VersionTuple(1, 0, 0)
50"""Version for the code controlling replication tables. This needs to be
51updated following compatibility rules when schema produced by this code
52changes.
53"""
56class ApdbSqlTableData(ApdbTableData):
57 """Implementation of ApdbTableData that wraps sqlalchemy Result."""
59 def __init__(self, result: sqlalchemy.engine.Result):
60 self._keys = list(result.keys())
61 self._rows: list[tuple] = cast(list[tuple], list(result.fetchall()))
63 def column_names(self) -> Sequence[str]:
64 return self._keys
66 def rows(self) -> Collection[tuple]:
67 return self._rows
70class ApdbSqlReplica(ApdbReplica):
71 """Implementation of `ApdbReplica` for SQL backend.
73 Parameters
74 ----------
75 schema : `ApdbSqlSchema`
76 Instance of `ApdbSqlSchema` class for APDB database.
77 engine : `sqlalchemy.engine.Engine`
78 Engine for database access.
79 timer : `bool`, optional
80 If `True` then log timing information.
81 """
83 def __init__(self, schema: ApdbSqlSchema, engine: sqlalchemy.engine.Engine, timer: bool = False):
84 self._schema = schema
85 self._engine = engine
86 self._timer = timer
88 @classmethod
89 def apdbReplicaImplementationVersion(cls) -> VersionTuple:
90 # Docstring inherited from base class.
91 return VERSION
93 def getReplicaChunks(self) -> list[ReplicaChunk] | None:
94 # docstring is inherited from a base class
95 if not self._schema.has_replica_chunks:
96 return None
98 table = self._schema.get_table(ExtraTables.ApdbReplicaChunks)
99 assert table is not None, "has_replica_chunks=True means it must be defined"
100 query = sql.select(
101 table.columns["apdb_replica_chunk"], table.columns["last_update_time"], table.columns["unique_id"]
102 ).order_by(table.columns["last_update_time"])
103 with Timer("DiaObject insert id select", self._timer):
104 with self._engine.connect() as conn:
105 result = conn.execution_options(stream_results=True, max_row_buffer=10000).execute(query)
106 ids = []
107 for row in result:
108 last_update_time = astropy.time.Time(row[1].timestamp(), format="unix_tai")
109 ids.append(ReplicaChunk(id=row[0], last_update_time=last_update_time, unique_id=row[2]))
110 return ids
112 def deleteReplicaChunks(self, chunks: Iterable[int]) -> None:
113 # docstring is inherited from a base class
114 if not self._schema.has_replica_chunks:
115 raise ValueError("APDB is not configured for replication")
117 table = self._schema.get_table(ExtraTables.ApdbReplicaChunks)
118 where_clause = table.columns["apdb_replica_chunk"].in_(chunks)
119 stmt = table.delete().where(where_clause)
120 with self._engine.begin() as conn:
121 conn.execute(stmt)
123 def getDiaObjectsChunks(self, chunks: Iterable[int]) -> ApdbTableData:
124 # docstring is inherited from a base class
125 return self._get_chunks(chunks, ApdbTables.DiaObject, ExtraTables.DiaObjectChunks)
127 def getDiaSourcesChunks(self, chunks: Iterable[int]) -> ApdbTableData:
128 # docstring is inherited from a base class
129 return self._get_chunks(chunks, ApdbTables.DiaSource, ExtraTables.DiaSourceChunks)
131 def getDiaForcedSourcesChunks(self, chunks: Iterable[int]) -> ApdbTableData:
132 # docstring is inherited from a base class
133 return self._get_chunks(chunks, ApdbTables.DiaForcedSource, ExtraTables.DiaForcedSourceChunks)
135 def _get_chunks(
136 self,
137 chunks: Iterable[int],
138 table_enum: ApdbTables,
139 chunk_table_enum: ExtraTables,
140 ) -> ApdbTableData:
141 """Return catalog of records for given insert identifiers, common
142 implementation for all DIA tables.
143 """
144 if not self._schema.has_replica_chunks:
145 raise ValueError("APDB is not configured for replication")
147 table = self._schema.get_table(table_enum)
148 chunk_table = self._schema.get_table(chunk_table_enum)
150 join = table.join(chunk_table)
151 chunk_id_column = chunk_table.columns["apdb_replica_chunk"]
152 apdb_columns = self._schema.get_apdb_columns(table_enum)
153 where_clause = chunk_id_column.in_(chunks)
154 query = sql.select(chunk_id_column, *apdb_columns).select_from(join).where(where_clause)
156 # execute select
157 with Timer(f"{table.name} replica chunk select", self._timer):
158 with self._engine.begin() as conn:
159 result = conn.execution_options(stream_results=True, max_row_buffer=10000).execute(query)
160 return ApdbSqlTableData(result)