Coverage for python/lsst/dax/apdb/sql/apdbSql.py: 88%
759 statements
« prev ^ index » next coverage.py v7.17.2, created at 2026-11-09 17:49 +0000
1# This file is part of dax_apdb.
2#
3# Developed for the LSST Data Management System.
4# This product includes software developed by the LSST Project
5# (http://www.lsst.org).
6# See the COPYRIGHT file at the top-level directory of this distribution
7# for details of code ownership.
8#
9# This program is free software: you can redistribute it and/or modify
10# it under the terms of the GNU General Public License as published by
11# the Free Software Foundation, either version 3 of the License, or
12# (at your option) any later version.
13#
14# This program is distributed in the hope that it will be useful,
15# but WITHOUT ANY WARRANTY; without even the implied warranty of
16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
17# GNU General Public License for more details.
18#
19# You should have received a copy of the GNU General Public License
20# along with this program. If not, see <http://www.gnu.org/licenses/>.
22"""Module defining Apdb class and related methods."""
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 ApdbWithdrawDiaForcedSourceRecord,
62 ApdbWithdrawDiaSourceRecord,
63)
64from ..config import ApdbConfig
65from ..monitor import MonAgent
66from ..recordIds import DiaForcedSourceId, DiaObjectId, DiaSourceId
67from ..schema_model import Table
68from ..timer import Timer
69from ..versionTuple import IncompatibleVersionError, VersionTuple
70from .apdbMetadataSql import ApdbMetadataSql
71from .apdbSqlAdmin import ApdbSqlAdmin
72from .apdbSqlReplica import ApdbSqlReplica, ApdbSqlTableData
73from .apdbSqlSchema import ApdbSqlSchema, ExtraTables
74from .config import ApdbSqlConfig
76if TYPE_CHECKING:
77 import sqlite3
79 from ..apdbMetadata import ApdbMetadata
80 from ..apdbUpdateRecord import ApdbUpdateRecord
82_LOG = logging.getLogger(__name__)
84_MON = MonAgent(__name__)
86VERSION = VersionTuple(1, 2, 1)
87"""Version for the code controlling non-replication tables. This needs to be
88updated following compatibility rules when schema produced by this code
89changes.
90"""
93def _coerce_uint64(df: pandas.DataFrame) -> pandas.DataFrame:
94 """Change the type of uint64 columns to int64, and return copy of data
95 frame.
96 """
97 names = [c[0] for c in df.dtypes.items() if c[1] == np.uint64]
98 return df.astype(dict.fromkeys(names, np.int64))
101def _make_midpointMjdTai_start(visit_time: astropy.time.Time, months: int) -> float:
102 """Calculate starting point for time-based source search.
104 Parameters
105 ----------
106 visit_time : `astropy.time.Time`
107 Time of current visit.
108 months : `int`
109 Number of months in the sources history.
111 Returns
112 -------
113 time : `float`
114 A ``midpointMjdTai`` starting point, MJD time.
115 """
116 # TODO: Use of MJD must be consistent with the code in ap_association
117 # (see DM-31996)
118 return float(visit_time.tai.mjd) - months * 30
121def _onSqlite3Connect(
122 dbapiConnection: sqlite3.Connection, connectionRecord: sqlalchemy.pool._ConnectionRecord
123) -> None:
124 # Enable foreign keys
125 with closing(dbapiConnection.cursor()) as cursor:
126 cursor.execute("PRAGMA foreign_keys=ON;")
129class ApdbSql(Apdb):
130 """Implementation of APDB interface based on SQL database.
132 The implementation is configured via standard ``pex_config`` mechanism
133 using `ApdbSqlConfig` configuration class. For an example of different
134 configurations check ``config/`` folder.
136 Parameters
137 ----------
138 config : `ApdbSqlConfig`
139 Configuration object.
140 """
142 metadataSchemaVersionKey = "version:schema"
143 """Name of the metadata key to store schema version number."""
145 metadataCodeVersionKey = "version:ApdbSql"
146 """Name of the metadata key to store code version number."""
148 metadataReplicaVersionKey = "version:ApdbSqlReplica"
149 """Name of the metadata key to store replica code version number."""
151 metadataConfigKey = "config:apdb-sql.json"
152 """Name of the metadata key to store code version number."""
154 metadataDedupKey = "status:deduplication.json"
155 """Name of the metadata key to store code version number."""
157 _frozen_parameters = (
158 "enable_replica",
159 "dia_object_index",
160 "pixelization.htm_level",
161 "pixelization.htm_index_column",
162 "ra_dec_columns",
163 )
164 """Names of the config parameters to be frozen in metadata table."""
166 def __init__(self, config: ApdbSqlConfig):
167 self._engine = self._makeEngine(config, create=False)
169 sa_metadata = sqlalchemy.MetaData(schema=config.namespace)
170 meta_table_name = ApdbTables.metadata.table_name(prefix=config.prefix)
171 meta_table = sqlalchemy.schema.Table(meta_table_name, sa_metadata, autoload_with=self._engine)
172 self._metadata = ApdbMetadataSql(self._engine, meta_table)
174 # Get tables schemas.
175 self._table_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
177 # Check that versions are compatible, must be the first thing before
178 # reading frozen config.
179 self._db_schema_version = self._versionCheck(self._metadata, self._table_schema.schemaVersion())
181 # Read frozen config from metadata.
182 config_json = self._metadata.get(self.metadataConfigKey)
183 if config_json is not None: 183 ↛ 188line 183 didn't jump to line 188 because the condition on line 183 was always true
184 # Update config from metadata.
185 freezer = ApdbConfigFreezer[ApdbSqlConfig](self._frozen_parameters)
186 self.config = freezer.update(config, config_json)
187 else:
188 self.config = config
190 self._schema = ApdbSqlSchema(
191 table_schema=self._table_schema,
192 engine=self._engine,
193 dia_object_index=self.config.dia_object_index,
194 prefix=self.config.prefix,
195 namespace=self.config.namespace,
196 htm_index_column=self.config.pixelization.htm_index_column,
197 enable_replica=self.config.enable_replica,
198 )
200 self.pixelator = HtmPixelization(self.config.pixelization.htm_level)
202 if _LOG.isEnabledFor(logging.DEBUG): 202 ↛ exitline 202 didn't return from function '__init__' because the condition on line 202 was always true
203 _LOG.debug("ApdbSql Configuration: %s", self.config.model_dump())
205 def close(self) -> None:
206 # docstring is inherited from a base class
207 self._engine.dispose()
209 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer:
210 """Create `Timer` instance given its name."""
211 return Timer(name, _MON, tags=tags)
213 @classmethod
214 def _makeEngine(cls, config: ApdbSqlConfig, *, create: bool) -> sqlalchemy.engine.Engine:
215 """Make SQLALchemy engine based on configured parameters.
217 Parameters
218 ----------
219 config : `ApdbSqlConfig`
220 Configuration object.
221 create : `bool`
222 Whether to try to create new database file, only relevant for
223 SQLite backend which always creates new files by default.
224 """
225 # engine is reused between multiple processes, make sure that we don't
226 # share connections by disabling pool (by using NullPool class)
227 kw: MutableMapping[str, Any] = dict(config.connection_config.extra_parameters)
228 conn_args: dict[str, Any] = {}
229 if not config.connection_config.connection_pool: 229 ↛ 230line 229 didn't jump to line 230 because the condition on line 229 was never true
230 kw.update(poolclass=NullPool)
231 if config.connection_config.isolation_level is not None: 231 ↛ 232line 231 didn't jump to line 232 because the condition on line 231 was never true
232 kw.update(isolation_level=config.connection_config.isolation_level)
233 elif config.db_url.startswith("sqlite"): 233 ↛ 236line 233 didn't jump to line 236 because the condition on line 233 was always true
234 # Use READ_UNCOMMITTED as default value for sqlite.
235 kw.update(isolation_level="READ_UNCOMMITTED")
236 if config.connection_config.connection_timeout is not None:
237 if config.db_url.startswith("sqlite"): 237 ↛ 239line 237 didn't jump to line 239 because the condition on line 237 was always true
238 conn_args.update(timeout=config.connection_config.connection_timeout)
239 elif config.db_url.startswith(("postgresql", "mysql")):
240 conn_args.update(connect_timeout=int(config.connection_config.connection_timeout))
241 kw.update(connect_args=conn_args)
242 engine = sqlalchemy.create_engine(cls._connection_url(config.db_url, create=create), **kw)
244 if engine.dialect.name == "sqlite": 244 ↛ 248line 244 didn't jump to line 248 because the condition on line 244 was always true
245 # Need to enable foreign keys on every new connection.
246 sqlalchemy.event.listen(engine, "connect", _onSqlite3Connect)
248 return engine
250 @classmethod
251 def _connection_url(cls, config_url: str, *, create: bool) -> sqlalchemy.engine.URL | str:
252 """Generate a complete URL for database with proper credentials.
254 Parameters
255 ----------
256 config_url : `str`
257 Database URL as specified in configuration.
258 create : `bool`
259 Whether to try to create new database file, only relevant for
260 SQLite backend which always creates new files by default.
262 Returns
263 -------
264 connection_url : `sqlalchemy.engine.URL` or `str`
265 Connection URL including credentials.
266 """
267 # Allow 3rd party authentication mechanisms by assuming connection
268 # string is correct when we can not recognize (dialect, host, database)
269 # matching keys.
270 components = urllib.parse.urlparse(config_url)
271 if all((components.scheme is not None, components.hostname is not None, components.path is not None)): 271 ↛ 272line 271 didn't jump to line 272 because the condition on line 271 was never true
272 try:
273 db_auth = DbAuth()
274 config_url = db_auth.getUrl(config_url)
275 except DbAuthNotFoundError:
276 # Credentials file doesn't exist or no matching credentials,
277 # use default auth.
278 pass
280 # SQLite has a nasty habit creating empty databases when they do not
281 # exist, tell it not to do that unless we do need to create it.
282 if not create:
283 config_url = cls._update_sqlite_url(config_url)
285 return config_url
287 @classmethod
288 def _update_sqlite_url(cls, url_string: str) -> str:
289 """If URL refers to sqlite dialect, update it so that the backend does
290 not try to create database file if it does not exist already.
292 Parameters
293 ----------
294 url_string : `str`
295 Connection string.
297 Returns
298 -------
299 url_string : `str`
300 Possibly updated connection string.
301 """
302 try:
303 url = sqlalchemy.make_url(url_string)
304 except sqlalchemy.exc.SQLAlchemyError:
305 # If parsing fails it means some special format, likely not
306 # sqlite so we just return it unchanged.
307 return url_string
309 if url.get_backend_name() == "sqlite": 309 ↛ 332line 309 didn't jump to line 332 because the condition on line 309 was always true
310 # Massage url so that database name starts with "file:" and
311 # option string has "mode=rw&uri=true". Database name
312 # should look like a path (:memory: is not supported by
313 # Apdb, but someone could still try to use it).
314 database = url.database
315 if database and not database.startswith((":", "file:")):
316 query = dict(url.query, mode="rw", uri="true")
317 # If ``database`` is an absolute path then original URL should
318 # include four slashes after "sqlite:". Humans are bad at
319 # counting things beyond four and sometimes an extra slash gets
320 # added unintentionally, which causes sqlite to treat initial
321 # element as "authority" and to complain. Strip extra slashes
322 # at the start of the path to avoid that (DM-46077).
323 if database.startswith("//"): 323 ↛ 324line 323 didn't jump to line 324 because the condition on line 323 was never true
324 warnings.warn(
325 f"Database URL contains extra leading slashes which will be removed: {url}",
326 stacklevel=3,
327 )
328 database = "/" + database.lstrip("/")
329 url = url.set(database=f"file:{database}", query=query)
330 url_string = url.render_as_string()
332 return url_string
334 @classmethod
335 def _versionCheck(cls, metadata: ApdbMetadataSql, schema_version: VersionTuple) -> VersionTuple:
336 """Check schema version compatibility and return the database schema
337 version.
338 """
340 def _get_version(key: str) -> VersionTuple:
341 """Retrieve version number from given metadata key."""
342 version_str = metadata.get(key)
343 if version_str is None: 343 ↛ 345line 343 didn't jump to line 345 because the condition on line 343 was never true
344 # Should not happen with existing metadata table.
345 raise RuntimeError(f"Version key {key!r} does not exist in metadata table.")
346 return VersionTuple.fromString(version_str)
348 db_schema_version = _get_version(cls.metadataSchemaVersionKey)
349 db_code_version = _get_version(cls.metadataCodeVersionKey)
351 # For now there is no way to make read-only APDB instances, assume that
352 # any access can do updates.
353 if not schema_version.checkCompatibility(db_schema_version):
354 raise IncompatibleVersionError(
355 f"Configured schema version {schema_version} "
356 f"is not compatible with database version {db_schema_version}"
357 )
358 if not cls.apdbImplementationVersion().checkCompatibility(db_code_version): 358 ↛ 359line 358 didn't jump to line 359 because the condition on line 358 was never true
359 raise IncompatibleVersionError(
360 f"Current code version {cls.apdbImplementationVersion()} "
361 f"is not compatible with database version {db_code_version}"
362 )
364 # Check replica code version only if replica is enabled. Sort of
365 # chicken and egg problem - `enable_replica` is a part of frozen
366 # configuration, but we cannot read frozen configuration until we
367 # validate versions. Assume that if the replica version is present
368 # then replication is enabled.
369 if metadata.get(cls.metadataReplicaVersionKey) is not None:
370 db_replica_version = _get_version(cls.metadataReplicaVersionKey)
371 code_replica_version = ApdbSqlReplica.apdbReplicaImplementationVersion()
372 if not code_replica_version.checkCompatibility(db_replica_version): 372 ↛ 373line 372 didn't jump to line 373 because the condition on line 372 was never true
373 raise IncompatibleVersionError(
374 f"Current replication code version {code_replica_version} "
375 f"is not compatible with database version {db_replica_version}"
376 )
378 return db_schema_version
380 @classmethod
381 def apdbImplementationVersion(cls) -> VersionTuple:
382 """Return version number for current APDB implementation.
384 Returns
385 -------
386 version : `VersionTuple`
387 Version of the code defined in implementation class.
388 """
389 return VERSION
391 @classmethod
392 def init_database(
393 cls,
394 db_url: str,
395 *,
396 schema_file: str | None = None,
397 ss_schema_file: str | None = None,
398 read_sources_months: int | None = None,
399 read_forced_sources_months: int | None = None,
400 enable_replica: bool = False,
401 connection_timeout: int | None = None,
402 dia_object_index: str | None = None,
403 htm_level: int | None = None,
404 htm_index_column: str | None = None,
405 ra_dec_columns: tuple[str, str] | None = None,
406 prefix: str | None = None,
407 namespace: str | None = None,
408 drop: bool = False,
409 ) -> ApdbSqlConfig:
410 """Initialize new APDB instance and make configuration object for it.
412 Parameters
413 ----------
414 db_url : `str`
415 SQLAlchemy database URL.
416 schema_file : `str`, optional
417 Location of (YAML) configuration file with APDB schema. If not
418 specified then default location will be used.
419 ss_schema_file : `str`, optional
420 Location of (YAML) configuration file with SSO schema. If not
421 specified then default location will be used.
422 read_sources_months : `int`, optional
423 Number of months of history to read from DiaSource.
424 read_forced_sources_months : `int`, optional
425 Number of months of history to read from DiaForcedSource.
426 enable_replica : `bool`, optional
427 If True, make additional tables used for replication to PPDB.
428 connection_timeout : `int`, optional
429 Database connection timeout in seconds.
430 dia_object_index : `str`, optional
431 Indexing mode for DiaObject table.
432 htm_level : `int`, optional
433 HTM indexing level.
434 htm_index_column : `str`, optional
435 Name of a HTM index column for DiaObject and DiaSource tables.
436 ra_dec_columns : `tuple` [`str`, `str`], optional
437 Names of ra/dec columns in DiaObject table.
438 prefix : `str`, optional
439 Optional prefix for all table names.
440 namespace : `str`, optional
441 Name of the database schema for all APDB tables. If not specified
442 then default schema is used.
443 drop : `bool`, optional
444 If `True` then drop existing tables before re-creating the schema.
446 Returns
447 -------
448 config : `ApdbSqlConfig`
449 Resulting configuration object for a created APDB instance.
450 """
451 config = ApdbSqlConfig(db_url=db_url, enable_replica=enable_replica)
452 if schema_file is not None: 452 ↛ 454line 452 didn't jump to line 454 because the condition on line 452 was always true
453 config.schema_file = schema_file
454 if ss_schema_file is not None:
455 config.ss_schema_file = ss_schema_file
456 if read_sources_months is not None:
457 config.read_sources_months = read_sources_months
458 if read_forced_sources_months is not None:
459 config.read_forced_sources_months = read_forced_sources_months
460 if connection_timeout is not None: 460 ↛ 461line 460 didn't jump to line 461 because the condition on line 460 was never true
461 config.connection_config.connection_timeout = connection_timeout
462 if dia_object_index is not None:
463 config.dia_object_index = dia_object_index
464 if htm_level is not None: 464 ↛ 465line 464 didn't jump to line 465 because the condition on line 464 was never true
465 config.pixelization.htm_level = htm_level
466 if htm_index_column is not None: 466 ↛ 467line 466 didn't jump to line 467 because the condition on line 466 was never true
467 config.pixelization.htm_index_column = htm_index_column
468 if ra_dec_columns is not None: 468 ↛ 469line 468 didn't jump to line 469 because the condition on line 468 was never true
469 config.ra_dec_columns = ra_dec_columns
470 if prefix is not None: 470 ↛ 471line 470 didn't jump to line 471 because the condition on line 470 was never true
471 config.prefix = prefix
472 if namespace is not None: 472 ↛ 473line 472 didn't jump to line 473 because the condition on line 472 was never true
473 config.namespace = namespace
475 cls._makeSchema(config, drop=drop)
477 # SQLite has a nasty habit of creating empty database by default,
478 # update URL in config file to disable that behavior.
479 config.db_url = cls._update_sqlite_url(config.db_url)
481 return config
483 def get_replica(self) -> ApdbSqlReplica:
484 """Return `ApdbReplica` instance for this database."""
485 return ApdbSqlReplica(self._schema, self._engine, self._db_schema_version)
487 def tableRowCount(self) -> dict[str, int]:
488 """Return dictionary with the table names and row counts.
490 Used by ``ap_proto`` to keep track of the size of the database tables.
491 Depending on database technology this could be expensive operation.
493 Returns
494 -------
495 row_counts : `dict`
496 Dict where key is a table name and value is a row count.
497 """
498 res = {}
499 tables = [ApdbTables.DiaObject, ApdbTables.DiaSource, ApdbTables.DiaForcedSource]
500 if self.config.dia_object_index == "last_object_table":
501 tables.append(ApdbTables.DiaObjectLast)
502 with self._engine.begin() as conn:
503 for table in tables:
504 sa_table = self._schema.get_table(table)
505 stmt = sql.select(func.count()).select_from(sa_table)
506 count: int = conn.execute(stmt).scalar_one()
507 res[table.name] = count
509 return res
511 def getConfig(self) -> ApdbSqlConfig:
512 # docstring is inherited from a base class
513 return self.config
515 def tableDef(self, table: ApdbTables) -> Table | None:
516 # docstring is inherited from a base class
517 return self._schema.tableSchemas.get(table)
519 @classmethod
520 def _makeSchema(cls, config: ApdbConfig, drop: bool = False) -> None:
521 # docstring is inherited from a base class
523 if not isinstance(config, ApdbSqlConfig): 523 ↛ 524line 523 didn't jump to line 524 because the condition on line 523 was never true
524 raise TypeError(f"Unexpected type of configuration object: {type(config)}")
526 engine = cls._makeEngine(config, create=True)
528 table_schema = ApdbSchema(config.schema_file, config.ss_schema_file)
530 # Ask schema class to create all tables.
531 schema = ApdbSqlSchema(
532 table_schema=table_schema,
533 engine=engine,
534 dia_object_index=config.dia_object_index,
535 prefix=config.prefix,
536 namespace=config.namespace,
537 htm_index_column=config.pixelization.htm_index_column,
538 enable_replica=config.enable_replica,
539 )
540 schema.makeSchema(drop=drop)
542 # Need metadata table to store few items in it.
543 meta_table = schema.get_table(ApdbTables.metadata)
544 apdb_meta = ApdbMetadataSql(engine, meta_table)
546 # Fill version numbers, overwrite if they are already there.
547 apdb_meta.set(cls.metadataSchemaVersionKey, str(table_schema.schemaVersion()), force=True)
548 apdb_meta.set(cls.metadataCodeVersionKey, str(cls.apdbImplementationVersion()), force=True)
549 if config.enable_replica:
550 # Only store replica code version if replica is enabled.
551 apdb_meta.set(
552 cls.metadataReplicaVersionKey,
553 str(ApdbSqlReplica.apdbReplicaImplementationVersion()),
554 force=True,
555 )
557 # Store frozen part of a configuration in metadata.
558 freezer = ApdbConfigFreezer[ApdbSqlConfig](cls._frozen_parameters)
559 apdb_meta.set(cls.metadataConfigKey, freezer.to_json(config), force=True)
561 engine.dispose()
563 def getDiaObjects(self, region: Region) -> pandas.DataFrame:
564 # docstring is inherited from a base class
566 # decide what columns we need
567 if self.config.dia_object_index == "last_object_table":
568 table_enum = ApdbTables.DiaObjectLast
569 else:
570 table_enum = ApdbTables.DiaObject
571 table = self._schema.get_table(table_enum)
572 if not self.config.dia_object_columns: 572 ↛ 575line 572 didn't jump to line 575 because the condition on line 572 was always true
573 columns = self._schema.get_apdb_columns(table_enum)
574 else:
575 columns = [table.c[col] for col in self.config.dia_object_columns]
576 query = sql.select(*columns)
578 # build selection
579 query = query.where(self._filterRegion(table, region))
581 validity_end_column = self._timestamp_column_name("validityEnd")
583 # select latest version of objects
584 if self.config.dia_object_index != "last_object_table":
585 query = query.where(table.columns[validity_end_column] == None) # noqa: E711
587 # _LOG.debug("query: %s", query)
589 # execute select
590 with self._timer("select_time", tags={"table": "DiaObject"}) as timer:
591 with self._engine.begin() as conn:
592 result = conn.execute(query)
593 column_defs = self._table_schema.tableSchemas[table_enum].columns
594 table_data = ApdbSqlTableData(result, column_defs)
595 objects = table_data.to_pandas()
596 timer.add_values(row_count=len(objects))
597 _LOG.debug("found %s DiaObjects", len(objects))
598 return self._fix_result_timestamps(objects)
600 def getDiaSources(
601 self,
602 region: Region,
603 object_ids: Iterable[int] | None,
604 visit_time: astropy.time.Time,
605 start_time: astropy.time.Time | None = None,
606 ) -> pandas.DataFrame | None:
607 # docstring is inherited from a base class
608 if start_time is None and self.config.read_sources_months == 0:
609 _LOG.debug("Skip DiaSources fetching")
610 return None
612 if start_time is None:
613 start_time_mjdTai = _make_midpointMjdTai_start(visit_time, self.config.read_sources_months)
614 else:
615 start_time_mjdTai = float(start_time.tai.mjd)
616 _LOG.debug("start_time_mjdTai = %.6f", start_time_mjdTai)
618 if object_ids is None:
619 # region-based select
620 return self._getDiaSourcesInRegion(region, start_time_mjdTai)
621 else:
622 return self._getDiaSourcesByIDs(list(object_ids), start_time_mjdTai)
624 def getDiaForcedSources(
625 self,
626 region: Region,
627 object_ids: Iterable[int] | None,
628 visit_time: astropy.time.Time,
629 start_time: astropy.time.Time | None = None,
630 ) -> pandas.DataFrame | None:
631 # docstring is inherited from a base class
632 if start_time is None and self.config.read_forced_sources_months == 0:
633 _LOG.debug("Skip DiaForceSources fetching")
634 return None
636 if object_ids is None:
637 # This implementation does not support region-based selection. In
638 # the past DiaForcedSource schema did not have ra/dec columns (it
639 # had x/y columns). ra/dec were added at some point, so we could
640 # add pixelOd column to this table if/when needed.
641 raise NotImplementedError("Region-based selection is not supported")
643 # TODO: DateTime.MJD must be consistent with code in ap_association,
644 # alternatively we can fill midpointMjdTai ourselves in store()
645 if start_time is None:
646 start_time_mjdTai = _make_midpointMjdTai_start(visit_time, self.config.read_forced_sources_months)
647 else:
648 start_time_mjdTai = float(start_time.tai.mjd)
649 _LOG.debug("start_time_mjdTai = %.6f", start_time_mjdTai)
651 with self._timer("select_time", tags={"table": "DiaForcedSource"}) as timer:
652 sources = self._getSourcesByIDs(ApdbTables.DiaForcedSource, list(object_ids), start_time_mjdTai)
653 timer.add_values(row_count=len(sources))
655 _LOG.debug("found %s DiaForcedSources", len(sources))
656 return sources
658 def getDiaObjectsForDedup(self, since: astropy.time.Time | None = None) -> pandas.DataFrame:
659 # docstring is inherited from a base class
661 if since is None:
662 # Read last deduplication time from metadata.
663 dedup_str = self._metadata.get(self.metadataDedupKey)
664 if dedup_str is not None:
665 dedup_state = json.loads(dedup_str)
666 dedup_time_str = dedup_state["dedup_time_iso_tai"]
667 since = astropy.time.Time(dedup_time_str, format="iso", scale="tai")
669 validity_start_column = self._timestamp_column_name("validityStart")
671 # decide what columns we need
672 if self.config.dia_object_index == "last_object_table":
673 table_enum = ApdbTables.DiaObjectLast
674 else:
675 table_enum = ApdbTables.DiaObject
676 table = self._schema.get_table(table_enum)
678 if not self.config.dia_object_columns_for_dedup: 678 ↛ 679line 678 didn't jump to line 679 because the condition on line 678 was never true
679 columns = self._schema.get_apdb_columns(table_enum)
680 else:
681 column_names = list(self.config.dia_object_columns_for_dedup)
682 if validity_start_column not in column_names: 682 ↛ 684line 682 didn't jump to line 684 because the condition on line 682 was always true
683 column_names.insert(0, validity_start_column)
684 if "diaObjectId" not in column_names: 684 ↛ 686line 684 didn't jump to line 686 because the condition on line 684 was always true
685 column_names.insert(0, "diaObjectId")
686 columns = [table.columns[col] for col in column_names]
688 query = sql.select(*columns)
690 # build selection
691 if since is not None:
692 timestamp = self._timestamp_column_value(since)
693 query = query.where(table.columns[validity_start_column] >= timestamp)
695 # execute select
696 with self._timer("select_time", tags={"table": "DiaObject"}) as timer:
697 with self._engine.begin() as conn:
698 result = conn.execute(query)
699 column_defs = self._table_schema.tableSchemas[table_enum].columns
700 table_data = ApdbSqlTableData(result, column_defs)
701 objects = table_data.to_pandas()
702 timer.add_values(row_count=len(objects))
703 _LOG.debug("found %s DiaObjects", len(objects))
704 return self._fix_result_timestamps(objects)
706 def getDiaSourcesForDiaObjects(
707 self, objects: list[DiaObjectId], start_time: astropy.time.Time, max_dist_arcsec: float = 1.0
708 ) -> pandas.DataFrame:
709 # docstring is inherited from a base class
710 object_ids = {object_id.diaObjectId for object_id in objects}
711 return self._getDiaSourcesByIDs(list(object_ids), float(start_time.tai.mjd))
713 def containsVisitDetector(
714 self,
715 visit: int,
716 detector: int,
717 region: Region | None = None,
718 visit_time: astropy.time.Time | None = None,
719 ) -> bool:
720 # docstring is inherited from a base class
721 src_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaSource)
722 frcsrc_table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaForcedSource)
723 # Query should load only one leaf page of the index
724 query1 = sql.select(src_table.c.visit).filter_by(visit=visit, detector=detector).limit(1)
726 with self._engine.begin() as conn:
727 result = conn.execute(query1).scalar_one_or_none()
728 if result is not None:
729 return True
730 else:
731 # Backup query if an image was processed but had no diaSources
732 query2 = sql.select(frcsrc_table.c.visit).filter_by(visit=visit, detector=detector).limit(1)
733 result = conn.execute(query2).scalar_one_or_none()
734 return result is not None
736 def store(
737 self,
738 visit_time: astropy.time.Time,
739 objects: pandas.DataFrame,
740 sources: pandas.DataFrame | None = None,
741 forced_sources: pandas.DataFrame | None = None,
742 ) -> None:
743 # docstring is inherited from a base class
744 objects = self._fix_input_timestamps(objects)
745 if sources is not None:
746 sources = self._fix_input_timestamps(sources)
747 if forced_sources is not None:
748 forced_sources = self._fix_input_timestamps(forced_sources)
750 # We want to run all inserts in one transaction.
751 with self._engine.begin() as connection:
752 replica_chunk: ReplicaChunk | None = None
753 if self._schema.replication_enabled:
754 replica_chunk = ReplicaChunk.make_replica_chunk(visit_time, self.config.replica_chunk_seconds)
755 self._storeReplicaChunk(replica_chunk, connection)
757 # fill pixelId column for DiaObjects
758 objects = self._add_spatial_index(objects)
759 self._storeDiaObjects(objects, visit_time, replica_chunk, connection)
761 if sources is not None:
762 # fill pixelId column for DiaSources
763 sources = self._add_spatial_index(sources)
764 self._storeDiaSources(sources, replica_chunk, connection)
766 if forced_sources is not None:
767 self._storeDiaForcedSources(forced_sources, replica_chunk, connection)
769 def reassignDiaSourcesToDiaObjects(
770 self,
771 idMap: Mapping[DiaSourceId, int],
772 *,
773 increment_nDiaSources: bool = True,
774 decrement_nDiaSources: bool = True,
775 ) -> None:
776 # docstring is inherited from a base class
778 new_object_ids = set(idMap.values())
779 source_ids = {source.diaSourceId for source in idMap}
781 current_time = self._current_time()
782 current_time_ns = int(current_time.unix_tai * 1e9)
784 with self._engine.begin() as conn:
785 # Make sure that all DiaSources exist.
786 found_sources = self._get_diasource_data(conn, source_ids, "diaObjectId")
787 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}): 787 ↛ 788line 787 didn't jump to line 788 because the condition on line 787 was never true
788 raise LookupError(f"Some source IDs are missing from DiaSource table: {missing_ids}")
789 original_object_ids = {row.diaSourceId: row.diaObjectId for row in found_sources}
791 # Make sure that all DiaObjects exist, we also want to know
792 # nDiaSources count for current and new records because we want to
793 # send updated values to replica.
794 all_object_ids = new_object_ids | set(original_object_ids.values())
795 found_objects = self._get_diaobject_data(conn, all_object_ids, "ra", "dec", "nDiaSources")
796 if missing_ids := (new_object_ids - {row.diaObjectId for row in found_objects}): 796 ↛ 797line 796 didn't jump to line 797 because the condition on line 796 was never true
797 raise LookupError(f"Some object IDs are missing from DiaObject table: {missing_ids}")
799 found_objects_by_id = {row.diaObjectId: row for row in found_objects}
801 update_records: list[ApdbUpdateRecord] = []
802 update_order = 0
804 # Update DiaSources.
805 source_table = self._schema.get_table(ApdbTables.DiaSource)
806 for source, diaObjectId in idMap.items():
807 update = (
808 source_table.update()
809 .where(source_table.columns["diaSourceId"] == source.diaSourceId)
810 .values(diaObjectId=diaObjectId)
811 )
812 conn.execute(update)
814 if self._schema.replication_enabled:
815 update_records.append(
816 ApdbReassignDiaSourceToDiaObjectRecord(
817 diaSourceId=source.diaSourceId,
818 ra=source.ra,
819 dec=source.dec,
820 midpointMjdTai=source.midpointMjdTai,
821 diaObjectId=diaObjectId,
822 update_time_ns=current_time_ns,
823 update_order=update_order,
824 )
825 )
826 update_order += 1
828 # DiaObject tables to update.
829 object_tables = [self._schema.get_table(ApdbTables.DiaObject)]
830 if self.config.dia_object_index == "last_object_table":
831 object_tables.append(self._schema.get_table(ApdbTables.DiaObjectLast))
833 # Things to increment/decrement.
834 increments: Counter = Counter()
835 if increment_nDiaSources:
836 increments.update(idMap.values())
837 if decrement_nDiaSources:
838 increments.subtract(original_object_ids[source_id.diaSourceId] for source_id in idMap)
840 if increments:
841 for table in object_tables:
842 for diaObjectId, increment in increments.items():
843 update = (
844 table.update()
845 .where(table.columns["diaObjectId"] == diaObjectId)
846 .values(nDiaSources=table.columns["nDiaSources"] + increment)
847 )
848 conn.execute(update)
850 # Also send updated values to replica.
851 if self._schema.replication_enabled:
852 for diaObjectId, increment in increments.items():
853 dia_object = found_objects_by_id[diaObjectId]
854 update_records.append(
855 ApdbUpdateNDiaSourcesRecord(
856 diaObjectId=diaObjectId,
857 ra=dia_object.ra,
858 dec=dia_object.dec,
859 nDiaSources=dia_object.nDiaSources + increment,
860 update_time_ns=current_time_ns,
861 update_order=update_order,
862 )
863 )
864 update_order += 1
866 if update_records:
867 replica_chunk = ReplicaChunk.make_replica_chunk(
868 current_time, self.config.replica_chunk_seconds
869 )
870 self._storeUpdateRecords(update_records, replica_chunk, connection=conn, store_chunk=True)
872 def setValidityEnd(
873 self, objects: list[DiaObjectId], validityEnd: astropy.time.Time, raise_on_missing_id: bool = False
874 ) -> int:
875 # docstring is inherited from a base class
877 with self._engine.begin() as conn:
878 return self._setValidityEnd(conn, objects, validityEnd, raise_on_missing_id=raise_on_missing_id)
880 def _setValidityEnd(
881 self,
882 conn: sqlalchemy.Connection,
883 objects: list[DiaObjectId],
884 validityEnd: astropy.time.Time,
885 raise_on_missing_id: bool = False,
886 ) -> int:
887 # docstring is inherited from a base class
888 if not objects: 888 ↛ 889line 888 didn't jump to line 889 because the condition on line 888 was never true
889 return 0
891 requested_ids = {obj.diaObjectId for obj in objects}
893 validity_end_column = self._timestamp_column_name("validityEnd")
894 validityEnd_value = self._timestamp_column_value(validityEnd)
896 # Find all matching DiaObjects with validityEnd = NULL.
897 table = self._schema.get_table(ApdbTables.DiaObject)
898 query = sql.select(table.columns["diaObjectId"]).where(
899 sqlalchemy.and_(
900 table.columns["diaObjectId"].in_(sorted(requested_ids)),
901 table.columns[validity_end_column].is_(None),
902 )
903 )
905 result = conn.execute(query)
906 found_ids = set(result.scalars())
908 # Check that we found all that is requested.
909 if raise_on_missing_id:
910 if missing_ids := (requested_ids - found_ids): 910 ↛ 914line 910 didn't jump to line 914 because the condition on line 910 was always true
911 raise LookupError(f"Some object IDs are missing from DiaObject table: {missing_ids}")
913 # Filter existing records.
914 if len(objects) != len(found_ids):
915 objects = [obj for obj in objects if obj.diaObjectId in found_ids]
917 if not objects:
918 return 0
920 values = {validity_end_column: validityEnd_value}
921 update = (
922 table.update()
923 .where(
924 sqlalchemy.and_(
925 table.columns["diaObjectId"].in_(sorted(found_ids)),
926 table.columns[validity_end_column].is_(None),
927 )
928 )
929 .values(**values)
930 )
931 result = conn.execute(update)
932 if result.rowcount != len(found_ids): 932 ↛ 933line 932 didn't jump to line 933 because the condition on line 932 was never true
933 raise RuntimeError(
934 f"Unexpected mismatch in the number of records updated. Object IDs = {found_ids}"
935 )
937 # Also drop them from DiaObjectLast.
938 if self.config.dia_object_index == "last_object_table":
939 last_table = self._schema.get_table(ApdbTables.DiaObjectLast)
940 delete = last_table.delete().where(last_table.columns["diaObjectId"].in_(sorted(found_ids)))
941 result = conn.execute(delete)
942 if result.rowcount != len(found_ids): 942 ↛ 943line 942 didn't jump to line 943 because the condition on line 942 was never true
943 raise RuntimeError(
944 f"Unexpected mismatch in the number of records deleted. Object IDs = {found_ids}"
945 )
947 # If replication is enabled then send all updates.
948 if self._schema.replication_enabled:
949 current_time = self._current_time()
950 current_time_ns = int(current_time.unix_tai * 1e9)
951 replica_chunk = ReplicaChunk.make_replica_chunk(current_time, self.config.replica_chunk_seconds)
953 update_records = [
954 ApdbCloseDiaObjectValidityRecord(
955 diaObjectId=obj.diaObjectId,
956 ra=obj.ra,
957 dec=obj.dec,
958 update_time_ns=current_time_ns,
959 update_order=index,
960 validityEndMjdTai=float(validityEnd.tai.mjd),
961 nDiaSources=None,
962 )
963 for index, obj in enumerate(objects)
964 ]
966 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True, connection=conn)
968 return len(objects)
970 def resetDedup(self, dedup_time: astropy.time.Time | None = None) -> None:
971 # docstring is inherited from a base class
973 # SQL backend does not have separate dedup tables, nothing to delete,
974 # only save last dedup time in metadata.
975 if dedup_time is None: 975 ↛ 976line 975 didn't jump to line 976 because the condition on line 975 was never true
976 dedup_time = self._current_time()
977 data = {"dedup_time_iso_tai": dedup_time.tai.to_value("iso")}
978 data_json = json.dumps(data)
979 self._metadata.set(self.metadataDedupKey, data_json, force=True)
981 def reassignDiaSources(self, idMap: Mapping[int, int]) -> None:
982 # docstring is inherited from a base class
984 timestamp: float | datetime.datetime
985 now = self._current_time()
986 timestamp_column = self._timestamp_column_name("ssObjectReassocTime")
987 timestamp = self._timestamp_column_value(now)
989 table = self._schema.get_table(ApdbTables.DiaSource)
990 query = table.update().where(table.columns["diaSourceId"] == sql.bindparam("srcId"))
992 with self._engine.begin() as conn:
993 # Need to make sure that every ID exists in the database, but
994 # executemany may not support rowcount, so iterate and check what
995 # is missing.
996 missing_ids: list[int] = []
997 for key, value in idMap.items():
998 params = {
999 "srcId": key,
1000 "diaObjectId": 0,
1001 "ssObjectId": value,
1002 timestamp_column: timestamp,
1003 }
1004 result = conn.execute(query, params)
1005 if result.rowcount == 0:
1006 missing_ids.append(key)
1007 if missing_ids:
1008 missing = ",".join(str(item) for item in missing_ids)
1009 raise ValueError(f"Following DiaSource IDs do not exist in the database: {missing}")
1011 def withdrawDiaSources(
1012 self,
1013 diaSourceIds: Iterable[DiaSourceId],
1014 *,
1015 timeWithdrawn: astropy.time.Time | None = None,
1016 ) -> None:
1017 # docstring is inherited from a base class
1019 diaSourceIds = list(diaSourceIds)
1020 source_ids = {source.diaSourceId for source in diaSourceIds}
1022 if timeWithdrawn is None: 1022 ↛ 1023line 1022 didn't jump to line 1023 because the condition on line 1022 was never true
1023 timeWithdrawn = self._current_time()
1024 time_value = self._timestamp_column_value(timeWithdrawn)
1025 column_name = self._timestamp_column_name("time_withdrawn")
1027 with self._engine.begin() as conn:
1028 # Make sure that all DiaSources exist.
1029 found_sources = self._get_diasource_data(conn, source_ids, "diaObjectId", column_name)
1030 if missing_ids := (source_ids - {row.diaSourceId for row in found_sources}): 1030 ↛ 1031line 1030 didn't jump to line 1031 because the condition on line 1030 was never true
1031 raise LookupError(f"Some source IDs are missing from DiaSource table: {missing_ids}")
1033 # Ignore sources already withdrawn.
1034 source_ids = {row.diaSourceId for row in found_sources if getattr(row, column_name) is None}
1035 diaSourceIds = [source for source in diaSourceIds if source.diaSourceId in source_ids]
1037 # Set time_withdrawn for sources.
1038 table = self._schema.get_table(ApdbTables.DiaSource)
1039 where = table.columns["diaSourceId"].in_(sorted(source_ids))
1040 update = table.update().where(where).values({column_name: time_value})
1041 conn.execute(update)
1043 # If replication is enabled then send all updates.
1044 if self._schema.replication_enabled:
1045 current_time = self._current_time()
1046 current_time_ns = int(current_time.unix_tai * 1e9)
1047 replica_chunk = ReplicaChunk.make_replica_chunk(
1048 current_time, self.config.replica_chunk_seconds
1049 )
1051 update_records = [
1052 ApdbWithdrawDiaSourceRecord(
1053 diaSourceId=source.diaSourceId,
1054 ra=source.ra,
1055 dec=source.dec,
1056 midpointMjdTai=source.midpointMjdTai,
1057 update_time_ns=current_time_ns,
1058 update_order=update_order,
1059 timeWithdrawnMjdTai=float(timeWithdrawn.tai.mjd),
1060 )
1061 for update_order, source in enumerate(diaSourceIds)
1062 ]
1063 if update_records: 1063 ↛ exitline 1063 didn't jump to the function exit
1064 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True, connection=conn)
1066 def withdrawDiaForcedSources(
1067 self,
1068 diaForcedSourceIds: Iterable[DiaForcedSourceId],
1069 *,
1070 timeWithdrawn: astropy.time.Time | None = None,
1071 ) -> None:
1072 # docstring is inherited from a base class
1074 def _fsrc_id(fsource: Any) -> tuple[int, int, int]:
1075 return (fsource.diaObjectId, fsource.visit, fsource.detector)
1077 diaForcedSourceIds = list(diaForcedSourceIds)
1078 source_ids = {_fsrc_id(source) for source in diaForcedSourceIds}
1080 if timeWithdrawn is None: 1080 ↛ 1081line 1080 didn't jump to line 1081 because the condition on line 1080 was never true
1081 timeWithdrawn = self._current_time()
1082 time_value = self._timestamp_column_value(timeWithdrawn)
1083 column_name = self._timestamp_column_name("time_withdrawn")
1085 with self._engine.begin() as conn:
1086 # Make sure that all DiaForcedSources exist.
1087 table = self._schema.get_table(ApdbTables.DiaForcedSource)
1088 id_columns = [table.columns["diaObjectId"], table.columns["visit"], table.columns["detector"]]
1089 where = sqlalchemy.tuple_(*id_columns).in_(sorted(source_ids))
1090 query = sql.select(table.columns[column_name], *id_columns).where(where)
1091 result = list(conn.execute(query))
1092 if missing_ids := (source_ids - {_fsrc_id(row) for row in result}): 1092 ↛ 1093line 1092 didn't jump to line 1093 because the condition on line 1092 was never true
1093 raise LookupError(f"Some source IDs are missing from DiaForcedSource table: {missing_ids}")
1095 # Ignore sources already withdrawn.
1096 source_ids = {_fsrc_id(row) for row in result if getattr(row, column_name) is None}
1097 diaForcedSourceIds = [source for source in diaForcedSourceIds if _fsrc_id(source) in source_ids]
1099 where = sqlalchemy.tuple_(*id_columns).in_(sorted(source_ids))
1100 update = table.update().where(where).values({column_name: time_value})
1101 conn.execute(update)
1103 # If replication is enabled then send all updates.
1104 if self._schema.replication_enabled:
1105 current_time = self._current_time()
1106 current_time_ns = int(current_time.unix_tai * 1e9)
1107 replica_chunk = ReplicaChunk.make_replica_chunk(
1108 current_time, self.config.replica_chunk_seconds
1109 )
1111 update_records = [
1112 ApdbWithdrawDiaForcedSourceRecord(
1113 diaObjectId=source.diaObjectId,
1114 visit=source.visit,
1115 detector=source.detector,
1116 ra=source.ra,
1117 dec=source.dec,
1118 midpointMjdTai=source.midpointMjdTai,
1119 update_time_ns=current_time_ns,
1120 update_order=index,
1121 timeWithdrawnMjdTai=float(timeWithdrawn.tai.mjd),
1122 )
1123 for index, source in enumerate(diaForcedSourceIds)
1124 ]
1126 self._storeUpdateRecords(update_records, replica_chunk, store_chunk=True, connection=conn)
1128 def countUnassociatedObjects(self) -> int:
1129 # docstring is inherited from a base class
1131 # Retrieve the DiaObject table.
1132 table: sqlalchemy.schema.Table = self._schema.get_table(ApdbTables.DiaObject)
1134 # Construct the sql statement.
1135 validity_end_column = self._timestamp_column_name("validityEnd")
1136 stmt = (
1137 sql.select(func.count())
1138 .select_from(table)
1139 .where(
1140 sqlalchemy.and_(
1141 table.columns["nDiaSources"] == 1,
1142 table.columns[validity_end_column].is_(None),
1143 )
1144 )
1145 )
1147 # Return the count.
1148 with self._engine.begin() as conn:
1149 count = conn.execute(stmt).scalar_one()
1151 return count
1153 @property
1154 def schema(self) -> ApdbSchema:
1155 # docstring is inherited from a base class
1156 return self._table_schema
1158 @property
1159 def metadata(self) -> ApdbMetadata:
1160 # docstring is inherited from a base class
1161 return self._metadata
1163 @property
1164 def admin(self) -> ApdbSqlAdmin:
1165 # docstring is inherited from a base class
1166 return ApdbSqlAdmin(self.pixelator)
1168 def _getDiaSourcesInRegion(self, region: Region, start_time_mjdTai: float) -> pandas.DataFrame:
1169 """Return catalog of DiaSource instances from given region.
1171 Parameters
1172 ----------
1173 region : `lsst.sphgeom.Region`
1174 Region to search for DIASources.
1175 start_time_mjdTai : `float`
1176 Lower bound of time window for the query.
1178 Returns
1179 -------
1180 catalog : `pandas.DataFrame`
1181 Catalog containing DiaSource records.
1182 """
1183 table = self._schema.get_table(ApdbTables.DiaSource)
1184 columns = self._schema.get_apdb_columns(ApdbTables.DiaSource)
1185 query = sql.select(*columns)
1187 # build selection
1188 time_filter = table.columns["midpointMjdTai"] > start_time_mjdTai
1189 where = sql.expression.and_(self._filterRegion(table, region), time_filter)
1190 query = query.where(where)
1192 # execute select
1193 with self._timer("DiaSource_select_time", tags={"table": "DiaSource"}) as timer:
1194 with self._engine.begin() as conn:
1195 result = conn.execute(query)
1196 column_defs = self._table_schema.tableSchemas[ApdbTables.DiaSource].columns
1197 table_data = ApdbSqlTableData(result, column_defs)
1198 sources = table_data.to_pandas()
1199 timer.add_values(row_counts=len(sources))
1200 _LOG.debug("found %s DiaSources", len(sources))
1201 return self._fix_result_timestamps(sources)
1203 def _getDiaSourcesByIDs(self, object_ids: list[int], start_time_mjdTai: float) -> pandas.DataFrame:
1204 """Return catalog of DiaSource instances given set of DiaObject IDs.
1206 Parameters
1207 ----------
1208 object_ids :
1209 Collection of DiaObject IDs
1210 start_time_mjdTai : `float`
1211 Lower bound of time window for the query.
1213 Returns
1214 -------
1215 catalog : `pandas.DataFrame`
1216 Catalog containing DiaSource records.
1217 """
1218 with self._timer("select_time", tags={"table": "DiaSource"}) as timer:
1219 sources = self._getSourcesByIDs(ApdbTables.DiaSource, object_ids, start_time_mjdTai)
1220 timer.add_values(row_count=len(sources))
1222 _LOG.debug("found %s DiaSources", len(sources))
1223 return sources
1225 def _getSourcesByIDs(
1226 self, table_enum: ApdbTables, object_ids: list[int], midpointMjdTai_start: float
1227 ) -> pandas.DataFrame:
1228 """Return catalog of DiaSource or DiaForcedSource instances given set
1229 of DiaObject IDs.
1231 Parameters
1232 ----------
1233 table : `sqlalchemy.schema.Table`
1234 Database table.
1235 object_ids :
1236 Collection of DiaObject IDs
1237 midpointMjdTai_start : `float`
1238 Earliest midpointMjdTai to retrieve.
1240 Returns
1241 -------
1242 catalog : `pandas.DataFrame`
1243 Catalog contaning DiaSource records. `None` is returned if
1244 ``read_sources_months`` configuration parameter is set to 0 or
1245 when ``object_ids`` is empty.
1246 """
1247 table = self._schema.get_table(table_enum)
1248 columns = self._schema.get_apdb_columns(table_enum)
1249 column_defs = self._table_schema.tableSchemas[table_enum].columns
1251 sources: pandas.DataFrame | None = None
1252 if len(object_ids) <= 0:
1253 _LOG.debug("ID list is empty, just fetch empty result")
1254 query = sql.select(*columns).where(sql.literal(False))
1255 with self._engine.begin() as conn:
1256 result = conn.execute(query)
1257 table_data = ApdbSqlTableData(result, column_defs)
1258 sources = table_data.to_pandas()
1259 else:
1260 data_frames: list[pandas.DataFrame] = []
1261 for ids in chunk_iterable(sorted(object_ids), 1000):
1262 query = sql.select(*columns)
1264 # Some types like np.int64 can cause issues with
1265 # sqlalchemy, convert them to int.
1266 int_ids = [int(oid) for oid in ids]
1268 # select by object id
1269 query = query.where(
1270 sql.expression.and_(
1271 table.columns["diaObjectId"].in_(int_ids),
1272 table.columns["midpointMjdTai"] >= midpointMjdTai_start,
1273 )
1274 )
1276 # execute select
1277 with self._engine.begin() as conn:
1278 result = conn.execute(query)
1279 table_data = ApdbSqlTableData(result, column_defs)
1280 data_frames.append(table_data.to_pandas())
1282 if len(data_frames) == 1: 1282 ↛ 1285line 1282 didn't jump to line 1285 because the condition on line 1282 was always true
1283 sources = data_frames[0]
1284 else:
1285 sources = pandas.concat(data_frames)
1286 assert sources is not None, "Catalog cannot be None"
1287 return self._fix_result_timestamps(sources)
1289 def _storeReplicaChunk(
1290 self,
1291 replica_chunk: ReplicaChunk,
1292 connection: sqlalchemy.engine.Connection,
1293 ) -> None:
1294 # `visit_time.datetime` returns naive datetime, even though all astropy
1295 # times are in UTC. Add UTC timezone to timestamp so that database
1296 # can store a correct value.
1297 dt = datetime.datetime.fromtimestamp(replica_chunk.last_update_time.unix_tai, tz=datetime.UTC)
1299 table = self._schema.get_table(ExtraTables.ApdbReplicaChunks)
1301 # We need UPSERT which is dialect-specific construct
1302 values = {"last_update_time": dt, "unique_id": replica_chunk.unique_id}
1303 row = {"apdb_replica_chunk": replica_chunk.id} | values
1304 if connection.dialect.name == "sqlite": 1304 ↛ 1308line 1304 didn't jump to line 1308 because the condition on line 1304 was always true
1305 insert_sqlite = sqlalchemy.dialects.sqlite.insert(table)
1306 insert_sqlite = insert_sqlite.on_conflict_do_update(index_elements=table.primary_key, set_=values)
1307 connection.execute(insert_sqlite, row)
1308 elif connection.dialect.name == "postgresql":
1309 insert_pg = sqlalchemy.dialects.postgresql.dml.insert(table)
1310 insert_pg = insert_pg.on_conflict_do_update(constraint=table.primary_key, set_=values)
1311 connection.execute(insert_pg, row)
1312 else:
1313 raise TypeError(f"Unsupported dialect {connection.dialect.name} for upsert.")
1315 def _storeDiaObjects(
1316 self,
1317 objs: pandas.DataFrame,
1318 visit_time: astropy.time.Time,
1319 replica_chunk: ReplicaChunk | None,
1320 connection: sqlalchemy.engine.Connection,
1321 ) -> None:
1322 """Store catalog of DiaObjects from current visit.
1324 Parameters
1325 ----------
1326 objs : `pandas.DataFrame`
1327 Catalog with DiaObject records.
1328 visit_time : `astropy.time.Time`
1329 Time of the visit.
1330 replica_chunk : `ReplicaChunk`
1331 Insert identifier.
1332 """
1333 if len(objs) == 0:
1334 _LOG.debug("No objects to write to database.")
1335 return
1337 # Some types like np.int64 can cause issues with sqlalchemy, convert
1338 # them to int.
1339 ids = sorted(int(oid) for oid in objs["diaObjectId"])
1340 _LOG.debug("first object ID: %d", ids[0])
1342 validity_start_column = self._timestamp_column_name("validityStart")
1343 validity_end_column = self._timestamp_column_name("validityEnd")
1344 timestamp = self._timestamp_column_value(visit_time)
1346 # everything to be done in single transaction
1347 if self.config.dia_object_index == "last_object_table":
1348 # Insert and replace all records in LAST table.
1349 table = self._schema.get_table(ApdbTables.DiaObjectLast)
1351 # DiaObjectLast did not have this column in the past.
1352 use_validity_start = self._schema.check_column(ApdbTables.DiaObjectLast, validity_start_column)
1354 # Drop the previous objects (pandas cannot upsert).
1355 query = table.delete().where(table.columns["diaObjectId"].in_(ids))
1357 with self._timer("delete_time", tags={"table": table.name}) as timer:
1358 res = connection.execute(query)
1359 timer.add_values(row_count=res.rowcount)
1360 _LOG.debug("deleted %s objects", res.rowcount)
1362 # DiaObjectLast is a subset of DiaObject, strip missing columns
1363 last_column_names = [column.name for column in table.columns]
1364 if validity_start_column in last_column_names and validity_start_column not in objs.columns: 1364 ↛ 1366line 1364 didn't jump to line 1366 because the condition on line 1364 was always true
1365 last_column_names.remove(validity_start_column)
1366 last_objs = objs[last_column_names]
1367 last_objs = _coerce_uint64(last_objs)
1369 # Fill validityStart, only when it is in the schema.
1370 if use_validity_start: 1370 ↛ 1378line 1370 didn't jump to line 1378 because the condition on line 1370 was always true
1371 if validity_start_column in last_objs: 1371 ↛ 1372line 1371 didn't jump to line 1372 because the condition on line 1371 was never true
1372 last_objs[validity_start_column] = timestamp
1373 else:
1374 extra_column = pandas.Series([timestamp] * len(last_objs), name=validity_start_column)
1375 last_objs.set_index(extra_column.index, inplace=True)
1376 last_objs = pandas.concat([last_objs, extra_column], axis="columns")
1378 with self._timer("insert_time", tags={"table": "DiaObjectLast"}) as timer:
1379 last_objs.to_sql(
1380 table.name,
1381 connection,
1382 if_exists="append",
1383 index=False,
1384 schema=table.schema,
1385 )
1386 timer.add_values(row_count=len(last_objs))
1388 # truncate existing validity intervals
1389 table = self._schema.get_table(ApdbTables.DiaObject)
1391 update = (
1392 table.update()
1393 .values(**{validity_end_column: timestamp})
1394 .where(
1395 sql.expression.and_(
1396 table.columns["diaObjectId"].in_(ids),
1397 table.columns[validity_end_column].is_(None),
1398 )
1399 )
1400 )
1402 with self._timer("truncate_time", tags={"table": table.name}) as timer:
1403 res = connection.execute(update)
1404 timer.add_values(row_count=res.rowcount)
1405 _LOG.debug("truncated %s intervals", res.rowcount)
1407 objs = _coerce_uint64(objs)
1409 # Fill additional columns
1410 extra_columns: list[pandas.Series] = []
1411 if validity_start_column in objs.columns: 1411 ↛ 1412line 1411 didn't jump to line 1412 because the condition on line 1411 was never true
1412 objs[validity_start_column] = timestamp
1413 else:
1414 extra_columns.append(pandas.Series([timestamp] * len(objs), name=validity_start_column))
1415 if validity_end_column in objs.columns: 1415 ↛ 1416line 1415 didn't jump to line 1416 because the condition on line 1415 was never true
1416 objs[validity_end_column] = None
1417 else:
1418 extra_columns.append(pandas.Series([None] * len(objs), name=validity_end_column))
1419 if extra_columns: 1419 ↛ 1424line 1419 didn't jump to line 1424 because the condition on line 1419 was always true
1420 objs.set_index(extra_columns[0].index, inplace=True)
1421 objs = pandas.concat([objs] + extra_columns, axis="columns")
1423 # Insert replica data
1424 table = self._schema.get_table(ApdbTables.DiaObject)
1425 replica_data: list[dict] = []
1426 replica_stmt: Any = None
1427 replica_table_name = ""
1428 if replica_chunk is not None:
1429 pk_names = [column.name for column in table.primary_key]
1430 replica_data = objs[pk_names].to_dict("records")
1431 if replica_data: 1431 ↛ 1439line 1431 didn't jump to line 1439 because the condition on line 1431 was always true
1432 for row in replica_data:
1433 row["apdb_replica_chunk"] = replica_chunk.id
1434 replica_table = self._schema.get_table(ExtraTables.DiaObjectChunks)
1435 replica_table_name = replica_table.name
1436 replica_stmt = replica_table.insert()
1438 # insert new versions
1439 with self._timer("insert_time", tags={"table": table.name}) as timer:
1440 objs.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1441 timer.add_values(row_count=len(objs))
1442 if replica_stmt is not None:
1443 with self._timer("insert_time", tags={"table": replica_table_name}) as timer:
1444 connection.execute(replica_stmt, replica_data)
1445 timer.add_values(row_count=len(replica_data))
1447 def _storeDiaSources(
1448 self,
1449 sources: pandas.DataFrame,
1450 replica_chunk: ReplicaChunk | None,
1451 connection: sqlalchemy.engine.Connection,
1452 ) -> None:
1453 """Store catalog of DiaSources from current visit.
1455 Parameters
1456 ----------
1457 sources : `pandas.DataFrame`
1458 Catalog containing DiaSource records
1459 """
1460 table = self._schema.get_table(ApdbTables.DiaSource)
1462 # Insert replica data
1463 replica_data: list[dict] = []
1464 replica_stmt: Any = None
1465 replica_table_name = ""
1466 if replica_chunk is not None:
1467 pk_names = [column.name for column in table.primary_key]
1468 replica_data = sources[pk_names].to_dict("records")
1469 if replica_data: 1469 ↛ 1477line 1469 didn't jump to line 1477 because the condition on line 1469 was always true
1470 for row in replica_data:
1471 row["apdb_replica_chunk"] = replica_chunk.id
1472 replica_table = self._schema.get_table(ExtraTables.DiaSourceChunks)
1473 replica_table_name = replica_table.name
1474 replica_stmt = replica_table.insert()
1476 # everything to be done in single transaction
1477 with self._timer("insert_time", tags={"table": table.name}) as timer:
1478 sources = _coerce_uint64(sources)
1479 sources.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1480 timer.add_values(row_count=len(sources))
1481 if replica_stmt is not None:
1482 with self._timer("replica_insert_time", tags={"table": replica_table_name}) as timer:
1483 connection.execute(replica_stmt, replica_data)
1484 timer.add_values(row_count=len(replica_data))
1486 def _storeDiaForcedSources(
1487 self,
1488 sources: pandas.DataFrame,
1489 replica_chunk: ReplicaChunk | None,
1490 connection: sqlalchemy.engine.Connection,
1491 ) -> None:
1492 """Store a set of DiaForcedSources from current visit.
1494 Parameters
1495 ----------
1496 sources : `pandas.DataFrame`
1497 Catalog containing DiaForcedSource records
1498 """
1499 table = self._schema.get_table(ApdbTables.DiaForcedSource)
1501 # Insert replica data
1502 replica_data: list[dict] = []
1503 replica_stmt: Any = None
1504 replica_table_name = ""
1505 if replica_chunk is not None:
1506 pk_names = [column.name for column in table.primary_key]
1507 replica_data = sources[pk_names].to_dict("records")
1508 if replica_data: 1508 ↛ 1516line 1508 didn't jump to line 1516 because the condition on line 1508 was always true
1509 for row in replica_data:
1510 row["apdb_replica_chunk"] = replica_chunk.id
1511 replica_table = self._schema.get_table(ExtraTables.DiaForcedSourceChunks)
1512 replica_table_name = replica_table.name
1513 replica_stmt = replica_table.insert()
1515 # everything to be done in single transaction
1516 with self._timer("insert_time", tags={"table": table.name}) as timer:
1517 sources = _coerce_uint64(sources)
1518 sources.to_sql(table.name, connection, if_exists="append", index=False, schema=table.schema)
1519 timer.add_values(row_count=len(sources))
1520 if replica_stmt is not None:
1521 with self._timer("insert_time", tags={"table": replica_table_name}) as timer:
1522 connection.execute(replica_stmt, replica_data)
1523 timer.add_values(row_count=len(replica_data))
1525 def _storeUpdateRecords(
1526 self,
1527 records: Iterable[ApdbUpdateRecord],
1528 chunk: ReplicaChunk,
1529 *,
1530 store_chunk: bool = False,
1531 connection: sqlalchemy.engine.Connection | None = None,
1532 ) -> None:
1533 """Store ApdbUpdateRecords in the replica table for those records.
1535 Parameters
1536 ----------
1537 records : `list` [`ApdbUpdateRecord`]
1538 Records to store.
1539 chunk : `ReplicaChunk`
1540 Replica chunk for these records.
1541 store_chunk : `bool`
1542 If True then also store replica chunk.
1543 connection : `sqlalchemy.engine.Connection`
1544 SQLALchemy connection to use, if `None` the new connection will be
1545 made. `None` is useful for tests only, regular use will call this
1546 method in the same transaction that saves other types of records.
1548 Raises
1549 ------
1550 TypeError
1551 Raised if replication is not enabled for this instance.
1552 """
1553 if not self._schema.replication_enabled:
1554 raise TypeError("Replication is not enabled for this APDB instance.")
1556 apdb_replica_chunk = chunk.id
1557 # Do not use unique_if from ReplicaChunk as it could be reused in
1558 # multiple calls to this method.
1559 update_unique_id = uuid.uuid4()
1561 record_dicts = []
1562 for record in records:
1563 record_dicts.append(
1564 {
1565 "apdb_replica_chunk": apdb_replica_chunk,
1566 "update_time_ns": record.update_time_ns,
1567 "update_order": record.update_order,
1568 "update_unique_id": update_unique_id,
1569 "update_payload": record.to_json(),
1570 }
1571 )
1573 if not record_dicts: 1573 ↛ 1574line 1573 didn't jump to line 1574 because the condition on line 1573 was never true
1574 return
1576 # TODO: Need to check that table exists.
1577 table = self._schema.get_table(ExtraTables.ApdbUpdateRecordChunks)
1579 def _do_store(connection: sqlalchemy.engine.Connection) -> None:
1580 if store_chunk: 1580 ↛ 1582line 1580 didn't jump to line 1582 because the condition on line 1580 was always true
1581 self._storeReplicaChunk(chunk, connection)
1582 with self._timer("insert_time", tags={"table": table.name}) as timer:
1583 connection.execute(table.insert(), record_dicts)
1584 timer.add_values(row_count=len(record_dicts))
1586 if connection is None:
1587 with self._engine.begin() as connection:
1588 _do_store(connection)
1589 else:
1590 _do_store(connection)
1592 def _htm_indices(self, region: Region) -> list[tuple[int, int]]:
1593 """Generate a set of HTM indices covering specified region.
1595 Parameters
1596 ----------
1597 region: `sphgeom.Region`
1598 Region that needs to be indexed.
1600 Returns
1601 -------
1602 Sequence of ranges, range is a tuple (minHtmID, maxHtmID).
1603 """
1604 _LOG.debug("region: %s", region)
1605 indices = self.pixelator.envelope(region, self.config.pixelization.htm_max_ranges)
1607 return indices.ranges()
1609 def _filterRegion(self, table: sqlalchemy.schema.Table, region: Region) -> sql.ColumnElement:
1610 """Make SQLAlchemy expression for selecting records in a region."""
1611 htm_index_column = table.columns[self.config.pixelization.htm_index_column]
1612 exprlist = []
1613 pixel_ranges = self._htm_indices(region)
1614 for low, upper in pixel_ranges:
1615 upper -= 1
1616 if low == upper: 1616 ↛ 1617line 1616 didn't jump to line 1617 because the condition on line 1616 was never true
1617 exprlist.append(htm_index_column == low)
1618 else:
1619 exprlist.append(sql.expression.between(htm_index_column, low, upper))
1621 return sql.expression.or_(*exprlist)
1623 def _add_spatial_index(self, df: pandas.DataFrame) -> pandas.DataFrame:
1624 """Calculate spatial index for each record and add it to a DataFrame.
1626 Parameters
1627 ----------
1628 df : `pandas.DataFrame`
1629 DataFrame which has to contain ra/dec columns, names of these
1630 columns are defined by configuration ``ra_dec_columns`` field.
1632 Returns
1633 -------
1634 df : `pandas.DataFrame`
1635 DataFrame with ``pixelId`` column which contains pixel index
1636 for ra/dec coordinates.
1638 Notes
1639 -----
1640 This overrides any existing column in a DataFrame with the same name
1641 (pixelId). Original DataFrame is not changed, copy of a DataFrame is
1642 returned.
1643 """
1644 # calculate HTM index for every DiaObject
1645 htm_index = np.zeros(df.shape[0], dtype=np.int64)
1646 ra_col, dec_col = self.config.ra_dec_columns
1647 for i, (ra, dec) in enumerate(zip(df[ra_col], df[dec_col])):
1648 uv3d = UnitVector3d(LonLat.fromDegrees(ra, dec))
1649 idx = self.pixelator.index(uv3d)
1650 htm_index[i] = idx
1651 df = df.copy()
1652 df[self.config.pixelization.htm_index_column] = htm_index
1653 return df
1655 def _fix_input_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1656 """Update timestamp columns in input DataFrame to be aware datetime
1657 type in in UTC.
1659 AP pipeline generates naive datetime instances, we want them to be
1660 aware before they go to database. All naive timestamps are assumed to
1661 be in UTC timezone (they should be TAI).
1662 """
1663 # Find all columns with aware non-UTC timestamps and convert to UTC.
1664 columns = [
1665 column
1666 for column, dtype in df.dtypes.items()
1667 if isinstance(dtype, pandas.DatetimeTZDtype) and dtype.tz is not datetime.UTC
1668 ]
1669 for column in columns: 1669 ↛ 1670line 1669 didn't jump to line 1670 because the loop on line 1669 never started
1670 df[column] = df[column].dt.tz_convert(datetime.UTC)
1671 # Find all columns with naive timestamps and add UTC timezone.
1672 columns = [
1673 column for column, dtype in df.dtypes.items() if pandas.api.types.is_datetime64_dtype(dtype)
1674 ]
1675 for column in columns:
1676 df[column] = df[column].dt.tz_localize(datetime.UTC)
1677 return df
1679 def _fix_result_timestamps(self, df: pandas.DataFrame) -> pandas.DataFrame:
1680 """Update timestamp columns to be naive datetime type in returned
1681 DataFrame.
1683 AP pipeline code expects DataFrames to contain naive datetime columns,
1684 while Postgres queries return timezone-aware type. This method converts
1685 those columns to naive datetime in UTC timezone.
1686 """
1687 # Find all columns with aware timestamps.
1688 columns = [column for column, dtype in df.dtypes.items() if isinstance(dtype, pandas.DatetimeTZDtype)]
1689 for column in columns: 1689 ↛ 1691line 1689 didn't jump to line 1691 because the loop on line 1689 never started
1690 # tz_convert(None) will convert to UTC and drop timezone.
1691 df[column] = df[column].dt.tz_convert(None)
1692 return df
1694 def _timestamp_column_name(self, column: str) -> str:
1695 """Return column name before/after schema migration to MJD TAI."""
1696 return self.schema.timestamp_column_name(column)
1698 def _timestamp_column_value(self, time: astropy.time.Time) -> float | datetime.datetime:
1699 """Return column value before/after schema migration to MJD TAI."""
1700 if self.schema.has_mjd_timestamps:
1701 return float(time.tai.mjd)
1702 else:
1703 return time.datetime.astimezone(tz=datetime.UTC)
1705 def _get_diaobject_data(
1706 self, conn: sqlalchemy.engine.Connection, object_ids: Iterable[int], *columns: str
1707 ) -> list:
1708 """Select records from either DiaObject or DiaObjectLast and return
1709 selected rows as names tuples.
1710 """
1711 where: sqlalchemy.ColumnElement[bool]
1712 if self.config.dia_object_index == "last_object_table":
1713 table = self._schema.get_table(ApdbTables.DiaObjectLast)
1714 where = table.columns["diaObjectId"].in_(sorted(object_ids))
1715 else:
1716 table = self._schema.get_table(ApdbTables.DiaObject)
1717 validity_end_column = self._timestamp_column_name("validityEnd")
1718 where = sqlalchemy.and_(
1719 table.columns["diaObjectId"].in_(sorted(object_ids)),
1720 table.columns[validity_end_column].is_(None),
1721 )
1722 column_list = [table.columns["diaObjectId"]] + [table.columns[column] for column in columns]
1723 query = sql.select(*column_list).where(where)
1724 result = conn.execute(query)
1726 return list(result)
1728 def _get_diasource_data(
1729 self, conn: sqlalchemy.engine.Connection, source_ids: Iterable[int], *columns: str
1730 ) -> list:
1731 """Select records from DiaSource table by diaSourceId and return
1732 selected rows as named tuples.
1733 """
1734 table = self._schema.get_table(ApdbTables.DiaSource)
1735 where = table.columns["diaSourceId"].in_(sorted(source_ids))
1736 column_list = [table.columns["diaSourceId"]] + [table.columns[column] for column in columns]
1737 query = sql.select(*column_list).where(where)
1738 result = conn.execute(query)
1740 return list(result)