Coverage for python / lsst / dax / apdb / sql / apdbSql.py: 14%
551 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-20 01:35 -0800
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."""
24from __future__ import annotations
26__all__ = ["ApdbSql"]
28import datetime
29import logging
30import urllib.parse
31import uuid
32import warnings
33from collections.abc import Iterable, Mapping, MutableMapping
34from contextlib import closing
35from typing import TYPE_CHECKING, Any
37import astropy.time
38import numpy as np
39import pandas
40import sqlalchemy
41import sqlalchemy.dialects.postgresql
42import sqlalchemy.dialects.sqlite
43from sqlalchemy import func, sql
44from sqlalchemy.pool import NullPool
46from lsst.sphgeom import HtmPixelization, LonLat, Region, UnitVector3d
47from lsst.utils.db_auth import DbAuth, DbAuthNotFoundError
48from lsst.utils.iteration import chunk_iterable
50from ..apdb import Apdb
51from ..apdbConfigFreezer import ApdbConfigFreezer
52from ..apdbReplica import ReplicaChunk
53from ..apdbSchema import ApdbSchema, ApdbTables
54from ..config import ApdbConfig
55from ..monitor import MonAgent
56from ..schema_model import Table
57from ..timer import Timer
58from ..versionTuple import IncompatibleVersionError, VersionTuple
59from .apdbMetadataSql import ApdbMetadataSql
60from .apdbSqlAdmin import ApdbSqlAdmin
61from .apdbSqlReplica import ApdbSqlReplica
62from .apdbSqlSchema import ApdbSqlSchema, ExtraTables
63from .config import ApdbSqlConfig
65if TYPE_CHECKING:
66 import sqlite3
68 from ..apdbMetadata import ApdbMetadata
69 from ..apdbUpdateRecord import ApdbUpdateRecord
71_LOG = logging.getLogger(__name__)
73_MON = MonAgent(__name__)
75VERSION = VersionTuple(1, 2, 1)
76"""Version for the code controlling non-replication tables. This needs to be
77updated following compatibility rules when schema produced by this code
78changes.
79"""
82def _coerce_uint64(df: pandas.DataFrame) -> pandas.DataFrame:
83 """Change the type of uint64 columns to int64, and return copy of data
84 frame.
85 """
86 names = [c[0] for c in df.dtypes.items() if c[1] == np.uint64]
87 return df.astype(dict.fromkeys(names, np.int64))
90def _make_midpointMjdTai_start(visit_time: astropy.time.Time, months: int) -> float:
91 """Calculate starting point for time-based source search.
93 Parameters
94 ----------
95 visit_time : `astropy.time.Time`
96 Time of current visit.
97 months : `int`
98 Number of months in the sources history.
100 Returns
101 -------
102 time : `float`
103 A ``midpointMjdTai`` starting point, MJD time.
104 """
105 # TODO: Use of MJD must be consistent with the code in ap_association
106 # (see DM-31996)
107 return float(visit_time.tai.mjd) - months * 30
110def _onSqlite3Connect(
111 dbapiConnection: sqlite3.Connection, connectionRecord: sqlalchemy.pool._ConnectionRecord
112) -> None:
113 # Enable foreign keys
114 with closing(dbapiConnection.cursor()) as cursor:
115 cursor.execute("PRAGMA foreign_keys=ON;")
118class ApdbSql(Apdb):
119 """Implementation of APDB interface based on SQL database.
121 The implementation is configured via standard ``pex_config`` mechanism
122 using `ApdbSqlConfig` configuration class. For an example of different
123 configurations check ``config/`` folder.
125 Parameters
126 ----------
127 config : `ApdbSqlConfig`
128 Configuration object.
129 """
131 metadataSchemaVersionKey = "version:schema"
132 """Name of the metadata key to store schema version number."""
134 metadataCodeVersionKey = "version:ApdbSql"
135 """Name of the metadata key to store code version number."""
137 metadataReplicaVersionKey = "version:ApdbSqlReplica"
138 """Name of the metadata key to store replica code version number."""
140 metadataConfigKey = "config:apdb-sql.json"
141 """Name of the metadata key to store code version number."""
143 _frozen_parameters = (
144 "enable_replica",
145 "dia_object_index",
146 "pixelization.htm_level",
147 "pixelization.htm_index_column",
148 "ra_dec_columns",
149 )
150 """Names of the config parameters to be frozen in metadata table."""
152 def __init__(self, config: ApdbSqlConfig):
153 self._engine = self._makeEngine(config, create=False)
155 sa_metadata = sqlalchemy.MetaData(schema=config.namespace)
156 meta_table_name = ApdbTables.metadata.table_name(prefix=config.prefix)
157 meta_table = sqlalchemy.schema.Table(meta_table_name, sa_metadata, autoload_with=self._engine)
158 self._metadata = ApdbMetadataSql(self._engine, meta_table)
160 # Get tables schemas.
161 self._table_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
163 # Check that versions are compatible, must be the first thing before
164 # reading frozen config.
165 self._db_schema_version = self._versionCheck(self._metadata, self._table_schema.schemaVersion())
167 # Read frozen config from metadata.
168 config_json = self._metadata.get(self.metadataConfigKey)
169 if config_json is not None:
170 # Update config from metadata.
171 freezer = ApdbConfigFreezer[ApdbSqlConfig](self._frozen_parameters)
172 self.config = freezer.update(config, config_json)
173 else:
174 self.config = config
176 self._schema = ApdbSqlSchema(
177 table_schema=self._table_schema,
178 engine=self._engine,
179 dia_object_index=self.config.dia_object_index,
180 prefix=self.config.prefix,
181 namespace=self.config.namespace,
182 htm_index_column=self.config.pixelization.htm_index_column,
183 enable_replica=self.config.enable_replica,
184 )
186 self.pixelator = HtmPixelization(self.config.pixelization.htm_level)
188 if _LOG.isEnabledFor(logging.DEBUG):
189 _LOG.debug("ApdbSql Configuration: %s", self.config.model_dump())
191 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
192 """Create `Timer` instance given its name."""
193 return Timer(name, _MON, tags=tags)
195 @classmethod
196 def _makeEngine(cls, config: ApdbSqlConfig, *, create: bool) -> sqlalchemy.engine.Engine:
197 """Make SQLALchemy engine based on configured parameters.
199 Parameters
200 ----------
201 config : `ApdbSqlConfig`
202 Configuration object.
203 create : `bool`
204 Whether to try to create new database file, only relevant for
205 SQLite backend which always creates new files by default.
206 """
207 # engine is reused between multiple processes, make sure that we don't
208 # share connections by disabling pool (by using NullPool class)
209 kw: MutableMapping[str, Any] = dict(config.connection_config.extra_parameters)
210 conn_args: dict[str, Any] = {}
211 if not config.connection_config.connection_pool:
212 kw.update(poolclass=NullPool)
213 if config.connection_config.isolation_level is not None:
214 kw.update(isolation_level=config.connection_config.isolation_level)
215 elif config.db_url.startswith("sqlite"):
216 # Use READ_UNCOMMITTED as default value for sqlite.
217 kw.update(isolation_level="READ_UNCOMMITTED")
218 if config.connection_config.connection_timeout is not None:
219 if config.db_url.startswith("sqlite"):
220 conn_args.update(timeout=config.connection_config.connection_timeout)
221 elif config.db_url.startswith(("postgresql", "mysql")):
222 conn_args.update(connect_timeout=int(config.connection_config.connection_timeout))
223 kw.update(connect_args=conn_args)
224 engine = sqlalchemy.create_engine(cls._connection_url(config.db_url, create=create), **kw)
226 if engine.dialect.name == "sqlite":
227 # Need to enable foreign keys on every new connection.
228 sqlalchemy.event.listen(engine, "connect", _onSqlite3Connect)
230 return engine
232 @classmethod
233 def _connection_url(cls, config_url: str, *, create: bool) -> sqlalchemy.engine.URL | str:
234 """Generate a complete URL for database with proper credentials.
236 Parameters
237 ----------
238 config_url : `str`
239 Database URL as specified in configuration.
240 create : `bool`
241 Whether to try to create new database file, only relevant for
242 SQLite backend which always creates new files by default.
244 Returns
245 -------
246 connection_url : `sqlalchemy.engine.URL` or `str`
247 Connection URL including credentials.
248 """
249 # Allow 3rd party authentication mechanisms by assuming connection
250 # string is correct when we can not recognize (dialect, host, database)
251 # matching keys.
252 components = urllib.parse.urlparse(config_url)
253 if all((components.scheme is not None, components.hostname is not None, components.path is not None)):
254 try:
255 db_auth = DbAuth()
256 config_url = db_auth.getUrl(config_url)
257 except DbAuthNotFoundError:
258 # Credentials file doesn't exist or no matching credentials,
259 # use default auth.
260 pass
262 # SQLite has a nasty habit creating empty databases when they do not
263 # exist, tell it not to do that unless we do need to create it.
264 if not create:
265 config_url = cls._update_sqlite_url(config_url)
267 return config_url
269 @classmethod
270 def _update_sqlite_url(cls, url_string: str) -> str:
271 """If URL refers to sqlite dialect, update it so that the backend does
272 not try to create database file if it does not exist already.
274 Parameters
275 ----------
276 url_string : `str`
277 Connection string.
279 Returns
280 -------
281 url_string : `str`
282 Possibly updated connection string.
283 """
284 try:
285 url = sqlalchemy.make_url(url_string)
286 except sqlalchemy.exc.SQLAlchemyError:
287 # If parsing fails it means some special format, likely not
288 # sqlite so we just return it unchanged.
289 return url_string
291 if url.get_backend_name() == "sqlite":
292 # Massage url so that database name starts with "file:" and
293 # option string has "mode=rw&uri=true". Database name
294 # should look like a path (:memory: is not supported by
295 # Apdb, but someone could still try to use it).
296 database = url.database
297 if database and not database.startswith((":", "file:")):
298 query = dict(url.query, mode="rw", uri="true")
299 # If ``database`` is an absolute path then original URL should
300 # include four slashes after "sqlite:". Humans are bad at
301 # counting things beyond four and sometimes an extra slash gets
302 # added unintentionally, which causes sqlite to treat initial
303 # element as "authority" and to complain. Strip extra slashes
304 # at the start of the path to avoid that (DM-46077).
305 if database.startswith("//"):
306 warnings.warn(
307 f"Database URL contains extra leading slashes which will be removed: {url}",
308 stacklevel=3,
309 )
310 database = "/" + database.lstrip("/")
311 url = url.set(database=f"file:{database}", query=query)
312 url_string = url.render_as_string()
314 return url_string
316 @classmethod
317 def _versionCheck(cls, metadata: ApdbMetadataSql, schema_version: VersionTuple) -> VersionTuple:
318 """Check schema version compatibility and return the database schema
319 version.
320 """
322 def _get_version(key: str) -> VersionTuple:
323 """Retrieve version number from given metadata key."""
324 version_str = metadata.get(key)
325 if version_str is None:
326 # Should not happen with existing metadata table.
327 raise RuntimeError(f"Version key {key!r} does not exist in metadata table.")
328 return VersionTuple.fromString(version_str)
330 db_schema_version = _get_version(cls.metadataSchemaVersionKey)
331 db_code_version = _get_version(cls.metadataCodeVersionKey)
333 # For now there is no way to make read-only APDB instances, assume that
334 # any access can do updates.
335 if not schema_version.checkCompatibility(db_schema_version):
336 raise IncompatibleVersionError(
337 f"Configured schema version {schema_version} "
338 f"is not compatible with database version {db_schema_version}"
339 )
340 if not cls.apdbImplementationVersion().checkCompatibility(db_code_version):
341 raise IncompatibleVersionError(
342 f"Current code version {cls.apdbImplementationVersion()} "
343 f"is not compatible with database version {db_code_version}"
344 )
346 # Check replica code version only if replica is enabled. Sort of
347 # chicken and egg problem - `enable_replica` is a part of frozen
348 # configuration, but we cannot read frozen configuration until we
349 # validate versions. Assume that if the replica version is present
350 # then replication is enabled.
351 if metadata.get(cls.metadataReplicaVersionKey) is not None:
352 db_replica_version = _get_version(cls.metadataReplicaVersionKey)
353 code_replica_version = ApdbSqlReplica.apdbReplicaImplementationVersion()
354 if not code_replica_version.checkCompatibility(db_replica_version):
355 raise IncompatibleVersionError(
356 f"Current replication code version {code_replica_version} "
357 f"is not compatible with database version {db_replica_version}"
358 )
360 return db_schema_version
362 @classmethod
363 def apdbImplementationVersion(cls) -> VersionTuple:
364 """Return version number for current APDB implementation.
366 Returns
367 -------
368 version : `VersionTuple`
369 Version of the code defined in implementation class.
370 """
371 return VERSION
373 @classmethod
374 def init_database(
375 cls,
376 db_url: str,
377 *,
378 schema_file: str | None = None,
379 ss_schema_file: str | None = None,
380 read_sources_months: int | None = None,
381 read_forced_sources_months: int | None = None,
382 enable_replica: bool = False,
383 connection_timeout: int | None = None,
384 dia_object_index: str | None = None,
385 htm_level: int | None = None,
386 htm_index_column: str | None = None,
387 ra_dec_columns: tuple[str, str] | None = None,
388 prefix: str | None = None,
389 namespace: str | None = None,
390 drop: bool = False,
391 ) -> ApdbSqlConfig:
392 """Initialize new APDB instance and make configuration object for it.
394 Parameters
395 ----------
396 db_url : `str`
397 SQLAlchemy database URL.
398 schema_file : `str`, optional
399 Location of (YAML) configuration file with APDB schema. If not
400 specified then default location will be used.
401 ss_schema_file : `str`, optional
402 Location of (YAML) configuration file with SSO schema. If not
403 specified then default location will be used.
404 read_sources_months : `int`, optional
405 Number of months of history to read from DiaSource.
406 read_forced_sources_months : `int`, optional
407 Number of months of history to read from DiaForcedSource.
408 enable_replica : `bool`, optional
409 If True, make additional tables used for replication to PPDB.
410 connection_timeout : `int`, optional
411 Database connection timeout in seconds.
412 dia_object_index : `str`, optional
413 Indexing mode for DiaObject table.
414 htm_level : `int`, optional
415 HTM indexing level.
416 htm_index_column : `str`, optional
417 Name of a HTM index column for DiaObject and DiaSource tables.
418 ra_dec_columns : `tuple` [`str`, `str`], optional
419 Names of ra/dec columns in DiaObject table.
420 prefix : `str`, optional
421 Optional prefix for all table names.
422 namespace : `str`, optional
423 Name of the database schema for all APDB tables. If not specified
424 then default schema is used.
425 drop : `bool`, optional
426 If `True` then drop existing tables before re-creating the schema.
428 Returns
429 -------
430 config : `ApdbSqlConfig`
431 Resulting configuration object for a created APDB instance.
432 """
433 config = ApdbSqlConfig(db_url=db_url, enable_replica=enable_replica)
434 if schema_file is not None:
435 config.schema_file = schema_file
436 if ss_schema_file is not None:
437 config.ss_schema_file = ss_schema_file
438 if read_sources_months is not None:
439 config.read_sources_months = read_sources_months
440 if read_forced_sources_months is not None:
441 config.read_forced_sources_months = read_forced_sources_months
442 if connection_timeout is not None:
443 config.connection_config.connection_timeout = connection_timeout
444 if dia_object_index is not None:
445 config.dia_object_index = dia_object_index
446 if htm_level is not None:
447 config.pixelization.htm_level = htm_level
448 if htm_index_column is not None:
449 config.pixelization.htm_index_column = htm_index_column
450 if ra_dec_columns is not None:
451 config.ra_dec_columns = ra_dec_columns
452 if prefix is not None:
453 config.prefix = prefix
454 if namespace is not None:
455 config.namespace = namespace
457 cls._makeSchema(config, drop=drop)
459 # SQLite has a nasty habit of creating empty database by default,
460 # update URL in config file to disable that behavior.
461 config.db_url = cls._update_sqlite_url(config.db_url)
463 return config
465 def get_replica(self) -> ApdbSqlReplica:
466 """Return `ApdbReplica` instance for this database."""
467 return ApdbSqlReplica(self._schema, self._engine, self._db_schema_version)
469 def tableRowCount(self) -> dict[str, int]:
470 """Return dictionary with the table names and row counts.
472 Used by ``ap_proto`` to keep track of the size of the database tables.
473 Depending on database technology this could be expensive operation.
475 Returns
476 -------
477 row_counts : `dict`
478 Dict where key is a table name and value is a row count.
479 """
480 res = {}
481 tables = [ApdbTables.DiaObject, ApdbTables.DiaSource, ApdbTables.DiaForcedSource]
482 if self.config.dia_object_index == "last_object_table":
483 tables.append(ApdbTables.DiaObjectLast)
484 with self._engine.begin() as conn:
485 for table in tables:
486 sa_table = self._schema.get_table(table)
487 stmt = sql.select(func.count()).select_from(sa_table)
488 count: int = conn.execute(stmt).scalar_one()
489 res[table.name] = count
491 return res
493 def getConfig(self) -> ApdbSqlConfig:
494 # docstring is inherited from a base class
495 return self.config
497 def tableDef(self, table: ApdbTables) -> Table | None:
498 # docstring is inherited from a base class
499 return self._schema.tableSchemas.get(table)
501 @classmethod
502 def _makeSchema(cls, config: ApdbConfig, drop: bool = False) -> None:
503 # docstring is inherited from a base class
505 if not isinstance(config, ApdbSqlConfig):
506 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
508 engine = cls._makeEngine(config, create=True)
510 table_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
512 # Ask schema class to create all tables.
513 schema = ApdbSqlSchema(
514 table_schema=table_schema,
515 engine=engine,
516 dia_object_index=config.dia_object_index,
517 prefix=config.prefix,
518 namespace=config.namespace,
519 htm_index_column=config.pixelization.htm_index_column,
520 enable_replica=config.enable_replica,
521 )
522 schema.makeSchema(drop=drop)
524 # Need metadata table to store few items in it.
525 meta_table = schema.get_table(ApdbTables.metadata)
526 apdb_meta = ApdbMetadataSql(engine, meta_table)
528 # Fill version numbers, overwrite if they are already there.
529 apdb_meta.set(cls.metadataSchemaVersionKey, str(table_schema.schemaVersion()), force=True)
530 apdb_meta.set(cls.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True)
531 if config.enable_replica:
532 # Only store replica code version if replica is enabled.
533 apdb_meta.set(
534 cls.metadataReplicaVersionKey,
535 str(ApdbSqlReplica.apdbReplicaImplementationVersion()),
536 force=True,
537 )
539 # Store frozen part of a configuration in metadata.
540 freezer = ApdbConfigFreezer[ApdbSqlConfig](cls._frozen_parameters)
541 apdb_meta.set(cls.metadataConfigKey, freezer.to_json(config), force=True)
543 def getDiaObjects(self, region: Region) -> pandas.DataFrame:
544 # docstring is inherited from a base class
546 # decide what columns we need
547 if self.config.dia_object_index == "last_object_table":
548 table_enum = ApdbTables.DiaObjectLast
549 else:
550 table_enum = ApdbTables.DiaObject
551 table = self._schema.get_table(table_enum)
552 if not self.config.dia_object_columns:
553 columns = self._schema.get_apdb_columns(table_enum)
554 else:
555 columns = [table.c[col] for col in self.config.dia_object_columns]
556 query = sql.select(*columns)
558 # build selection
559 query = query.where(self._filterRegion(table, region))
561 if self.schema.has_mjd_timestamps:
562 validity_end_column = "validityEndMjdTai"
563 else:
564 validity_end_column = "validityEnd"
566 # select latest version of objects
567 if self.config.dia_object_index != "last_object_table":
568 query = query.where(table.columns[validity_end_column] == None) # noqa: E711
570 # _LOG.debug("query: %s", query)
572 # execute select
573 with self._timer("select_time", tags={"table": "DiaObject"}) as timer:
574 with self._engine.begin() as conn:
575 objects = pandas.read_sql_query(query, conn)
576 timer.add_values(row_count=len(objects))
577 _LOG.debug("found %s DiaObjects", len(objects))
578 return self._fix_result_timestamps(objects)
580 def getDiaSources(
581 self,
582 region: Region,
583 object_ids: Iterable[int] | None,
584 visit_time: astropy.time.Time,
585 start_time: astropy.time.Time | None = None,
586 ) -> pandas.DataFrame | None:
587 # docstring is inherited from a base class
588 if start_time is None and self.config.read_sources_months == 0:
589 _LOG.debug("Skip DiaSources fetching")
590 return None
592 if start_time is None:
593 start_time_mjdTai = _make_midpointMjdTai_start(visit_time, self.config.read_sources_months)
594 else:
595 start_time_mjdTai = float(start_time.tai.mjd)
596 _LOG.debug("start_time_mjdTai = %.6f", start_time_mjdTai)
598 if object_ids is None:
599 # region-based select
600 return self._getDiaSourcesInRegion(region, start_time_mjdTai)
601 else:
602 return self._getDiaSourcesByIDs(list(object_ids), start_time_mjdTai)
604 def getDiaForcedSources(
605 self,
606 region: Region,
607 object_ids: Iterable[int] | None,
608 visit_time: astropy.time.Time,
609 start_time: astropy.time.Time | None = None,
610 ) -> pandas.DataFrame | None:
611 # docstring is inherited from a base class
612 if start_time is None and self.config.read_forced_sources_months == 0:
613 _LOG.debug("Skip DiaForceSources fetching")
614 return None
616 if object_ids is None:
617 # This implementation does not support region-based selection. In
618 # the past DiaForcedSource schema did not have ra/dec columns (it
619 # had x/y columns). ra/dec were added at some point, so we could
620 # add pixelOd column to this table if/when needed.
621 raise NotImplementedError("Region-based selection is not supported")
623 # TODO: DateTime.MJD must be consistent with code in ap_association,
624 # alternatively we can fill midpointMjdTai ourselves in store()
625 if start_time is None:
626 start_time_mjdTai = _make_midpointMjdTai_start(visit_time, self.config.read_forced_sources_months)
627 else:
628 start_time_mjdTai = float(start_time.tai.mjd)
629 _LOG.debug("start_time_mjdTai = %.6f", start_time_mjdTai)
631 with self._timer("select_time", tags={"table": "DiaForcedSource"}) as timer:
632 sources = self._getSourcesByIDs(ApdbTables.DiaForcedSource, list(object_ids), start_time_mjdTai)
633 timer.add_values(row_count=len(sources))
635 _LOG.debug("found %s DiaForcedSources", len(sources))
636 return sources
638 def containsVisitDetector(
639 self,
640 visit: int,
641 detector: int,
642 region: Region,
643 visit_time: astropy.time.Time,
644 ) -> bool:
645 # docstring is inherited from a base class
646 src_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaSource)
647 frcsrc_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaForcedSource)
648 # Query should load only one leaf page of the index
649 query1 = sql.select(src_table.c.visit).filter_by(visit=visit, detector=detector).limit(1)
651 with self._engine.begin() as conn:
652 result = conn.execute(query1).scalar_one_or_none()
653 if result is not None:
654 return True
655 else:
656 # Backup query if an image was processed but had no diaSources
657 query2 = sql.select(frcsrc_table.c.visit).filter_by(visit=visit, detector=detector).limit(1)
658 result = conn.execute(query2).scalar_one_or_none()
659 return result is not None
661 def store(
662 self,
663 visit_time: astropy.time.Time,
664 objects: pandas.DataFrame,
665 sources: pandas.DataFrame | None = None,
666 forced_sources: pandas.DataFrame | None = None,
667 ) -> None:
668 # docstring is inherited from a base class
669 objects = self._fix_input_timestamps(objects)
670 if sources is not None:
671 sources = self._fix_input_timestamps(sources)
672 if forced_sources is not None:
673 forced_sources = self._fix_input_timestamps(forced_sources)
675 # We want to run all inserts in one transaction.
676 with self._engine.begin() as connection:
677 replica_chunk: ReplicaChunk | None = None
678 if self._schema.replication_enabled:
679 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, self.config.replica_chunk_seconds)
680 self._storeReplicaChunk(replica_chunk, connection)
682 # fill pixelId column for DiaObjects
683 objects = self._add_spatial_index(objects)
684 self._storeDiaObjects(objects, visit_time, replica_chunk, connection)
686 if sources is not None:
687 # fill pixelId column for DiaSources
688 sources = self._add_spatial_index(sources)
689 self._storeDiaSources(sources, replica_chunk, connection)
691 if forced_sources is not None:
692 self._storeDiaForcedSources(forced_sources, replica_chunk, connection)
694 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
695 # docstring is inherited from a base class
697 timestamp: float | datetime.datetime
698 if self.schema.has_mjd_timestamps:
699 timestamp_column = "ssObjectReassocTimeMjdTai"
700 timestamp = float(astropy.time.Time.now().tai.mjd)
701 else:
702 timestamp_column = "ssObjectReassocTime"
703 timestamp = datetime.datetime.now(tz=datetime.UTC)
705 table = self._schema.get_table(ApdbTables.DiaSource)
706 query = table.update().where(table.columns["diaSourceId"] == sql.bindparam("srcId"))
708 with self._engine.begin() as conn:
709 # Need to make sure that every ID exists in the database, but
710 # executemany may not support rowcount, so iterate and check what
711 # is missing.
712 missing_ids: list[int] = []
713 for key, value in idMap.items():
714 params = {
715 "srcId": key,
716 "diaObjectId": 0,
717 "ssObjectId": value,
718 timestamp_column: timestamp,
719 }
720 result = conn.execute(query, params)
721 if result.rowcount == 0:
722 missing_ids.append(key)
723 if missing_ids:
724 missing = ",".join(str(item) for item in missing_ids)
725 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
727 def dailyJob(self) -> None:
728 # docstring is inherited from a base class
729 pass
731 def countUnassociatedObjects(self) -> int:
732 # docstring is inherited from a base class
734 # Retrieve the DiaObject table.
735 table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaObject)
737 if self.schema.has_mjd_timestamps:
738 validity_end_column = "validityEndMjdTai"
739 else:
740 validity_end_column = "validityEnd"
742 # Construct the sql statement.
743 stmt = sql.select(func.count()).select_from(table).where(table.c.nDiaSources == 1)
744 stmt = stmt.where(table.columns[validity_end_column] == None) # noqa: E711
746 # Return the count.
747 with self._engine.begin() as conn:
748 count = conn.execute(stmt).scalar_one()
750 return count
752 @property
753 def schema(self) -> ApdbSchema:
754 # docstring is inherited from a base class
755 return self._table_schema
757 @property
758 def metadata(self) -> ApdbMetadata:
759 # docstring is inherited from a base class
760 return self._metadata
762 @property
763 def admin(self) -> ApdbSqlAdmin:
764 # docstring is inherited from a base class
765 return ApdbSqlAdmin(self.pixelator)
767 def _getDiaSourcesInRegion(self, region: Region, start_time_mjdTai: float) -> pandas.DataFrame:
768 """Return catalog of DiaSource instances from given region.
770 Parameters
771 ----------
772 region : `lsst.sphgeom.Region`
773 Region to search for DIASources.
774 start_time_mjdTai : `float`
775 Lower bound of time window for the query.
777 Returns
778 -------
779 catalog : `pandas.DataFrame`
780 Catalog containing DiaSource records.
781 """
782 table = self._schema.get_table(ApdbTables.DiaSource)
783 columns = self._schema.get_apdb_columns(ApdbTables.DiaSource)
784 query = sql.select(*columns)
786 # build selection
787 time_filter = table.columns["midpointMjdTai"] > start_time_mjdTai
788 where = sql.expression.and_(self._filterRegion(table, region), time_filter)
789 query = query.where(where)
791 # execute select
792 with self._timer("DiaSource_select_time", tags={"table": "DiaSource"}) as timer:
793 with self._engine.begin() as conn:
794 sources = pandas.read_sql_query(query, conn)
795 timer.add_values(row_counts=len(sources))
796 _LOG.debug("found %s DiaSources", len(sources))
797 return self._fix_result_timestamps(sources)
799 def _getDiaSourcesByIDs(self, object_ids: list[int], start_time_mjdTai: float) -> pandas.DataFrame:
800 """Return catalog of DiaSource instances given set of DiaObject IDs.
802 Parameters
803 ----------
804 object_ids :
805 Collection of DiaObject IDs
806 start_time_mjdTai : `float`
807 Lower bound of time window for the query.
809 Returns
810 -------
811 catalog : `pandas.DataFrame`
812 Catalog containing DiaSource records.
813 """
814 with self._timer("select_time", tags={"table": "DiaSource"}) as timer:
815 sources = self._getSourcesByIDs(ApdbTables.DiaSource, object_ids, start_time_mjdTai)
816 timer.add_values(row_count=len(sources))
818 _LOG.debug("found %s DiaSources", len(sources))
819 return sources
821 def _getSourcesByIDs(
822 self, table_enum: ApdbTables, object_ids: list[int], midpointMjdTai_start: float
823 ) -> pandas.DataFrame:
824 """Return catalog of DiaSource or DiaForcedSource instances given set
825 of DiaObject IDs.
827 Parameters
828 ----------
829 table : `sqlalchemy.schema.Table`
830 Database table.
831 object_ids :
832 Collection of DiaObject IDs
833 midpointMjdTai_start : `float`
834 Earliest midpointMjdTai to retrieve.
836 Returns
837 -------
838 catalog : `pandas.DataFrame`
839 Catalog contaning DiaSource records. `None` is returned if
840 ``read_sources_months`` configuration parameter is set to 0 or
841 when ``object_ids`` is empty.
842 """
843 table = self._schema.get_table(table_enum)
844 columns = self._schema.get_apdb_columns(table_enum)
846 sources: pandas.DataFrame | None = None
847 if len(object_ids) <= 0:
848 _LOG.debug("ID list is empty, just fetch empty result")
849 query = sql.select(*columns).where(sql.literal(False))
850 with self._engine.begin() as conn:
851 sources = pandas.read_sql_query(query, conn)
852 else:
853 data_frames: list[pandas.DataFrame] = []
854 for ids in chunk_iterable(sorted(object_ids), 1000):
855 query = sql.select(*columns)
857 # Some types like np.int64 can cause issues with
858 # sqlalchemy, convert them to int.
859 int_ids = [int(oid) for oid in ids]
861 # select by object id
862 query = query.where(
863 sql.expression.and_(
864 table.columns["diaObjectId"].in_(int_ids),
865 table.columns["midpointMjdTai"] > midpointMjdTai_start,
866 )
867 )
869 # execute select
870 with self._engine.begin() as conn:
871 data_frames.append(pandas.read_sql_query(query, conn))
873 if len(data_frames) == 1:
874 sources = data_frames[0]
875 else:
876 sources = pandas.concat(data_frames)
877 assert sources is not None, "Catalog cannot be None"
878 return self._fix_result_timestamps(sources)
880 def _storeReplicaChunk(
881 self,
882 replica_chunk: ReplicaChunk,
883 connection: sqlalchemy.engine.Connection,
884 ) -> None:
885 # `visit_time.datetime` returns naive datetime, even though all astropy
886 # times are in UTC. Add UTC timezone to timestamp so that database
887 # can store a correct value.
888 dt = datetime.datetime.fromtimestamp(replica_chunk.last_update_time.unix_tai, tz=datetime.UTC)
890 table = self._schema.get_table(ExtraTables.ApdbReplicaChunks)
892 # We need UPSERT which is dialect-specific construct
893 values = {"last_update_time": dt, "unique_id": replica_chunk.unique_id}
894 row = {"apdb_replica_chunk": replica_chunk.id} | values
895 if connection.dialect.name == "sqlite":
896 insert_sqlite = sqlalchemy.dialects.sqlite.insert(table)
897 insert_sqlite = insert_sqlite.on_conflict_do_update(index_elements=table.primary_key, set_=values)
898 connection.execute(insert_sqlite, row)
899 elif connection.dialect.name == "postgresql":
900 insert_pg = sqlalchemy.dialects.postgresql.dml.insert(table)
901 insert_pg = insert_pg.on_conflict_do_update(constraint=table.primary_key, set_=values)
902 connection.execute(insert_pg, row)
903 else:
904 raise TypeError(f"Unsupported dialect {connection.dialect.name} for upsert.")
906 def _storeDiaObjects(
907 self,
908 objs: pandas.DataFrame,
909 visit_time: astropy.time.Time,
910 replica_chunk: ReplicaChunk | None,
911 connection: sqlalchemy.engine.Connection,
912 ) -> None:
913 """Store catalog of DiaObjects from current visit.
915 Parameters
916 ----------
917 objs : `pandas.DataFrame`
918 Catalog with DiaObject records.
919 visit_time : `astropy.time.Time`
920 Time of the visit.
921 replica_chunk : `ReplicaChunk`
922 Insert identifier.
923 """
924 if len(objs) == 0:
925 _LOG.debug("No objects to write to database.")
926 return
928 # Some types like np.int64 can cause issues with sqlalchemy, convert
929 # them to int.
930 ids = sorted(int(oid) for oid in objs["diaObjectId"])
931 _LOG.debug("first object ID: %d", ids[0])
933 if self.schema.has_mjd_timestamps:
934 validity_start_column = "validityStartMjdTai"
935 validity_end_column = "validityEndMjdTai"
936 timestamp = float(visit_time.tai.mjd)
937 else:
938 validity_start_column = "validityStart"
939 validity_end_column = "validityEnd"
940 timestamp = visit_time.datetime
942 # everything to be done in single transaction
943 if self.config.dia_object_index == "last_object_table":
944 # Insert and replace all records in LAST table.
945 table = self._schema.get_table(ApdbTables.DiaObjectLast)
947 # DiaObjectLast did not have this column in the past.
948 use_validity_start = self._schema.check_column(ApdbTables.DiaObjectLast, validity_start_column)
950 # Drop the previous objects (pandas cannot upsert).
951 query = table.delete().where(table.columns["diaObjectId"].in_(ids))
953 with self._timer("delete_time", tags={"table": table.name}) as timer:
954 res = connection.execute(query)
955 timer.add_values(row_count=res.rowcount)
956 _LOG.debug("deleted %s objects", res.rowcount)
958 # DiaObjectLast is a subset of DiaObject, strip missing columns
959 last_column_names = [column.name for column in table.columns]
960 if validity_start_column in last_column_names and validity_start_column not in objs.columns:
961 last_column_names.remove(validity_start_column)
962 last_objs = objs[last_column_names]
963 last_objs = _coerce_uint64(last_objs)
965 # Fill validityStart, only when it is in the schema.
966 if use_validity_start:
967 if validity_start_column in last_objs:
968 last_objs[validity_start_column] = timestamp
969 else:
970 extra_column = pandas.Series([timestamp] * len(last_objs), name=validity_start_column)
971 last_objs.set_index(extra_column.index, inplace=True)
972 last_objs = pandas.concat([last_objs, extra_column], axis="columns")
974 with self._timer("insert_time", tags={"table": "DiaObjectLast"}) as timer:
975 last_objs.to_sql(
976 table.name,
977 connection,
978 if_exists="append",
979 index=False,
980 schema=table.schema,
981 )
982 timer.add_values(row_count=len(last_objs))
983 else:
984 # truncate existing validity intervals
985 table = self._schema.get_table(ApdbTables.DiaObject)
987 update = (
988 table.update()
989 .values(**{validity_end_column: timestamp})
990 .where(
991 sql.expression.and_(
992 table.columns["diaObjectId"].in_(ids),
993 table.columns[validity_end_column].is_(None),
994 )
995 )
996 )
998 with self._timer("truncate_time", tags={"table": table.name}) as timer:
999 res = connection.execute(update)
1000 timer.add_values(row_count=res.rowcount)
1001 _LOG.debug("truncated %s intervals", res.rowcount)
1003 objs = _coerce_uint64(objs)
1005 # Fill additional columns
1006 extra_columns: list[pandas.Series] = []
1007 if validity_start_column in objs.columns:
1008 objs[validity_start_column] = timestamp
1009 else:
1010 extra_columns.append(pandas.Series([timestamp] * len(objs), name=validity_start_column))
1011 if validity_end_column in objs.columns:
1012 objs[validity_end_column] = None
1013 else:
1014 extra_columns.append(pandas.Series([None] * len(objs), name=validity_end_column))
1015 if extra_columns:
1016 objs.set_index(extra_columns[0].index, inplace=True)
1017 objs = pandas.concat([objs] + extra_columns, axis="columns")
1019 # Insert replica data
1020 table = self._schema.get_table(ApdbTables.DiaObject)
1021 replica_data: list[dict] = []
1022 replica_stmt: Any = None
1023 replica_table_name = ""
1024 if replica_chunk is not None:
1025 pk_names = [column.name for column in table.primary_key]
1026 replica_data = objs[pk_names].to_dict("records")
1027 if replica_data:
1028 for row in replica_data:
1029 row["apdb_replica_chunk"] = replica_chunk.id
1030 replica_table = self._schema.get_table(ExtraTables.DiaObjectChunks)
1031 replica_table_name = replica_table.name
1032 replica_stmt = replica_table.insert()
1034 # insert new versions
1035 with self._timer("insert_time", tags={"table": table.name}) as timer:
1036 objs.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1037 timer.add_values(row_count=len(objs))
1038 if replica_stmt is not None:
1039 with self._timer("insert_time", tags={"table": replica_table_name}) as timer:
1040 connection.execute(replica_stmt, replica_data)
1041 timer.add_values(row_count=len(replica_data))
1043 def _storeDiaSources(
1044 self,
1045 sources: pandas.DataFrame,
1046 replica_chunk: ReplicaChunk | None,
1047 connection: sqlalchemy.engine.Connection,
1048 ) -> None:
1049 """Store catalog of DiaSources from current visit.
1051 Parameters
1052 ----------
1053 sources : `pandas.DataFrame`
1054 Catalog containing DiaSource records
1055 """
1056 table = self._schema.get_table(ApdbTables.DiaSource)
1058 # Insert replica data
1059 replica_data: list[dict] = []
1060 replica_stmt: Any = None
1061 replica_table_name = ""
1062 if replica_chunk is not None:
1063 pk_names = [column.name for column in table.primary_key]
1064 replica_data = sources[pk_names].to_dict("records")
1065 if replica_data:
1066 for row in replica_data:
1067 row["apdb_replica_chunk"] = replica_chunk.id
1068 replica_table = self._schema.get_table(ExtraTables.DiaSourceChunks)
1069 replica_table_name = replica_table.name
1070 replica_stmt = replica_table.insert()
1072 # everything to be done in single transaction
1073 with self._timer("insert_time", tags={"table": table.name}) as timer:
1074 sources = _coerce_uint64(sources)
1075 sources.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1076 timer.add_values(row_count=len(sources))
1077 if replica_stmt is not None:
1078 with self._timer("replica_insert_time", tags={"table": replica_table_name}) as timer:
1079 connection.execute(replica_stmt, replica_data)
1080 timer.add_values(row_count=len(replica_data))
1082 def _storeDiaForcedSources(
1083 self,
1084 sources: pandas.DataFrame,
1085 replica_chunk: ReplicaChunk | None,
1086 connection: sqlalchemy.engine.Connection,
1087 ) -> None:
1088 """Store a set of DiaForcedSources from current visit.
1090 Parameters
1091 ----------
1092 sources : `pandas.DataFrame`
1093 Catalog containing DiaForcedSource records
1094 """
1095 table = self._schema.get_table(ApdbTables.DiaForcedSource)
1097 # Insert replica data
1098 replica_data: list[dict] = []
1099 replica_stmt: Any = None
1100 replica_table_name = ""
1101 if replica_chunk is not None:
1102 pk_names = [column.name for column in table.primary_key]
1103 replica_data = sources[pk_names].to_dict("records")
1104 if replica_data:
1105 for row in replica_data:
1106 row["apdb_replica_chunk"] = replica_chunk.id
1107 replica_table = self._schema.get_table(ExtraTables.DiaForcedSourceChunks)
1108 replica_table_name = replica_table.name
1109 replica_stmt = replica_table.insert()
1111 # everything to be done in single transaction
1112 with self._timer("insert_time", tags={"table": table.name}) as timer:
1113 sources = _coerce_uint64(sources)
1114 sources.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1115 timer.add_values(row_count=len(sources))
1116 if replica_stmt is not None:
1117 with self._timer("insert_time", tags={"table": replica_table_name}) as timer:
1118 connection.execute(replica_stmt, replica_data)
1119 timer.add_values(row_count=len(replica_data))
1121 def _storeUpdateRecords(
1122 self,
1123 records: Iterable[ApdbUpdateRecord],
1124 chunk: ReplicaChunk,
1125 *,
1126 store_chunk: bool = False,
1127 connection: sqlalchemy.engine.Connection | None = None,
1128 ) -> None:
1129 """Store ApdbUpdateRecords in the replica table for those records.
1131 Parameters
1132 ----------
1133 records : `list` [`ApdbUpdateRecord`]
1134 Records to store.
1135 chunk : `ReplicaChunk`
1136 Replica chunk for these records.
1137 store_chunk : `bool`
1138 If True then also store replica chunk.
1139 connection : `sqlalchemy.engine.Connection`
1140 SQLALchemy connection to use, if `None` the new connection will be
1141 made. `None` is useful for tests only, regular use will call this
1142 method in the same transaction that saves other types of records.
1144 Raises
1145 ------
1146 TypeError
1147 Raised if replication is not enabled for this instance.
1148 """
1149 if not self._schema.replication_enabled:
1150 raise TypeError("Replication is not enabled for this APDB instance.")
1152 apdb_replica_chunk = chunk.id
1153 # Do not use unique_if from ReplicaChunk as it could be reused in
1154 # multiple calls to this method.
1155 update_unique_id = uuid.uuid4()
1157 record_dicts = []
1158 for record in records:
1159 record_dicts.append(
1160 {
1161 "apdb_replica_chunk": apdb_replica_chunk,
1162 "update_time_ns": record.update_time_ns,
1163 "update_order": record.update_order,
1164 "update_unique_id": update_unique_id,
1165 "update_payload": record.to_json(),
1166 }
1167 )
1169 if not record_dicts:
1170 return
1172 # TODO: Need to check that table exists.
1173 table = self._schema.get_table(ExtraTables.ApdbUpdateRecordChunks)
1175 def _do_store(connection: sqlalchemy.engine.Connection) -> None:
1176 if store_chunk:
1177 self._storeReplicaChunk(chunk, connection)
1178 with self._timer("insert_time", tags={"table": table.name}) as timer:
1179 connection.execute(table.insert(), record_dicts)
1180 timer.add_values(row_count=len(record_dicts))
1182 if connection is None:
1183 with self._engine.begin() as connection:
1184 _do_store(connection)
1185 else:
1186 _do_store(connection)
1188 def _htm_indices(self, region: Region) -> list[tuple[int, int]]:
1189 """Generate a set of HTM indices covering specified region.
1191 Parameters
1192 ----------
1193 region: `sphgeom.Region`
1194 Region that needs to be indexed.
1196 Returns
1197 -------
1198 Sequence of ranges, range is a tuple (minHtmID, maxHtmID).
1199 """
1200 _LOG.debug("region: %s", region)
1201 indices = self.pixelator.envelope(region, self.config.pixelization.htm_max_ranges)
1203 return indices.ranges()
1205 def _filterRegion(self, table: sqlalchemy.schema.Table, region: Region) -> sql.ColumnElement:
1206 """Make SQLAlchemy expression for selecting records in a region."""
1207 htm_index_column = table.columns[self.config.pixelization.htm_index_column]
1208 exprlist = []
1209 pixel_ranges = self._htm_indices(region)
1210 for low, upper in pixel_ranges:
1211 upper -= 1
1212 if low == upper:
1213 exprlist.append(htm_index_column == low)
1214 else:
1215 exprlist.append(sql.expression.between(htm_index_column, low, upper))
1217 return sql.expression.or_(*exprlist)
1219 def _add_spatial_index(self, df: pandas.DataFrame) -> pandas.DataFrame:
1220 """Calculate spatial index for each record and add it to a DataFrame.
1222 Parameters
1223 ----------
1224 df : `pandas.DataFrame`
1225 DataFrame which has to contain ra/dec columns, names of these
1226 columns are defined by configuration ``ra_dec_columns`` field.
1228 Returns
1229 -------
1230 df : `pandas.DataFrame`
1231 DataFrame with ``pixelId`` column which contains pixel index
1232 for ra/dec coordinates.
1234 Notes
1235 -----
1236 This overrides any existing column in a DataFrame with the same name
1237 (pixelId). Original DataFrame is not changed, copy of a DataFrame is
1238 returned.
1239 """
1240 # calculate HTM index for every DiaObject
1241 htm_index = np.zeros(df.shape[0], dtype=np.int64)
1242 ra_col, dec_col = self.config.ra_dec_columns
1243 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1244 uv3d = UnitVector3d(LonLat.fromDegrees(ra, dec))
1245 idx = self.pixelator.index(uv3d)
1246 htm_index[i] = idx
1247 df = df.copy()
1248 df[self.config.pixelization.htm_index_column] = htm_index
1249 return df
1251 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1252 """Update timestamp columns in input DataFrame to be aware datetime
1253 type in in UTC.
1255 AP pipeline generates naive datetime instances, we want them to be
1256 aware before they go to database. All naive timestamps are assumed to
1257 be in UTC timezone (they should be TAI).
1258 """
1259 # Find all columns with aware non-UTC timestamps and convert to UTC.
1260 columns = [
1261 column
1262 for column, dtype in df.dtypes.items()
1263 if isinstance(dtype, pandas.DatetimeTZDtype) and dtype.tz is not datetime.UTC
1264 ]
1265 for column in columns:
1266 df[column] = df[column].dt.tz_convert(datetime.UTC)
1267 # Find all columns with naive timestamps and add UTC timezone.
1268 columns = [
1269 column for column, dtype in df.dtypes.items() if pandas.api.types.is_datetime64_dtype(dtype)
1270 ]
1271 for column in columns:
1272 df[column] = df[column].dt.tz_localize(datetime.UTC)
1273 return df
1275 def _fix_result_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1276 """Update timestamp columns to be naive datetime type in returned
1277 DataFrame.
1279 AP pipeline code expects DataFrames to contain naive datetime columns,
1280 while Postgres queries return timezone-aware type. This method converts
1281 those columns to naive datetime in UTC timezone.
1282 """
1283 # Find all columns with aware timestamps.
1284 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1285 for column in columns:
1286 # tz_convert(None) will convert to UTC and drop timezone.
1287 df[column] = df[column].dt.tz_convert(None)
1288 return df