Coverage for python / lsst / dax / apdb / sql / apdbSql.py: 12%
688 statements
« prev ^ index » next coverage.py v7.13.1, created at 2026-01-22 02:00 -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 json
30import logging
31import urllib.parse
32import uuid
33import warnings
34from collections import Counter
35from collections.abc import Iterable, Mapping, MutableMapping
36from contextlib import closing
37from typing import TYPE_CHECKING, Any
39import astropy.time
40import numpy as np
41import pandas
42import sqlalchemy
43import sqlalchemy.dialects.postgresql
44import sqlalchemy.dialects.sqlite
45from sqlalchemy import func, sql
46from sqlalchemy.pool import NullPool
48from lsst.sphgeom import HtmPixelization, LonLat, Region, UnitVector3d
49from lsst.utils.db_auth import DbAuth, DbAuthNotFoundError
50from lsst.utils.iteration import chunk_iterable
52from ..apdb import Apdb
53from ..apdbConfigFreezer import ApdbConfigFreezer
54from ..apdbReplica import ReplicaChunk
55from ..apdbSchema import ApdbSchema, ApdbTables
56from ..apdbUpdateRecord import (
57 ApdbCloseDiaObjectValidityRecord,
58 ApdbReassignDiaSourceToDiaObjectRecord,
59 ApdbUpdateNDiaSourcesRecord,
60 ApdbUpdateRecord,
61)
62from ..config import ApdbConfig
63from ..monitor import MonAgent
64from ..recordIds import DiaObjectId, DiaSourceId
65from ..schema_model import Table
66from ..timer import Timer
67from ..versionTuple import IncompatibleVersionError, VersionTuple
68from .apdbMetadataSql import ApdbMetadataSql
69from .apdbSqlAdmin import ApdbSqlAdmin
70from .apdbSqlReplica import ApdbSqlReplica
71from .apdbSqlSchema import ApdbSqlSchema, ExtraTables
72from .config import ApdbSqlConfig
74if TYPE_CHECKING:
75 import sqlite3
77 from ..apdbMetadata import ApdbMetadata
78 from ..apdbUpdateRecord import ApdbUpdateRecord
80_LOG = logging.getLogger(__name__)
82_MON = MonAgent(__name__)
84VERSION = VersionTuple(1, 2, 1)
85"""Version for the code controlling non-replication tables. This needs to be
86updated following compatibility rules when schema produced by this code
87changes.
88"""
91def _coerce_uint64(df: pandas.DataFrame) -> pandas.DataFrame:
92 """Change the type of uint64 columns to int64, and return copy of data
93 frame.
94 """
95 names = [c[0] for c in df.dtypes.items() if c[1] == np.uint64]
96 return df.astype(dict.fromkeys(names, np.int64))
99def _make_midpointMjdTai_start(visit_time: astropy.time.Time, months: int) -> float:
100 """Calculate starting point for time-based source search.
102 Parameters
103 ----------
104 visit_time : `astropy.time.Time`
105 Time of current visit.
106 months : `int`
107 Number of months in the sources history.
109 Returns
110 -------
111 time : `float`
112 A ``midpointMjdTai`` starting point, MJD time.
113 """
114 # TODO: Use of MJD must be consistent with the code in ap_association
115 # (see DM-31996)
116 return float(visit_time.tai.mjd) - months * 30
119def _onSqlite3Connect(
120 dbapiConnection: sqlite3.Connection, connectionRecord: sqlalchemy.pool._ConnectionRecord
121) -> None:
122 # Enable foreign keys
123 with closing(dbapiConnection.cursor()) as cursor:
124 cursor.execute("PRAGMA foreign_keys=ON;")
127class ApdbSql(Apdb):
128 """Implementation of APDB interface based on SQL database.
130 The implementation is configured via standard ``pex_config`` mechanism
131 using `ApdbSqlConfig` configuration class. For an example of different
132 configurations check ``config/`` folder.
134 Parameters
135 ----------
136 config : `ApdbSqlConfig`
137 Configuration object.
138 """
140 metadataSchemaVersionKey = "version:schema"
141 """Name of the metadata key to store schema version number."""
143 metadataCodeVersionKey = "version:ApdbSql"
144 """Name of the metadata key to store code version number."""
146 metadataReplicaVersionKey = "version:ApdbSqlReplica"
147 """Name of the metadata key to store replica code version number."""
149 metadataConfigKey = "config:apdb-sql.json"
150 """Name of the metadata key to store code version number."""
152 metadataDedupKey = "status:deduplication.json"
153 """Name of the metadata key to store code version number."""
155 _frozen_parameters = (
156 "enable_replica",
157 "dia_object_index",
158 "pixelization.htm_level",
159 "pixelization.htm_index_column",
160 "ra_dec_columns",
161 )
162 """Names of the config parameters to be frozen in metadata table."""
164 def __init__(self, config: ApdbSqlConfig):
165 self._engine = self._makeEngine(config, create=False)
167 sa_metadata = sqlalchemy.MetaData(schema=config.namespace)
168 meta_table_name = ApdbTables.metadata.table_name(prefix=config.prefix)
169 meta_table = sqlalchemy.schema.Table(meta_table_name, sa_metadata, autoload_with=self._engine)
170 self._metadata = ApdbMetadataSql(self._engine, meta_table)
172 # Get tables schemas.
173 self._table_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
175 # Check that versions are compatible, must be the first thing before
176 # reading frozen config.
177 self._db_schema_version = self._versionCheck(self._metadata, self._table_schema.schemaVersion())
179 # Read frozen config from metadata.
180 config_json = self._metadata.get(self.metadataConfigKey)
181 if config_json is not None:
182 # Update config from metadata.
183 freezer = ApdbConfigFreezer[ApdbSqlConfig](self._frozen_parameters)
184 self.config = freezer.update(config, config_json)
185 else:
186 self.config = config
188 self._schema = ApdbSqlSchema(
189 table_schema=self._table_schema,
190 engine=self._engine,
191 dia_object_index=self.config.dia_object_index,
192 prefix=self.config.prefix,
193 namespace=self.config.namespace,
194 htm_index_column=self.config.pixelization.htm_index_column,
195 enable_replica=self.config.enable_replica,
196 )
198 self.pixelator = HtmPixelization(self.config.pixelization.htm_level)
200 if _LOG.isEnabledFor(logging.DEBUG):
201 _LOG.debug("ApdbSql Configuration: %s", self.config.model_dump())
203 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
204 """Create `Timer` instance given its name."""
205 return Timer(name, _MON, tags=tags)
207 @classmethod
208 def _makeEngine(cls, config: ApdbSqlConfig, *, create: bool) -> sqlalchemy.engine.Engine:
209 """Make SQLALchemy engine based on configured parameters.
211 Parameters
212 ----------
213 config : `ApdbSqlConfig`
214 Configuration object.
215 create : `bool`
216 Whether to try to create new database file, only relevant for
217 SQLite backend which always creates new files by default.
218 """
219 # engine is reused between multiple processes, make sure that we don't
220 # share connections by disabling pool (by using NullPool class)
221 kw: MutableMapping[str, Any] = dict(config.connection_config.extra_parameters)
222 conn_args: dict[str, Any] = {}
223 if not config.connection_config.connection_pool:
224 kw.update(poolclass=NullPool)
225 if config.connection_config.isolation_level is not None:
226 kw.update(isolation_level=config.connection_config.isolation_level)
227 elif config.db_url.startswith("sqlite"):
228 # Use READ_UNCOMMITTED as default value for sqlite.
229 kw.update(isolation_level="READ_UNCOMMITTED")
230 if config.connection_config.connection_timeout is not None:
231 if config.db_url.startswith("sqlite"):
232 conn_args.update(timeout=config.connection_config.connection_timeout)
233 elif config.db_url.startswith(("postgresql", "mysql")):
234 conn_args.update(connect_timeout=int(config.connection_config.connection_timeout))
235 kw.update(connect_args=conn_args)
236 engine = sqlalchemy.create_engine(cls._connection_url(config.db_url, create=create), **kw)
238 if engine.dialect.name == "sqlite":
239 # Need to enable foreign keys on every new connection.
240 sqlalchemy.event.listen(engine, "connect", _onSqlite3Connect)
242 return engine
244 @classmethod
245 def _connection_url(cls, config_url: str, *, create: bool) -> sqlalchemy.engine.URL | str:
246 """Generate a complete URL for database with proper credentials.
248 Parameters
249 ----------
250 config_url : `str`
251 Database URL as specified in configuration.
252 create : `bool`
253 Whether to try to create new database file, only relevant for
254 SQLite backend which always creates new files by default.
256 Returns
257 -------
258 connection_url : `sqlalchemy.engine.URL` or `str`
259 Connection URL including credentials.
260 """
261 # Allow 3rd party authentication mechanisms by assuming connection
262 # string is correct when we can not recognize (dialect, host, database)
263 # matching keys.
264 components = urllib.parse.urlparse(config_url)
265 if all((components.scheme is not None, components.hostname is not None, components.path is not None)):
266 try:
267 db_auth = DbAuth()
268 config_url = db_auth.getUrl(config_url)
269 except DbAuthNotFoundError:
270 # Credentials file doesn't exist or no matching credentials,
271 # use default auth.
272 pass
274 # SQLite has a nasty habit creating empty databases when they do not
275 # exist, tell it not to do that unless we do need to create it.
276 if not create:
277 config_url = cls._update_sqlite_url(config_url)
279 return config_url
281 @classmethod
282 def _update_sqlite_url(cls, url_string: str) -> str:
283 """If URL refers to sqlite dialect, update it so that the backend does
284 not try to create database file if it does not exist already.
286 Parameters
287 ----------
288 url_string : `str`
289 Connection string.
291 Returns
292 -------
293 url_string : `str`
294 Possibly updated connection string.
295 """
296 try:
297 url = sqlalchemy.make_url(url_string)
298 except sqlalchemy.exc.SQLAlchemyError:
299 # If parsing fails it means some special format, likely not
300 # sqlite so we just return it unchanged.
301 return url_string
303 if url.get_backend_name() == "sqlite":
304 # Massage url so that database name starts with "file:" and
305 # option string has "mode=rw&uri=true". Database name
306 # should look like a path (:memory: is not supported by
307 # Apdb, but someone could still try to use it).
308 database = url.database
309 if database and not database.startswith((":", "file:")):
310 query = dict(url.query, mode="rw", uri="true")
311 # If ``database`` is an absolute path then original URL should
312 # include four slashes after "sqlite:". Humans are bad at
313 # counting things beyond four and sometimes an extra slash gets
314 # added unintentionally, which causes sqlite to treat initial
315 # element as "authority" and to complain. Strip extra slashes
316 # at the start of the path to avoid that (DM-46077).
317 if database.startswith("//"):
318 warnings.warn(
319 f"Database URL contains extra leading slashes which will be removed: {url}",
320 stacklevel=3,
321 )
322 database = "/" + database.lstrip("/")
323 url = url.set(database=f"file:{database}", query=query)
324 url_string = url.render_as_string()
326 return url_string
328 @classmethod
329 def _versionCheck(cls, metadata: ApdbMetadataSql, schema_version: VersionTuple) -> VersionTuple:
330 """Check schema version compatibility and return the database schema
331 version.
332 """
334 def _get_version(key: str) -> VersionTuple:
335 """Retrieve version number from given metadata key."""
336 version_str = metadata.get(key)
337 if version_str is None:
338 # Should not happen with existing metadata table.
339 raise RuntimeError(f"Version key {key!r} does not exist in metadata table.")
340 return VersionTuple.fromString(version_str)
342 db_schema_version = _get_version(cls.metadataSchemaVersionKey)
343 db_code_version = _get_version(cls.metadataCodeVersionKey)
345 # For now there is no way to make read-only APDB instances, assume that
346 # any access can do updates.
347 if not schema_version.checkCompatibility(db_schema_version):
348 raise IncompatibleVersionError(
349 f"Configured schema version {schema_version} "
350 f"is not compatible with database version {db_schema_version}"
351 )
352 if not cls.apdbImplementationVersion().checkCompatibility(db_code_version):
353 raise IncompatibleVersionError(
354 f"Current code version {cls.apdbImplementationVersion()} "
355 f"is not compatible with database version {db_code_version}"
356 )
358 # Check replica code version only if replica is enabled. Sort of
359 # chicken and egg problem - `enable_replica` is a part of frozen
360 # configuration, but we cannot read frozen configuration until we
361 # validate versions. Assume that if the replica version is present
362 # then replication is enabled.
363 if metadata.get(cls.metadataReplicaVersionKey) is not None:
364 db_replica_version = _get_version(cls.metadataReplicaVersionKey)
365 code_replica_version = ApdbSqlReplica.apdbReplicaImplementationVersion()
366 if not code_replica_version.checkCompatibility(db_replica_version):
367 raise IncompatibleVersionError(
368 f"Current replication code version {code_replica_version} "
369 f"is not compatible with database version {db_replica_version}"
370 )
372 return db_schema_version
374 @classmethod
375 def apdbImplementationVersion(cls) -> VersionTuple:
376 """Return version number for current APDB implementation.
378 Returns
379 -------
380 version : `VersionTuple`
381 Version of the code defined in implementation class.
382 """
383 return VERSION
385 @classmethod
386 def init_database(
387 cls,
388 db_url: str,
389 *,
390 schema_file: str | None = None,
391 ss_schema_file: str | None = None,
392 read_sources_months: int | None = None,
393 read_forced_sources_months: int | None = None,
394 enable_replica: bool = False,
395 connection_timeout: int | None = None,
396 dia_object_index: str | None = None,
397 htm_level: int | None = None,
398 htm_index_column: str | None = None,
399 ra_dec_columns: tuple[str, str] | None = None,
400 prefix: str | None = None,
401 namespace: str | None = None,
402 drop: bool = False,
403 ) -> ApdbSqlConfig:
404 """Initialize new APDB instance and make configuration object for it.
406 Parameters
407 ----------
408 db_url : `str`
409 SQLAlchemy database URL.
410 schema_file : `str`, optional
411 Location of (YAML) configuration file with APDB schema. If not
412 specified then default location will be used.
413 ss_schema_file : `str`, optional
414 Location of (YAML) configuration file with SSO schema. If not
415 specified then default location will be used.
416 read_sources_months : `int`, optional
417 Number of months of history to read from DiaSource.
418 read_forced_sources_months : `int`, optional
419 Number of months of history to read from DiaForcedSource.
420 enable_replica : `bool`, optional
421 If True, make additional tables used for replication to PPDB.
422 connection_timeout : `int`, optional
423 Database connection timeout in seconds.
424 dia_object_index : `str`, optional
425 Indexing mode for DiaObject table.
426 htm_level : `int`, optional
427 HTM indexing level.
428 htm_index_column : `str`, optional
429 Name of a HTM index column for DiaObject and DiaSource tables.
430 ra_dec_columns : `tuple` [`str`, `str`], optional
431 Names of ra/dec columns in DiaObject table.
432 prefix : `str`, optional
433 Optional prefix for all table names.
434 namespace : `str`, optional
435 Name of the database schema for all APDB tables. If not specified
436 then default schema is used.
437 drop : `bool`, optional
438 If `True` then drop existing tables before re-creating the schema.
440 Returns
441 -------
442 config : `ApdbSqlConfig`
443 Resulting configuration object for a created APDB instance.
444 """
445 config = ApdbSqlConfig(db_url=db_url, enable_replica=enable_replica)
446 if schema_file is not None:
447 config.schema_file = schema_file
448 if ss_schema_file is not None:
449 config.ss_schema_file = ss_schema_file
450 if read_sources_months is not None:
451 config.read_sources_months = read_sources_months
452 if read_forced_sources_months is not None:
453 config.read_forced_sources_months = read_forced_sources_months
454 if connection_timeout is not None:
455 config.connection_config.connection_timeout = connection_timeout
456 if dia_object_index is not None:
457 config.dia_object_index = dia_object_index
458 if htm_level is not None:
459 config.pixelization.htm_level = htm_level
460 if htm_index_column is not None:
461 config.pixelization.htm_index_column = htm_index_column
462 if ra_dec_columns is not None:
463 config.ra_dec_columns = ra_dec_columns
464 if prefix is not None:
465 config.prefix = prefix
466 if namespace is not None:
467 config.namespace = namespace
469 cls._makeSchema(config, drop=drop)
471 # SQLite has a nasty habit of creating empty database by default,
472 # update URL in config file to disable that behavior.
473 config.db_url = cls._update_sqlite_url(config.db_url)
475 return config
477 def get_replica(self) -> ApdbSqlReplica:
478 """Return `ApdbReplica` instance for this database."""
479 return ApdbSqlReplica(self._schema, self._engine, self._db_schema_version)
481 def tableRowCount(self) -> dict[str, int]:
482 """Return dictionary with the table names and row counts.
484 Used by ``ap_proto`` to keep track of the size of the database tables.
485 Depending on database technology this could be expensive operation.
487 Returns
488 -------
489 row_counts : `dict`
490 Dict where key is a table name and value is a row count.
491 """
492 res = {}
493 tables = [ApdbTables.DiaObject, ApdbTables.DiaSource, ApdbTables.DiaForcedSource]
494 if self.config.dia_object_index == "last_object_table":
495 tables.append(ApdbTables.DiaObjectLast)
496 with self._engine.begin() as conn:
497 for table in tables:
498 sa_table = self._schema.get_table(table)
499 stmt = sql.select(func.count()).select_from(sa_table)
500 count: int = conn.execute(stmt).scalar_one()
501 res[table.name] = count
503 return res
505 def getConfig(self) -> ApdbSqlConfig:
506 # docstring is inherited from a base class
507 return self.config
509 def tableDef(self, table: ApdbTables) -> Table | None:
510 # docstring is inherited from a base class
511 return self._schema.tableSchemas.get(table)
513 @classmethod
514 def _makeSchema(cls, config: ApdbConfig, drop: bool = False) -> None:
515 # docstring is inherited from a base class
517 if not isinstance(config, ApdbSqlConfig):
518 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
520 engine = cls._makeEngine(config, create=True)
522 table_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
524 # Ask schema class to create all tables.
525 schema = ApdbSqlSchema(
526 table_schema=table_schema,
527 engine=engine,
528 dia_object_index=config.dia_object_index,
529 prefix=config.prefix,
530 namespace=config.namespace,
531 htm_index_column=config.pixelization.htm_index_column,
532 enable_replica=config.enable_replica,
533 )
534 schema.makeSchema(drop=drop)
536 # Need metadata table to store few items in it.
537 meta_table = schema.get_table(ApdbTables.metadata)
538 apdb_meta = ApdbMetadataSql(engine, meta_table)
540 # Fill version numbers, overwrite if they are already there.
541 apdb_meta.set(cls.metadataSchemaVersionKey, str(table_schema.schemaVersion()), force=True)
542 apdb_meta.set(cls.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True)
543 if config.enable_replica:
544 # Only store replica code version if replica is enabled.
545 apdb_meta.set(
546 cls.metadataReplicaVersionKey,
547 str(ApdbSqlReplica.apdbReplicaImplementationVersion()),
548 force=True,
549 )
551 # Store frozen part of a configuration in metadata.
552 freezer = ApdbConfigFreezer[ApdbSqlConfig](cls._frozen_parameters)
553 apdb_meta.set(cls.metadataConfigKey, freezer.to_json(config), force=True)
555 def getDiaObjects(self, region: Region) -> pandas.DataFrame:
556 # docstring is inherited from a base class
558 # decide what columns we need
559 if self.config.dia_object_index == "last_object_table":
560 table_enum = ApdbTables.DiaObjectLast
561 else:
562 table_enum = ApdbTables.DiaObject
563 table = self._schema.get_table(table_enum)
564 if not self.config.dia_object_columns:
565 columns = self._schema.get_apdb_columns(table_enum)
566 else:
567 columns = [table.c[col] for col in self.config.dia_object_columns]
568 query = sql.select(*columns)
570 # build selection
571 query = query.where(self._filterRegion(table, region))
573 validity_end_column = self._timestamp_column_name("validityEnd")
575 # select latest version of objects
576 if self.config.dia_object_index != "last_object_table":
577 query = query.where(table.columns[validity_end_column] == None) # noqa: E711
579 # _LOG.debug("query: %s", query)
581 # execute select
582 with self._timer("select_time", tags={"table": "DiaObject"}) as timer:
583 with self._engine.begin() as conn:
584 objects = pandas.read_sql_query(query, conn)
585 timer.add_values(row_count=len(objects))
586 _LOG.debug("found %s DiaObjects", len(objects))
587 return self._fix_result_timestamps(objects)
589 def getDiaSources(
590 self,
591 region: Region,
592 object_ids: Iterable[int] | None,
593 visit_time: astropy.time.Time,
594 start_time: astropy.time.Time | None = None,
595 ) -> pandas.DataFrame | None:
596 # docstring is inherited from a base class
597 if start_time is None and self.config.read_sources_months == 0:
598 _LOG.debug("Skip DiaSources fetching")
599 return None
601 if start_time is None:
602 start_time_mjdTai = _make_midpointMjdTai_start(visit_time, self.config.read_sources_months)
603 else:
604 start_time_mjdTai = float(start_time.tai.mjd)
605 _LOG.debug("start_time_mjdTai = %.6f", start_time_mjdTai)
607 if object_ids is None:
608 # region-based select
609 return self._getDiaSourcesInRegion(region, start_time_mjdTai)
610 else:
611 return self._getDiaSourcesByIDs(list(object_ids), start_time_mjdTai)
613 def getDiaForcedSources(
614 self,
615 region: Region,
616 object_ids: Iterable[int] | None,
617 visit_time: astropy.time.Time,
618 start_time: astropy.time.Time | None = None,
619 ) -> pandas.DataFrame | None:
620 # docstring is inherited from a base class
621 if start_time is None and self.config.read_forced_sources_months == 0:
622 _LOG.debug("Skip DiaForceSources fetching")
623 return None
625 if object_ids is None:
626 # This implementation does not support region-based selection. In
627 # the past DiaForcedSource schema did not have ra/dec columns (it
628 # had x/y columns). ra/dec were added at some point, so we could
629 # add pixelOd column to this table if/when needed.
630 raise NotImplementedError("Region-based selection is not supported")
632 # TODO: DateTime.MJD must be consistent with code in ap_association,
633 # alternatively we can fill midpointMjdTai ourselves in store()
634 if start_time is None:
635 start_time_mjdTai = _make_midpointMjdTai_start(visit_time, self.config.read_forced_sources_months)
636 else:
637 start_time_mjdTai = float(start_time.tai.mjd)
638 _LOG.debug("start_time_mjdTai = %.6f", start_time_mjdTai)
640 with self._timer("select_time", tags={"table": "DiaForcedSource"}) as timer:
641 sources = self._getSourcesByIDs(ApdbTables.DiaForcedSource, list(object_ids), start_time_mjdTai)
642 timer.add_values(row_count=len(sources))
644 _LOG.debug("found %s DiaForcedSources", len(sources))
645 return sources
647 def getDiaObjectsForDedup(self, since: astropy.time.Time | None = None) -> pandas.DataFrame:
648 # docstring is inherited from a base class
650 if since is None:
651 # Read last deduplication time from metadata.
652 dedup_str = self._metadata.get(self.metadataDedupKey)
653 if dedup_str is not None:
654 dedup_state = json.loads(dedup_str)
655 dedup_time_str = dedup_state["dedup_time_iso_tai"]
656 since = astropy.time.Time(dedup_time_str, format="iso", scale="tai")
658 validity_start_column = self._timestamp_column_name("validityStart")
660 # decide what columns we need
661 if self.config.dia_object_index == "last_object_table":
662 table_enum = ApdbTables.DiaObjectLast
663 else:
664 table_enum = ApdbTables.DiaObject
665 table = self._schema.get_table(table_enum)
667 if not self.config.dia_object_columns_for_dedup:
668 columns = self._schema.get_apdb_columns(table_enum)
669 else:
670 column_names = list(self.config.dia_object_columns_for_dedup)
671 if validity_start_column not in column_names:
672 column_names.insert(0, validity_start_column)
673 if "diaObjectId" not in column_names:
674 column_names.insert(0, "diaObjectId")
675 columns = [table.columns[col] for col in column_names]
677 query = sql.select(*columns)
679 # build selection
680 if since is not None:
681 timestamp = self._timestamp_column_value(since)
682 query = query.where(table.columns[validity_start_column] >= timestamp)
684 # execute select
685 with self._timer("select_time", tags={"table": "DiaObject"}) as timer:
686 with self._engine.begin() as conn:
687 objects = pandas.read_sql_query(query, conn)
688 timer.add_values(row_count=len(objects))
689 _LOG.debug("found %s DiaObjects", len(objects))
690 return self._fix_result_timestamps(objects)
692 def getDiaSourcesForDiaObjects(
693 self, objects: list[DiaObjectId], start_time: astropy.time.Time, max_dist_arcsec: float = 1.0
694 ) -> pandas.DataFrame:
695 # docstring is inherited from a base class
696 object_ids = {object_id.diaObjectId for object_id in objects}
697 return self._getDiaSourcesByIDs(list(object_ids), float(start_time.tai.mjd))
699 def containsVisitDetector(
700 self,
701 visit: int,
702 detector: int,
703 region: Region,
704 visit_time: astropy.time.Time,
705 ) -> bool:
706 # docstring is inherited from a base class
707 src_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaSource)
708 frcsrc_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaForcedSource)
709 # Query should load only one leaf page of the index
710 query1 = sql.select(src_table.c.visit).filter_by(visit=visit, detector=detector).limit(1)
712 with self._engine.begin() as conn:
713 result = conn.execute(query1).scalar_one_or_none()
714 if result is not None:
715 return True
716 else:
717 # Backup query if an image was processed but had no diaSources
718 query2 = sql.select(frcsrc_table.c.visit).filter_by(visit=visit, detector=detector).limit(1)
719 result = conn.execute(query2).scalar_one_or_none()
720 return result is not None
722 def store(
723 self,
724 visit_time: astropy.time.Time,
725 objects: pandas.DataFrame,
726 sources: pandas.DataFrame | None = None,
727 forced_sources: pandas.DataFrame | None = None,
728 ) -> None:
729 # docstring is inherited from a base class
730 objects = self._fix_input_timestamps(objects)
731 if sources is not None:
732 sources = self._fix_input_timestamps(sources)
733 if forced_sources is not None:
734 forced_sources = self._fix_input_timestamps(forced_sources)
736 # We want to run all inserts in one transaction.
737 with self._engine.begin() as connection:
738 replica_chunk: ReplicaChunk | None = None
739 if self._schema.replication_enabled:
740 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, self.config.replica_chunk_seconds)
741 self._storeReplicaChunk(replica_chunk, connection)
743 # fill pixelId column for DiaObjects
744 objects = self._add_spatial_index(objects)
745 self._storeDiaObjects(objects, visit_time, replica_chunk, connection)
747 if sources is not None:
748 # fill pixelId column for DiaSources
749 sources = self._add_spatial_index(sources)
750 self._storeDiaSources(sources, replica_chunk, connection)
752 if forced_sources is not None:
753 self._storeDiaForcedSources(forced_sources, replica_chunk, connection)
755 def reassignDiaSourcesToDiaObjects(
756 self,
757 idMap: Mapping[DiaSourceId, int],
758 *,
759 increment_nDiaSources: bool = True,
760 decrement_nDiaSources: bool = True,
761 ) -> None:
762 # docstring is inherited from a base class
764 new_object_ids = set(idMap.values())
765 source_ids = {source.diaSourceId for source in idMap}
767 current_time = self._current_time()
768 current_time_ns = int(current_time.unix_tai * 1e9)
770 with self._engine.begin() as conn:
771 # Make sure that all DiaSources exist.
772 found_sources = self._get_diasource_data(conn, source_ids, "diaObjectId")
773 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}):
774 raise LookupError(f"Some source IDs are missing from DiaSource table: {missing_ids}")
775 original_object_ids = {row.diaSourceId: row.diaObjectId for row in found_sources}
777 # Make sure that all DiaObjects exist, we also want to know
778 # nDiaSources count for current and new records because we want to
779 # send updated values to replica.
780 all_object_ids = new_object_ids | set(original_object_ids.values())
781 found_objects = self._get_diaobject_data(conn, all_object_ids, "ra", "dec", "nDiaSources")
782 if missing_ids := (new_object_ids - {row.diaObjectId for row in found_objects}):
783 raise LookupError(f"Some object IDs are missing from DiaObject table: {missing_ids}")
785 found_objects_by_id = {row.diaObjectId: row for row in found_objects}
787 update_records: list[ApdbUpdateRecord] = []
788 update_order = 0
790 # Update DiaSources.
791 source_table = self._schema.get_table(ApdbTables.DiaSource)
792 for source, diaObjectId in idMap.items():
793 update = (
794 source_table.update()
795 .where(source_table.columns["diaSourceId"] == source.diaSourceId)
796 .values(diaObjectId=diaObjectId)
797 )
798 conn.execute(update)
800 if self._schema.replication_enabled:
801 update_records.append(
802 ApdbReassignDiaSourceToDiaObjectRecord(
803 diaSourceId=source.diaSourceId,
804 ra=source.ra,
805 dec=source.dec,
806 midpointMjdTai=source.midpointMjdTai,
807 diaObjectId=diaObjectId,
808 update_time_ns=current_time_ns,
809 update_order=update_order,
810 )
811 )
812 update_order += 1
814 # DiaObject tables to update.
815 object_tables = [self._schema.get_table(ApdbTables.DiaObject)]
816 if self.config.dia_object_index == "last_object_table":
817 object_tables.append(self._schema.get_table(ApdbTables.DiaObjectLast))
819 # Things to increment/decrement.
820 increments: Counter = Counter()
821 if increment_nDiaSources:
822 increments.update(idMap.values())
823 if decrement_nDiaSources:
824 increments.subtract(original_object_ids[source_id.diaSourceId] for source_id in idMap)
826 if increments:
827 for table in object_tables:
828 for diaObjectId, increment in increments.items():
829 update = (
830 table.update()
831 .where(table.columns["diaObjectId"] == diaObjectId)
832 .values(nDiaSources=table.columns["nDiaSources"] + increment)
833 )
834 conn.execute(update)
836 # Also send updated values to replica.
837 if self._schema.replication_enabled:
838 for diaObjectId, increment in increments.items():
839 dia_object = found_objects_by_id[diaObjectId]
840 update_records.append(
841 ApdbUpdateNDiaSourcesRecord(
842 diaObjectId=diaObjectId,
843 ra=dia_object.ra,
844 dec=dia_object.dec,
845 nDiaSources=dia_object.nDiaSources + increment,
846 update_time_ns=current_time_ns,
847 update_order=update_order,
848 )
849 )
850 update_order += 1
852 if update_records:
853 replica_chunk = ReplicaChunk.make_replica_chunk(
854 current_time, self.config.replica_chunk_seconds
855 )
856 self._storeUpdateRecords(update_records, replica_chunk, connection=conn, store_chunk=True)
858 def setValidityEnd(
859 self, objects: list[DiaObjectId], validityEnd: astropy.time.Time, raise_on_missing_id: bool = False
860 ) -> int:
861 # docstring is inherited from a base class
862 if not objects:
863 return 0
865 requested_ids = {obj.diaObjectId for obj in objects}
867 validity_end_column = self._timestamp_column_name("validityEnd")
868 validityEnd_value = self._timestamp_column_value(validityEnd)
870 # Find all matching DiaObjects with validityEnd = NULL.
871 table = self._schema.get_table(ApdbTables.DiaObject)
872 query = sql.select(table.columns["diaObjectId"]).where(
873 sqlalchemy.and_(
874 table.columns["diaObjectId"].in_(sorted(requested_ids)),
875 table.columns[validity_end_column].is_(None),
876 )
877 )
879 with self._engine.begin() as conn:
880 result = conn.execute(query)
881 found_ids = set(result.scalars())
883 # Check that we found all that is requested.
884 if raise_on_missing_id:
885 if missing_ids := (requested_ids - found_ids):
886 raise LookupError(f"Some object IDs are missing from DiaObjectLast table: {missing_ids}")
888 # Filter existing records.
889 if len(objects) != len(found_ids):
890 objects = [obj for obj in objects if obj.diaObjectId in found_ids]
892 if not objects:
893 return 0
895 values = {validity_end_column: validityEnd_value}
896 update = (
897 table.update()
898 .where(
899 sqlalchemy.and_(
900 table.columns["diaObjectId"].in_(sorted(found_ids)),
901 table.columns[validity_end_column].is_(None),
902 )
903 )
904 .values(**values)
905 )
906 result = conn.execute(update)
907 if result.rowcount != len(found_ids):
908 raise RuntimeError(
909 f"Unexpected mismatch in the number of records updated. Object IDs = {found_ids}"
910 )
912 # Also drop them from DiaObjectLast.
913 if self.config.dia_object_index == "last_object_table":
914 last_table = self._schema.get_table(ApdbTables.DiaObjectLast)
915 delete = last_table.delete().where(last_table.columns["diaObjectId"].in_(sorted(found_ids)))
916 result = conn.execute(delete)
917 if result.rowcount != len(found_ids):
918 raise RuntimeError(
919 f"Unexpected mismatch in the number of records deleted. Object IDs = {found_ids}"
920 )
922 # If replication is enabled then send all updates.
923 if self._schema.replication_enabled:
924 current_time = self._current_time()
925 current_time_ns = int(current_time.unix_tai * 1e9)
926 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, self.config.replica_chunk_seconds)
928 update_records = [
929 ApdbCloseDiaObjectValidityRecord(
930 diaObjectId=obj.diaObjectId,
931 ra=obj.ra,
932 dec=obj.dec,
933 update_time_ns=current_time_ns,
934 update_order=index,
935 validityEndMjdTai=float(validityEnd.tai.mjd),
936 nDiaSources=None,
937 )
938 for index, obj in enumerate(objects)
939 ]
941 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True)
943 return len(objects)
945 def resetDedup(self, dedup_time: astropy.time.Time | None = None) -> None:
946 # docstring is inherited from a base class
948 # SQL backend does not have separate dedup tables, nothing to delete,
949 # only save last dedup time in metadata.
950 if dedup_time is None:
951 dedup_time = self._current_time()
952 data = {"dedup_time_iso_tai": dedup_time.tai.to_value("iso")}
953 data_json = json.dumps(data)
954 self._metadata.set(self.metadataDedupKey, data_json, force=True)
956 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
957 # docstring is inherited from a base class
959 timestamp: float | datetime.datetime
960 now = self._current_time()
961 timestamp_column = self._timestamp_column_name("ssObjectReassocTime")
962 timestamp = self._timestamp_column_value(now)
964 table = self._schema.get_table(ApdbTables.DiaSource)
965 query = table.update().where(table.columns["diaSourceId"] == sql.bindparam("srcId"))
967 with self._engine.begin() as conn:
968 # Need to make sure that every ID exists in the database, but
969 # executemany may not support rowcount, so iterate and check what
970 # is missing.
971 missing_ids: list[int] = []
972 for key, value in idMap.items():
973 params = {
974 "srcId": key,
975 "diaObjectId": 0,
976 "ssObjectId": value,
977 timestamp_column: timestamp,
978 }
979 result = conn.execute(query, params)
980 if result.rowcount == 0:
981 missing_ids.append(key)
982 if missing_ids:
983 missing = ",".join(str(item) for item in missing_ids)
984 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
986 def countUnassociatedObjects(self) -> int:
987 # docstring is inherited from a base class
989 # Retrieve the DiaObject table.
990 table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaObject)
992 # Construct the sql statement.
993 validity_end_column = self._timestamp_column_name("validityEnd")
994 stmt = (
995 sql.select(func.count())
996 .select_from(table)
997 .where(
998 sqlalchemy.and_(
999 table.columns["nDiaSources"] == 1,
1000 table.columns[validity_end_column].is_(None),
1001 )
1002 )
1003 )
1005 # Return the count.
1006 with self._engine.begin() as conn:
1007 count = conn.execute(stmt).scalar_one()
1009 return count
1011 @property
1012 def schema(self) -> ApdbSchema:
1013 # docstring is inherited from a base class
1014 return self._table_schema
1016 @property
1017 def metadata(self) -> ApdbMetadata:
1018 # docstring is inherited from a base class
1019 return self._metadata
1021 @property
1022 def admin(self) -> ApdbSqlAdmin:
1023 # docstring is inherited from a base class
1024 return ApdbSqlAdmin(self.pixelator)
1026 def _getDiaSourcesInRegion(self, region: Region, start_time_mjdTai: float) -> pandas.DataFrame:
1027 """Return catalog of DiaSource instances from given region.
1029 Parameters
1030 ----------
1031 region : `lsst.sphgeom.Region`
1032 Region to search for DIASources.
1033 start_time_mjdTai : `float`
1034 Lower bound of time window for the query.
1036 Returns
1037 -------
1038 catalog : `pandas.DataFrame`
1039 Catalog containing DiaSource records.
1040 """
1041 table = self._schema.get_table(ApdbTables.DiaSource)
1042 columns = self._schema.get_apdb_columns(ApdbTables.DiaSource)
1043 query = sql.select(*columns)
1045 # build selection
1046 time_filter = table.columns["midpointMjdTai"] > start_time_mjdTai
1047 where = sql.expression.and_(self._filterRegion(table, region), time_filter)
1048 query = query.where(where)
1050 # execute select
1051 with self._timer("DiaSource_select_time", tags={"table": "DiaSource"}) as timer:
1052 with self._engine.begin() as conn:
1053 sources = pandas.read_sql_query(query, conn)
1054 timer.add_values(row_counts=len(sources))
1055 _LOG.debug("found %s DiaSources", len(sources))
1056 return self._fix_result_timestamps(sources)
1058 def _getDiaSourcesByIDs(self, object_ids: list[int], start_time_mjdTai: float) -> pandas.DataFrame:
1059 """Return catalog of DiaSource instances given set of DiaObject IDs.
1061 Parameters
1062 ----------
1063 object_ids :
1064 Collection of DiaObject IDs
1065 start_time_mjdTai : `float`
1066 Lower bound of time window for the query.
1068 Returns
1069 -------
1070 catalog : `pandas.DataFrame`
1071 Catalog containing DiaSource records.
1072 """
1073 with self._timer("select_time", tags={"table": "DiaSource"}) as timer:
1074 sources = self._getSourcesByIDs(ApdbTables.DiaSource, object_ids, start_time_mjdTai)
1075 timer.add_values(row_count=len(sources))
1077 _LOG.debug("found %s DiaSources", len(sources))
1078 return sources
1080 def _getSourcesByIDs(
1081 self, table_enum: ApdbTables, object_ids: list[int], midpointMjdTai_start: float
1082 ) -> pandas.DataFrame:
1083 """Return catalog of DiaSource or DiaForcedSource instances given set
1084 of DiaObject IDs.
1086 Parameters
1087 ----------
1088 table : `sqlalchemy.schema.Table`
1089 Database table.
1090 object_ids :
1091 Collection of DiaObject IDs
1092 midpointMjdTai_start : `float`
1093 Earliest midpointMjdTai to retrieve.
1095 Returns
1096 -------
1097 catalog : `pandas.DataFrame`
1098 Catalog contaning DiaSource records. `None` is returned if
1099 ``read_sources_months`` configuration parameter is set to 0 or
1100 when ``object_ids`` is empty.
1101 """
1102 table = self._schema.get_table(table_enum)
1103 columns = self._schema.get_apdb_columns(table_enum)
1105 sources: pandas.DataFrame | None = None
1106 if len(object_ids) <= 0:
1107 _LOG.debug("ID list is empty, just fetch empty result")
1108 query = sql.select(*columns).where(sql.literal(False))
1109 with self._engine.begin() as conn:
1110 sources = pandas.read_sql_query(query, conn)
1111 else:
1112 data_frames: list[pandas.DataFrame] = []
1113 for ids in chunk_iterable(sorted(object_ids), 1000):
1114 query = sql.select(*columns)
1116 # Some types like np.int64 can cause issues with
1117 # sqlalchemy, convert them to int.
1118 int_ids = [int(oid) for oid in ids]
1120 # select by object id
1121 query = query.where(
1122 sql.expression.and_(
1123 table.columns["diaObjectId"].in_(int_ids),
1124 table.columns["midpointMjdTai"] >= midpointMjdTai_start,
1125 )
1126 )
1128 # execute select
1129 with self._engine.begin() as conn:
1130 data_frames.append(pandas.read_sql_query(query, conn))
1132 if len(data_frames) == 1:
1133 sources = data_frames[0]
1134 else:
1135 sources = pandas.concat(data_frames)
1136 assert sources is not None, "Catalog cannot be None"
1137 return self._fix_result_timestamps(sources)
1139 def _storeReplicaChunk(
1140 self,
1141 replica_chunk: ReplicaChunk,
1142 connection: sqlalchemy.engine.Connection,
1143 ) -> None:
1144 # `visit_time.datetime` returns naive datetime, even though all astropy
1145 # times are in UTC. Add UTC timezone to timestamp so that database
1146 # can store a correct value.
1147 dt = datetime.datetime.fromtimestamp(replica_chunk.last_update_time.unix_tai, tz=datetime.UTC)
1149 table = self._schema.get_table(ExtraTables.ApdbReplicaChunks)
1151 # We need UPSERT which is dialect-specific construct
1152 values = {"last_update_time": dt, "unique_id": replica_chunk.unique_id}
1153 row = {"apdb_replica_chunk": replica_chunk.id} | values
1154 if connection.dialect.name == "sqlite":
1155 insert_sqlite = sqlalchemy.dialects.sqlite.insert(table)
1156 insert_sqlite = insert_sqlite.on_conflict_do_update(index_elements=table.primary_key, set_=values)
1157 connection.execute(insert_sqlite, row)
1158 elif connection.dialect.name == "postgresql":
1159 insert_pg = sqlalchemy.dialects.postgresql.dml.insert(table)
1160 insert_pg = insert_pg.on_conflict_do_update(constraint=table.primary_key, set_=values)
1161 connection.execute(insert_pg, row)
1162 else:
1163 raise TypeError(f"Unsupported dialect {connection.dialect.name} for upsert.")
1165 def _storeDiaObjects(
1166 self,
1167 objs: pandas.DataFrame,
1168 visit_time: astropy.time.Time,
1169 replica_chunk: ReplicaChunk | None,
1170 connection: sqlalchemy.engine.Connection,
1171 ) -> None:
1172 """Store catalog of DiaObjects from current visit.
1174 Parameters
1175 ----------
1176 objs : `pandas.DataFrame`
1177 Catalog with DiaObject records.
1178 visit_time : `astropy.time.Time`
1179 Time of the visit.
1180 replica_chunk : `ReplicaChunk`
1181 Insert identifier.
1182 """
1183 if len(objs) == 0:
1184 _LOG.debug("No objects to write to database.")
1185 return
1187 # Some types like np.int64 can cause issues with sqlalchemy, convert
1188 # them to int.
1189 ids = sorted(int(oid) for oid in objs["diaObjectId"])
1190 _LOG.debug("first object ID: %d", ids[0])
1192 validity_start_column = self._timestamp_column_name("validityStart")
1193 validity_end_column = self._timestamp_column_name("validityEnd")
1194 timestamp = self._timestamp_column_value(visit_time)
1196 # everything to be done in single transaction
1197 if self.config.dia_object_index == "last_object_table":
1198 # Insert and replace all records in LAST table.
1199 table = self._schema.get_table(ApdbTables.DiaObjectLast)
1201 # DiaObjectLast did not have this column in the past.
1202 use_validity_start = self._schema.check_column(ApdbTables.DiaObjectLast, validity_start_column)
1204 # Drop the previous objects (pandas cannot upsert).
1205 query = table.delete().where(table.columns["diaObjectId"].in_(ids))
1207 with self._timer("delete_time", tags={"table": table.name}) as timer:
1208 res = connection.execute(query)
1209 timer.add_values(row_count=res.rowcount)
1210 _LOG.debug("deleted %s objects", res.rowcount)
1212 # DiaObjectLast is a subset of DiaObject, strip missing columns
1213 last_column_names = [column.name for column in table.columns]
1214 if validity_start_column in last_column_names and validity_start_column not in objs.columns:
1215 last_column_names.remove(validity_start_column)
1216 last_objs = objs[last_column_names]
1217 last_objs = _coerce_uint64(last_objs)
1219 # Fill validityStart, only when it is in the schema.
1220 if use_validity_start:
1221 if validity_start_column in last_objs:
1222 last_objs[validity_start_column] = timestamp
1223 else:
1224 extra_column = pandas.Series([timestamp] * len(last_objs), name=validity_start_column)
1225 last_objs.set_index(extra_column.index, inplace=True)
1226 last_objs = pandas.concat([last_objs, extra_column], axis="columns")
1228 with self._timer("insert_time", tags={"table": "DiaObjectLast"}) as timer:
1229 last_objs.to_sql(
1230 table.name,
1231 connection,
1232 if_exists="append",
1233 index=False,
1234 schema=table.schema,
1235 )
1236 timer.add_values(row_count=len(last_objs))
1237 else:
1238 # truncate existing validity intervals
1239 table = self._schema.get_table(ApdbTables.DiaObject)
1241 update = (
1242 table.update()
1243 .values(**{validity_end_column: timestamp})
1244 .where(
1245 sql.expression.and_(
1246 table.columns["diaObjectId"].in_(ids),
1247 table.columns[validity_end_column].is_(None),
1248 )
1249 )
1250 )
1252 with self._timer("truncate_time", tags={"table": table.name}) as timer:
1253 res = connection.execute(update)
1254 timer.add_values(row_count=res.rowcount)
1255 _LOG.debug("truncated %s intervals", res.rowcount)
1257 objs = _coerce_uint64(objs)
1259 # Fill additional columns
1260 extra_columns: list[pandas.Series] = []
1261 if validity_start_column in objs.columns:
1262 objs[validity_start_column] = timestamp
1263 else:
1264 extra_columns.append(pandas.Series([timestamp] * len(objs), name=validity_start_column))
1265 if validity_end_column in objs.columns:
1266 objs[validity_end_column] = None
1267 else:
1268 extra_columns.append(pandas.Series([None] * len(objs), name=validity_end_column))
1269 if extra_columns:
1270 objs.set_index(extra_columns[0].index, inplace=True)
1271 objs = pandas.concat([objs] + extra_columns, axis="columns")
1273 # Insert replica data
1274 table = self._schema.get_table(ApdbTables.DiaObject)
1275 replica_data: list[dict] = []
1276 replica_stmt: Any = None
1277 replica_table_name = ""
1278 if replica_chunk is not None:
1279 pk_names = [column.name for column in table.primary_key]
1280 replica_data = objs[pk_names].to_dict("records")
1281 if replica_data:
1282 for row in replica_data:
1283 row["apdb_replica_chunk"] = replica_chunk.id
1284 replica_table = self._schema.get_table(ExtraTables.DiaObjectChunks)
1285 replica_table_name = replica_table.name
1286 replica_stmt = replica_table.insert()
1288 # insert new versions
1289 with self._timer("insert_time", tags={"table": table.name}) as timer:
1290 objs.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1291 timer.add_values(row_count=len(objs))
1292 if replica_stmt is not None:
1293 with self._timer("insert_time", tags={"table": replica_table_name}) as timer:
1294 connection.execute(replica_stmt, replica_data)
1295 timer.add_values(row_count=len(replica_data))
1297 def _storeDiaSources(
1298 self,
1299 sources: pandas.DataFrame,
1300 replica_chunk: ReplicaChunk | None,
1301 connection: sqlalchemy.engine.Connection,
1302 ) -> None:
1303 """Store catalog of DiaSources from current visit.
1305 Parameters
1306 ----------
1307 sources : `pandas.DataFrame`
1308 Catalog containing DiaSource records
1309 """
1310 table = self._schema.get_table(ApdbTables.DiaSource)
1312 # Insert replica data
1313 replica_data: list[dict] = []
1314 replica_stmt: Any = None
1315 replica_table_name = ""
1316 if replica_chunk is not None:
1317 pk_names = [column.name for column in table.primary_key]
1318 replica_data = sources[pk_names].to_dict("records")
1319 if replica_data:
1320 for row in replica_data:
1321 row["apdb_replica_chunk"] = replica_chunk.id
1322 replica_table = self._schema.get_table(ExtraTables.DiaSourceChunks)
1323 replica_table_name = replica_table.name
1324 replica_stmt = replica_table.insert()
1326 # everything to be done in single transaction
1327 with self._timer("insert_time", tags={"table": table.name}) as timer:
1328 sources = _coerce_uint64(sources)
1329 sources.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1330 timer.add_values(row_count=len(sources))
1331 if replica_stmt is not None:
1332 with self._timer("replica_insert_time", tags={"table": replica_table_name}) as timer:
1333 connection.execute(replica_stmt, replica_data)
1334 timer.add_values(row_count=len(replica_data))
1336 def _storeDiaForcedSources(
1337 self,
1338 sources: pandas.DataFrame,
1339 replica_chunk: ReplicaChunk | None,
1340 connection: sqlalchemy.engine.Connection,
1341 ) -> None:
1342 """Store a set of DiaForcedSources from current visit.
1344 Parameters
1345 ----------
1346 sources : `pandas.DataFrame`
1347 Catalog containing DiaForcedSource records
1348 """
1349 table = self._schema.get_table(ApdbTables.DiaForcedSource)
1351 # Insert replica data
1352 replica_data: list[dict] = []
1353 replica_stmt: Any = None
1354 replica_table_name = ""
1355 if replica_chunk is not None:
1356 pk_names = [column.name for column in table.primary_key]
1357 replica_data = sources[pk_names].to_dict("records")
1358 if replica_data:
1359 for row in replica_data:
1360 row["apdb_replica_chunk"] = replica_chunk.id
1361 replica_table = self._schema.get_table(ExtraTables.DiaForcedSourceChunks)
1362 replica_table_name = replica_table.name
1363 replica_stmt = replica_table.insert()
1365 # everything to be done in single transaction
1366 with self._timer("insert_time", tags={"table": table.name}) as timer:
1367 sources = _coerce_uint64(sources)
1368 sources.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1369 timer.add_values(row_count=len(sources))
1370 if replica_stmt is not None:
1371 with self._timer("insert_time", tags={"table": replica_table_name}) as timer:
1372 connection.execute(replica_stmt, replica_data)
1373 timer.add_values(row_count=len(replica_data))
1375 def _storeUpdateRecords(
1376 self,
1377 records: Iterable[ApdbUpdateRecord],
1378 chunk: ReplicaChunk,
1379 *,
1380 store_chunk: bool = False,
1381 connection: sqlalchemy.engine.Connection | None = None,
1382 ) -> None:
1383 """Store ApdbUpdateRecords in the replica table for those records.
1385 Parameters
1386 ----------
1387 records : `list` [`ApdbUpdateRecord`]
1388 Records to store.
1389 chunk : `ReplicaChunk`
1390 Replica chunk for these records.
1391 store_chunk : `bool`
1392 If True then also store replica chunk.
1393 connection : `sqlalchemy.engine.Connection`
1394 SQLALchemy connection to use, if `None` the new connection will be
1395 made. `None` is useful for tests only, regular use will call this
1396 method in the same transaction that saves other types of records.
1398 Raises
1399 ------
1400 TypeError
1401 Raised if replication is not enabled for this instance.
1402 """
1403 if not self._schema.replication_enabled:
1404 raise TypeError("Replication is not enabled for this APDB instance.")
1406 apdb_replica_chunk = chunk.id
1407 # Do not use unique_if from ReplicaChunk as it could be reused in
1408 # multiple calls to this method.
1409 update_unique_id = uuid.uuid4()
1411 record_dicts = []
1412 for record in records:
1413 record_dicts.append(
1414 {
1415 "apdb_replica_chunk": apdb_replica_chunk,
1416 "update_time_ns": record.update_time_ns,
1417 "update_order": record.update_order,
1418 "update_unique_id": update_unique_id,
1419 "update_payload": record.to_json(),
1420 }
1421 )
1423 if not record_dicts:
1424 return
1426 # TODO: Need to check that table exists.
1427 table = self._schema.get_table(ExtraTables.ApdbUpdateRecordChunks)
1429 def _do_store(connection: sqlalchemy.engine.Connection) -> None:
1430 if store_chunk:
1431 self._storeReplicaChunk(chunk, connection)
1432 with self._timer("insert_time", tags={"table": table.name}) as timer:
1433 connection.execute(table.insert(), record_dicts)
1434 timer.add_values(row_count=len(record_dicts))
1436 if connection is None:
1437 with self._engine.begin() as connection:
1438 _do_store(connection)
1439 else:
1440 _do_store(connection)
1442 def _htm_indices(self, region: Region) -> list[tuple[int, int]]:
1443 """Generate a set of HTM indices covering specified region.
1445 Parameters
1446 ----------
1447 region: `sphgeom.Region`
1448 Region that needs to be indexed.
1450 Returns
1451 -------
1452 Sequence of ranges, range is a tuple (minHtmID, maxHtmID).
1453 """
1454 _LOG.debug("region: %s", region)
1455 indices = self.pixelator.envelope(region, self.config.pixelization.htm_max_ranges)
1457 return indices.ranges()
1459 def _filterRegion(self, table: sqlalchemy.schema.Table, region: Region) -> sql.ColumnElement:
1460 """Make SQLAlchemy expression for selecting records in a region."""
1461 htm_index_column = table.columns[self.config.pixelization.htm_index_column]
1462 exprlist = []
1463 pixel_ranges = self._htm_indices(region)
1464 for low, upper in pixel_ranges:
1465 upper -= 1
1466 if low == upper:
1467 exprlist.append(htm_index_column == low)
1468 else:
1469 exprlist.append(sql.expression.between(htm_index_column, low, upper))
1471 return sql.expression.or_(*exprlist)
1473 def _add_spatial_index(self, df: pandas.DataFrame) -> pandas.DataFrame:
1474 """Calculate spatial index for each record and add it to a DataFrame.
1476 Parameters
1477 ----------
1478 df : `pandas.DataFrame`
1479 DataFrame which has to contain ra/dec columns, names of these
1480 columns are defined by configuration ``ra_dec_columns`` field.
1482 Returns
1483 -------
1484 df : `pandas.DataFrame`
1485 DataFrame with ``pixelId`` column which contains pixel index
1486 for ra/dec coordinates.
1488 Notes
1489 -----
1490 This overrides any existing column in a DataFrame with the same name
1491 (pixelId). Original DataFrame is not changed, copy of a DataFrame is
1492 returned.
1493 """
1494 # calculate HTM index for every DiaObject
1495 htm_index = np.zeros(df.shape[0], dtype=np.int64)
1496 ra_col, dec_col = self.config.ra_dec_columns
1497 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1498 uv3d = UnitVector3d(LonLat.fromDegrees(ra, dec))
1499 idx = self.pixelator.index(uv3d)
1500 htm_index[i] = idx
1501 df = df.copy()
1502 df[self.config.pixelization.htm_index_column] = htm_index
1503 return df
1505 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1506 """Update timestamp columns in input DataFrame to be aware datetime
1507 type in in UTC.
1509 AP pipeline generates naive datetime instances, we want them to be
1510 aware before they go to database. All naive timestamps are assumed to
1511 be in UTC timezone (they should be TAI).
1512 """
1513 # Find all columns with aware non-UTC timestamps and convert to UTC.
1514 columns = [
1515 column
1516 for column, dtype in df.dtypes.items()
1517 if isinstance(dtype, pandas.DatetimeTZDtype) and dtype.tz is not datetime.UTC
1518 ]
1519 for column in columns:
1520 df[column] = df[column].dt.tz_convert(datetime.UTC)
1521 # Find all columns with naive timestamps and add UTC timezone.
1522 columns = [
1523 column for column, dtype in df.dtypes.items() if pandas.api.types.is_datetime64_dtype(dtype)
1524 ]
1525 for column in columns:
1526 df[column] = df[column].dt.tz_localize(datetime.UTC)
1527 return df
1529 def _fix_result_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1530 """Update timestamp columns to be naive datetime type in returned
1531 DataFrame.
1533 AP pipeline code expects DataFrames to contain naive datetime columns,
1534 while Postgres queries return timezone-aware type. This method converts
1535 those columns to naive datetime in UTC timezone.
1536 """
1537 # Find all columns with aware timestamps.
1538 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1539 for column in columns:
1540 # tz_convert(None) will convert to UTC and drop timezone.
1541 df[column] = df[column].dt.tz_convert(None)
1542 return df
1544 def _timestamp_column_name(self, column: str) -> str:
1545 """Return column name before/after schema migration to MJD TAI."""
1546 return self.schema.timestamp_column_name(column)
1548 def _timestamp_column_value(self, time: astropy.time.Time) -> float | datetime.datetime:
1549 """Return column value before/after schema migration to MJD TAI."""
1550 if self.schema.has_mjd_timestamps:
1551 return float(time.tai.mjd)
1552 else:
1553 return time.datetime.astimezone(tz=datetime.UTC)
1555 def _get_diaobject_data(
1556 self, conn: sqlalchemy.engine.Connection, object_ids: Iterable[int], *columns: str
1557 ) -> list:
1558 """Select records from either DiaObject or DiaObjectLast and return
1559 selected rows as names tuples.
1560 """
1561 where: sqlalchemy.ColumnElement[bool]
1562 if self.config.dia_object_index == "last_object_table":
1563 table = self._schema.get_table(ApdbTables.DiaObjectLast)
1564 where = table.columns["diaObjectId"].in_(sorted(object_ids))
1565 else:
1566 table = self._schema.get_table(ApdbTables.DiaObject)
1567 validity_end_column = self._timestamp_column_name("validityEnd")
1568 where = sqlalchemy.and_(
1569 table.columns["diaObjectId"].in_(sorted(object_ids)),
1570 table.columns[validity_end_column].is_(None),
1571 )
1572 column_list = [table.columns["diaObjectId"]] + [table.columns[column] for column in columns]
1573 query = sql.select(*column_list).where(where)
1574 result = conn.execute(query)
1576 return list(result)
1578 def _get_diasource_data(
1579 self, conn: sqlalchemy.engine.Connection, source_ids: Iterable[int], *columns: str
1580 ) -> list:
1581 """Select records from DiaSource table by diaSourceId and return
1582 selected rows as named tuples.
1583 """
1584 table = self._schema.get_table(ApdbTables.DiaSource)
1585 where = table.columns["diaSourceId"].in_(sorted(source_ids))
1586 column_list = [table.columns["diaSourceId"]] + [table.columns[column] for column in columns]
1587 query = sql.select(*column_list).where(where)
1588 result = conn.execute(query)
1590 return list(result)