diff --git a/pyproject.toml b/pyproject.toml index 3126ba6..614ec71 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -22,11 +22,13 @@ dependencies = [ "pydantic_settings", "asyncer", "astropy", + "numpy", "scipy", "astroquery", "astropydantic", "pandas", "pyarrow", + "tqdm", "uuid7-standard", ] @@ -48,6 +50,7 @@ socat-act-fits = "socat.ingest.actfits:main" socat-websky-csv = "socat.ingest.webskycsv:main" socat-migrate = "socat.migrate:migrate" socat-jpl-parqet = "socat.ingest.jplparquet:main" +socat-build-db = "socat.ingest.bulk_sqlite:main" [tool.ruff.lint] preview = true diff --git a/socat/alembic/versions/1bfc5561403e_add_sso_id_time_composite_index_to_.py b/socat/alembic/versions/1bfc5561403e_add_sso_id_time_composite_index_to_.py new file mode 100644 index 0000000..619ccd1 --- /dev/null +++ b/socat/alembic/versions/1bfc5561403e_add_sso_id_time_composite_index_to_.py @@ -0,0 +1,38 @@ +"""Add sso_id/time composite index to moving_sources + +Revision ID: 1bfc5561403e +Revises: c3d4e5f6a7b8 +Create Date: 2026-07-17 00:00:00.000000 + +get_ephem_points()/get_source_generator() filter moving_sources by a known +sso_id first, then narrow by time. moving_sources.time already has its own +index (added by 35a6a33e0a34's initial op.create_table), but sso_id was +never indexed -- at real ephemeris-table scale (tens of millions of rows, +see simonsobs/sotrplib's scripts/solar_system/build_socat_db.py), that made +every per-object ephemeris lookup a full table scan. This mirrors the +RegisteredMovingSourceTable.__table_args__ index added to the ORM model in +the same change, so a database provisioned via `alembic upgrade head` ends +up with the same indexes as one provisioned via SQLModel.metadata.create_all(). +""" + +from collections.abc import Sequence + +from alembic import op + +# revision identifiers, used by Alembic. +revision: str = "1bfc5561403e" +down_revision: str | None = "c3d4e5f6a7b8" +branch_labels: str | Sequence[str] | None = None +depends_on: str | Sequence[str] | None = None + + +def upgrade() -> None: + op.create_index( + "idx_moving_sources_sso_time", + "moving_sources", + ["sso_id", "time"], + ) + + +def downgrade() -> None: + op.drop_index("idx_moving_sources_sso_time", table_name="moving_sources") diff --git a/socat/database/sources.py b/socat/database/sources.py index 33d348a..2b3d9e8 100644 --- a/socat/database/sources.py +++ b/socat/database/sources.py @@ -6,6 +6,7 @@ from astropy.time import Time from astropydantic import AstroPydanticICRS, AstroPydanticQuantity, AstroPydanticTime from pydantic import BaseModel +from sqlalchemy import Index from sqlmodel import Field, SQLModel @@ -190,6 +191,13 @@ class RegisteredMovingSourceTable(SQLModel, table=True): """ __tablename__ = "moving_sources" + __table_args__ = ( + # get_ephem_points()/get_source_generator() filter by a known + # sso_id first, then narrow by time -- this composite index + # serves both that and plain sso_id-only lookups (leftmost + # prefix), so no separate single-column sso_id index is needed. + Index("idx_moving_sources_sso_time", "sso_id", "time"), + ) ephem_id: uuid.UUID = Field(primary_key=True, default_factory=uuid.create) sso_id: uuid.UUID = Field( @@ -207,7 +215,11 @@ class RegisteredMovingSourceTable(SQLModel, table=True): nullable=False, ondelete="CASCADE", ) - time: datetime + # get_box_sso()/get_monitored_ssos()/get_pointing_ssos() have no + # sso_id to filter on (that's what they're looking up) -- they filter + # by time directly, so it needs its own leading index too, not just + # as the second column of the composite above. + time: datetime = Field(index=True, nullable=False) ra_deg: float = Field(nullable=False) dec_deg: float = Field(nullable=False) flux_mJy: float | None = Field(nullable=True) diff --git a/socat/database/statements.py b/socat/database/statements.py index b202024..d7a99bf 100644 --- a/socat/database/statements.py +++ b/socat/database/statements.py @@ -11,7 +11,7 @@ from astropy.time import Time from astropy.units import Quantity from astroquery.query import BaseVOQuery -from sqlmodel import select, union_all, update +from sqlmodel import select, union, union_all, update from socat.database.services import AstroqueryServiceTable from socat.database.sources import ( @@ -146,6 +146,15 @@ def get_box_sso( Database statement. """ if lower_left.ra > upper_right.ra: + # Each branch joins against every matching ephemeris point, so + # without its own .distinct() a single object with N matching + # ephem points produces N duplicate rows here -- and the + # .distinct() on the outer `select(...).from_statement(union_stmt)` + # below does NOT fix that: from_statement() replaces the executed + # SQL wholesale, so that outer .distinct() is a no-op. union_all + # (rather than union) is also insufficient on its own: an object + # can legitimately have ephem points on both sides of RA=0 within + # the same time window, and union_all wouldn't merge those. right_box = ( select(SolarSystemObjectTable) .outerjoin( @@ -163,6 +172,7 @@ def get_box_sso( RegisteredMovingSourceTable.dec_deg <= float(upper_right.dec.to_value("deg")), ) + .distinct() ) left_box = ( select(SolarSystemObjectTable) @@ -181,9 +191,10 @@ def get_box_sso( RegisteredMovingSourceTable.dec_deg <= float(upper_right.dec.to_value("deg")), ) + .distinct() ) - union_stmt = union_all(right_box, left_box) - return select(SolarSystemObjectTable).distinct().from_statement(union_stmt) + union_stmt = union(right_box, left_box) + return select(SolarSystemObjectTable).from_statement(union_stmt) else: return ( select(SolarSystemObjectTable) diff --git a/socat/ingest/bulk_sqlite.py b/socat/ingest/bulk_sqlite.py new file mode 100644 index 0000000..c40f3bd --- /dev/null +++ b/socat/ingest/bulk_sqlite.py @@ -0,0 +1,328 @@ +""" +Build a fresh SOCat SQLite database from an ACT-format point-source FITS +catalog (fixed sources) and JPL Horizons batched-ephemeris parquet files +(solar system objects), e.g. the ones produced by sotrplib's +sotrplib/solar_system/download_ephem_from_horizons.py. + +Why not the other ingest modules (socat.ingest.actfits.ingest_fits_file, +socat.ingest.jplparquet.ingest_jpl_parquet_file)? Both go through a +ClientBase (create_source()/create_sso()/create_ephem()), which for the DB +client means one INSERT + COMMIT per row -- appropriate for adding a +handful of sources to an existing catalog, but not for building a database +from scratch. A real ephemeris ingest is a different scale: at these +files' native ~2-hour cadence over 2015-2033, two ephemeris files hold on +the order of tens of millions of rows across a few hundred objects, and +one-commit-per-row would take on the order of days. This module instead +bulk-loads via raw sqlite3.executemany with SQLite tuned for bulk writes +(no WAL/fsync, foreign keys off during load), which gets the whole build +down to minutes. + +The table schema itself is still created from socat's own SQLModel +metadata (see socat.database.sources), so its indexes -- moving_sources +(sso_id, time) and (time) -- are always whatever socat currently defines, +with no separate index-creation step needed here. + +The database file is built on local node disk and then copied to the +requested destination -- SQLite (especially anything that touches the WAL +file) doesn't behave reliably over network filesystems like Lustre/NFS, +and a scratch/home mount is exactly where the final --output usually lives. +""" + +import shutil +import sqlite3 +import tempfile +import time +from pathlib import Path + +import numpy as np +import pyarrow.parquet as pq +import uuid7 +from astropy.io import fits +from sqlalchemy import create_engine +from sqlmodel import SQLModel +from tqdm import tqdm + +_TIME_FMT = "%Y-%m-%d %H:%M:%S.%f" + + +def create_schema(db_path: Path) -> None: + """ + Create SOCat's tables via its own SQLModel metadata, so the on-disk + schema (including indexes) always matches whatever socat currently + defines. + """ + from socat.database import ALL_TABLES # noqa: F401 (registers metadata) + + engine = create_engine(f"sqlite:///{db_path}", future=True) + SQLModel.metadata.create_all(bind=engine) + engine.dispose() + + +def _tune_for_bulk_load(conn: sqlite3.Connection) -> None: + conn.execute("PRAGMA journal_mode=MEMORY") + conn.execute("PRAGMA synchronous=OFF") + conn.execute("PRAGMA foreign_keys=OFF") + conn.execute("PRAGMA temp_store=MEMORY") + conn.execute("PRAGMA cache_size=-500000") # ~500MB page cache + + +def _finalize_pragmas(conn: sqlite3.Connection) -> None: + conn.execute("PRAGMA foreign_keys=ON") + conn.execute("PRAGMA journal_mode=DELETE") + conn.execute("PRAGMA synchronous=NORMAL") + + +def load_fixed_sources( + conn: sqlite3.Connection, + fits_path: Path, + hdu: int = 1, + monitored_flux_threshold_mJy: float = 20.0, + pointing_flux_threshold_mJy: float = 300.0, +) -> int: + """ + Load an ACT point-source FITS catalog into the `fixed_sources` table. + Mirrors socat.ingest.actfits.ingest_fits_file's flagging logic, just + vectorized and bulk-inserted instead of one create_source() per row. + """ + data = fits.open(fits_path)[hdu].data + ra = np.asarray(data["RADeg"], dtype=float) + dec = np.asarray(data["decDeg"], dtype=float) + flux_mJy = np.asarray(data["fluxJy"], dtype=float) * 1000.0 + names = [str(n).strip() for n in data["name"]] + monitored = flux_mJy >= monitored_flux_threshold_mJy + pointing = flux_mJy >= pointing_flux_threshold_mJy + + rows = [ + ( + uuid7.create().hex, + float(ra[i]), + float(dec[i]), + float(flux_mJy[i]), + names[i], + bool(monitored[i]), + bool(pointing[i]), + ) + for i in range(len(data)) + ] + conn.executemany( + "INSERT INTO fixed_sources " + "(source_id, ra_deg, dec_deg, flux_mJy, name, monitored, pointing) " + "VALUES (?,?,?,?,?,?,?)", + rows, + ) + conn.commit() + return len(rows) + + +def _parse_designation(designation: str) -> tuple[int | None, str]: + """Split a designation like '1 Ceres' into MPC ID and name.""" + parts = str(designation).split(maxsplit=1) + try: + mpc_id = int(parts[0]) + except ValueError: + mpc_id = None + name = parts[1] if len(parts) > 1 else str(mpc_id) + return mpc_id, name + + +def collect_sso_designations( + parquet_paths: list[Path], +) -> dict[str, tuple[int | None, str]]: + """ + Enumerate unique SSOs across all ephemeris files without reading the + full ephemeris columns. These files have exactly one row group per + object, so reading just the 'designation' column's first value per + row group is enough -- pyarrow only decodes that one column, not + ra/dec/time. + """ + designations: dict[str, tuple[int | None, str]] = {} + for path in parquet_paths: + pf = pq.ParquetFile(path) + for i in range(pf.metadata.num_row_groups): + designation = pf.read_row_group(i, columns=["designation"])["designation"][ + 0 + ].as_py() + if designation not in designations: + designations[designation] = _parse_designation(designation) + return designations + + +def load_solar_system_objects( + conn: sqlite3.Connection, + designations: dict[str, tuple[int | None, str]], +) -> dict[str, str]: + """Insert one row per unique SSO, return designation -> sso_id.hex.""" + sso_ids: dict[str, str] = {} + rows = [] + for designation, (mpc_id, name) in designations.items(): + sso_id_hex = uuid7.create().hex + sso_ids[designation] = sso_id_hex + rows.append((sso_id_hex, mpc_id, name, True, False)) + conn.executemany( + "INSERT INTO solarsystem_objects (sso_id, MPC_id, name, monitored, pointing) " + "VALUES (?,?,?,?,?)", + rows, + ) + conn.commit() + return sso_ids + + +def load_ephemerides( + conn: sqlite3.Connection, + parquet_paths: list[Path], + sso_ids: dict[str, str], + designations: dict[str, tuple[int | None, str]], + chunk_rows: int = 500_000, +) -> int: + """ + Bulk-load ephemeris points into `moving_sources`, one object (row + group) at a time so peak memory stays bounded regardless of the + total row count across all files. + """ + total = 0 + insert_sql = ( + "INSERT INTO moving_sources " + "(ephem_id, sso_id, MPC_id, name, time, ra_deg, dec_deg, flux_mJy) " + "VALUES (?,?,?,?,?,?,?,?)" + ) + for path in parquet_paths: + pf = pq.ParquetFile(path) + n_groups = pf.metadata.num_row_groups + for i in tqdm(range(n_groups), desc=f"Ingesting {path.name}"): + table = pf.read_row_group( + i, columns=["designation", "datetime_utc", "ra_deg", "dec_deg"] + ) + designation = table.column("designation")[0].as_py() + mpc_id, name = designations[designation] + sso_id_hex = sso_ids[designation] + + times = table.column("datetime_utc").to_pylist() + ras = table.column("ra_deg").to_pylist() + decs = table.column("dec_deg").to_pylist() + + rows = [ + ( + uuid7.create().hex, + sso_id_hex, + mpc_id, + name, + t.strftime(_TIME_FMT), + ra, + dec, + None, + ) + for t, ra, dec in zip(times, ras, decs) + ] + + for start in range(0, len(rows), chunk_rows): + conn.executemany(insert_sql, rows[start : start + chunk_rows]) + conn.commit() + + total += len(rows) + return total + + +def build( + fits_path: Path, + ephem_paths: list[Path], + output_path: Path, + build_dir: Path | None = None, +) -> dict: + """ + Build a full SOCat SQLite database from an ACT FITS catalog and one + or more JPL ephemeris parquet files, writing the result to + `output_path`. Returns a dict of row counts and elapsed time. + """ + build_dir = Path(build_dir or tempfile.mkdtemp(prefix="socat_build_", dir="/tmp")) + build_dir.mkdir(parents=True, exist_ok=True) + build_path = build_dir / output_path.name + + if build_path.exists(): + build_path.unlink() + + t0 = time.time() + create_schema(build_path) + + conn = sqlite3.connect(str(build_path)) + _tune_for_bulk_load(conn) + + n_fixed = load_fixed_sources(conn, fits_path) + print(f"Loaded {n_fixed} fixed sources ({time.time() - t0:.1f}s elapsed)") + + designations = collect_sso_designations(ephem_paths) + sso_ids = load_solar_system_objects(conn, designations) + print( + f"Loaded {len(sso_ids)} solar system objects ({time.time() - t0:.1f}s elapsed)" + ) + + n_ephem = load_ephemerides(conn, ephem_paths, sso_ids, designations) + print(f"Loaded {n_ephem} ephemeris points ({time.time() - t0:.1f}s elapsed)") + + _finalize_pragmas(conn) + conn.execute("ANALYZE") + conn.commit() + conn.close() + + output_path.parent.mkdir(parents=True, exist_ok=True) + shutil.copy2(build_path, output_path) + build_path.unlink() + try: + build_dir.rmdir() + except OSError: + pass + + elapsed = time.time() - t0 + print(f"Wrote {output_path} in {elapsed:.1f}s") + + return { + "n_fixed_sources": n_fixed, + "n_sso": len(sso_ids), + "n_ephem": n_ephem, + "elapsed_sec": elapsed, + } + + +def main(): # pragma: no cover + import argparse as ap + + parser = ap.ArgumentParser( + prog="socat-build-db", + description=( + "Build a SOCat SQLite database from an ACT FITS point-source " + "catalog and one or more JPL Horizons ephemeris parquet files" + ), + ) + parser.add_argument( + "--fits-file", + type=Path, + required=True, + help="ACT-compatible FITS point source catalog (fixed sources)", + ) + parser.add_argument( + "--ephem-file", + type=Path, + action="append", + required=True, + dest="ephem_files", + help="JPL Horizons batched ephemeris parquet file; repeatable", + ) + parser.add_argument("--output", type=Path, required=True) + parser.add_argument( + "--build-dir", + type=Path, + default=None, + help="Local scratch dir to build in before copying to --output " + "(default: a fresh temp dir under /tmp)", + ) + args = parser.parse_args() + + build( + fits_path=args.fits_file, + ephem_paths=args.ephem_files, + output_path=args.output, + build_dir=args.build_dir, + ) + + +if __name__ == "__main__": # pragma: no cover + main() diff --git a/tests/test_bulk_sqlite.py b/tests/test_bulk_sqlite.py new file mode 100644 index 0000000..6c85201 --- /dev/null +++ b/tests/test_bulk_sqlite.py @@ -0,0 +1,192 @@ +""" +Test the fast bulk-sqlite ingest path (socat.ingest.bulk_sqlite), used for +building a full database from files rather than adding sources one at a +time through a live client. +""" + +import sqlite3 + +import astropy.units as u +import numpy as np +import pandas as pd +import pyarrow as pa +import pyarrow.parquet as pq +import pytest +from astropy.coordinates import ICRS +from astropy.io import fits +from astropy.time import Time + +from socat.client.db import Client +from socat.ingest import bulk_sqlite + + +@pytest.fixture +def act_fits_catalog(tmp_path): + n_sources = 10 + data = np.zeros( + n_sources, + dtype=[("raDeg", "f8"), ("decDeg", "f8"), ("fluxJy", "f8"), ("name", "S20")], + ) + data["raDeg"] = np.linspace(0, 359, n_sources) + data["decDeg"] = np.linspace(-80, 80, n_sources) + # Two sources deliberately above the monitored/pointing thresholds. + data["fluxJy"] = np.full(n_sources, 0.001) + data["fluxJy"][0] = 0.025 # 25 mJy -- above monitored (20 mJy) + data["fluxJy"][1] = 0.5 # 500 mJy -- above pointing (300 mJy) + data["name"] = np.array([f"Source_{i}" for i in range(n_sources)], dtype="S20") + + hdu = fits.BinTableHDU(data) + fits_path = tmp_path / "act_catalog.fits" + hdu.writeto(fits_path) + return fits_path + + +@pytest.fixture +def jpl_ephem_parquet(tmp_path): + """Two objects, 5 two-hour-cadence points each, written as one row + group per designation -- matching the real JPL Horizons + batched-ephemeris parquet files (verified against production data: + each file has exactly one row group per object), which is the + precondition collect_sso_designations()/load_ephemerides() rely on + to avoid reading full ephemeris columns just to enumerate objects. + A plain single-shot df.to_parquet() would NOT reproduce this (pandas + writes one row group for the whole frame), so each object's rows are + written as a separate row group explicitly. + """ + start = Time("2025-01-01T00:00:00") + parquet_path = tmp_path / "ephem.parquet" + writer = None + try: + for designation, ra0 in [("1 Ceres", 10.0), ("2 Pallas", 200.0)]: + rows = [] + for i in range(5): + t = start + i * 2 * u.hour + rows.append( + { + "designation": designation, + "datetime_utc": t.datetime, + "julian_day": t.jd, + "ra_deg": ra0 + i * 0.01, + "dec_deg": 5.0 + i * 0.01, + "distance_au": 2.5, + } + ) + table = pa.Table.from_pandas(pd.DataFrame(rows), preserve_index=False) + if writer is None: + writer = pq.ParquetWriter(parquet_path, table.schema) + writer.write_table(table) + finally: + if writer is not None: + writer.close() + return parquet_path + + +def test_create_schema_has_expected_indexes(tmp_path): + db_path = tmp_path / "socat.db" + bulk_sqlite.create_schema(db_path) + + conn = sqlite3.connect(str(db_path)) + index_names = { + row[0] + for row in conn.execute( + "SELECT name FROM sqlite_master WHERE type='index' AND tbl_name='moving_sources'" + ) + } + conn.close() + assert "idx_moving_sources_sso_time" in index_names + assert "ix_moving_sources_time" in index_names + + +def test_load_fixed_sources(tmp_path, act_fits_catalog): + db_path = tmp_path / "socat.db" + bulk_sqlite.create_schema(db_path) + conn = sqlite3.connect(str(db_path)) + bulk_sqlite._tune_for_bulk_load(conn) + + n_loaded = bulk_sqlite.load_fixed_sources(conn, act_fits_catalog) + assert n_loaded == 10 + + rows = conn.execute( + "SELECT name, flux_mJy, monitored, pointing FROM fixed_sources ORDER BY name" + ).fetchall() + conn.close() + assert len(rows) == 10 + + by_name = {r[0]: r for r in rows} + assert by_name["Source_0"][1] == pytest.approx(25.0) + assert by_name["Source_0"][2] == 1 # monitored + assert by_name["Source_0"][3] == 0 # not pointing + assert by_name["Source_1"][1] == pytest.approx(500.0) + assert by_name["Source_1"][3] == 1 # pointing + assert by_name["Source_2"][2] == 0 + assert by_name["Source_2"][3] == 0 + + +def test_parse_designation(): + assert bulk_sqlite._parse_designation("1 Ceres") == (1, "Ceres") + assert bulk_sqlite._parse_designation("7335 (1989 JA)") == (7335, "(1989 JA)") + + +def test_collect_and_load_sso_and_ephemerides(tmp_path, jpl_ephem_parquet): + db_path = tmp_path / "socat.db" + bulk_sqlite.create_schema(db_path) + conn = sqlite3.connect(str(db_path)) + bulk_sqlite._tune_for_bulk_load(conn) + + designations = bulk_sqlite.collect_sso_designations([jpl_ephem_parquet]) + assert designations == {"1 Ceres": (1, "Ceres"), "2 Pallas": (2, "Pallas")} + + sso_ids = bulk_sqlite.load_solar_system_objects(conn, designations) + assert set(sso_ids) == {"1 Ceres", "2 Pallas"} + + n_ephem = bulk_sqlite.load_ephemerides( + conn, [jpl_ephem_parquet], sso_ids, designations + ) + assert n_ephem == 10 + + ceres_rows = conn.execute( + "SELECT ra_deg, dec_deg, time FROM moving_sources WHERE sso_id = ? ORDER BY time", + (sso_ids["1 Ceres"],), + ).fetchall() + conn.close() + assert len(ceres_rows) == 5 + assert ceres_rows[0][0] == pytest.approx(10.0) + assert ceres_rows[0][2] == "2025-01-01 00:00:00.000000" + assert ceres_rows[-1][0] == pytest.approx(10.04) + + +def test_build_end_to_end(tmp_path, act_fits_catalog, jpl_ephem_parquet): + output_path = tmp_path / "output" / "socat.db" + build_dir = tmp_path / "build" + + summary = bulk_sqlite.build( + fits_path=act_fits_catalog, + ephem_paths=[jpl_ephem_parquet], + output_path=output_path, + build_dir=build_dir, + ) + assert summary["n_fixed_sources"] == 10 + assert summary["n_sso"] == 2 + assert summary["n_ephem"] == 10 + assert output_path.exists() + # build_dir should be cleaned up after the copy. + assert not (build_dir / output_path.name).exists() + + conn = sqlite3.connect(str(output_path)) + assert conn.execute("PRAGMA foreign_key_check").fetchall() == [] + conn.close() + + client = Client(db_url=f"sqlite:///{output_path}") + fixed = client.get_box_fixed( + lower_left=ICRS(ra=0.0 * u.deg, dec=-89.9 * u.deg), + upper_right=ICRS(ra=359.999 * u.deg, dec=89.9 * u.deg), + ) + assert len(fixed) == 10 + + sso = client.get_box_sso( + lower_left=ICRS(ra=0.0 * u.deg, dec=-89.9 * u.deg), + upper_right=ICRS(ra=359.999 * u.deg, dec=89.9 * u.deg), + t_min=Time("2025-01-01T00:00:00"), + t_max=Time("2025-01-01T10:00:00"), + ) + assert sorted(s.name for s in sso) == ["Ceres", "Pallas"] diff --git a/tests/test_db_client.py b/tests/test_db_client.py index 200c030..87fb004 100644 --- a/tests/test_db_client.py +++ b/tests/test_db_client.py @@ -362,6 +362,107 @@ def test_box(db_client): client.delete_source(source_id=source_2.source_id) +def test_box_ra_wraparound_dedup(db_client): + """ + get_box_sso/get_box_fixed take a different code path when the box + wraps across RA=0/360 (lower_left.ra > upper_right.ra), splitting the + query into two RA sub-ranges combined with a union. A source with + several matching ephemeris points -- or points landing in both + sub-ranges -- must still come back as exactly one row. + + db_client is session-scoped and shared across this module's tests, so + this uses a dec band (40-43) that no other test in this file touches, + to avoid colliding with fixtures other tests leave behind. + """ + client = db_client + + t_min = Time("2025-01-01T00:00:00.00") + t_max = t_min + 10 * u.h + + fixed_in_box = client.create_source( + position=ICRS(355.0 * u.deg, 41.0 * u.deg), + name="wrap-fixed", + flux=1.0 * u.mJy, + ) + fixed_out_of_box = client.create_source( + position=ICRS(180.0 * u.deg, 41.0 * u.deg), + name="not-wrapped-fixed", + flux=1.0 * u.mJy, + ) + + # Several ephem points, all within the "right" (350-360) sub-range -- + # exercises per-branch dedup. + right_only = client.create_sso(name="wrap-right-only", MPC_id=9001) + for i in range(4): + client.create_ephem( + sso_id=right_only.sso_id, + MPC_id=right_only.MPC_id, + name=right_only.name, + time=t_min + i * u.h, + position=ICRS((355.0 + i) * u.deg, 41.0 * u.deg), + flux=1.0 * u.mJy, + ) + + # Ephem points on both sides of RA=0 within the window -- exercises + # cross-branch dedup (union vs union_all). + crosses_zero = client.create_sso(name="wrap-crosses-zero", MPC_id=9002) + for i, ra in enumerate([356.0, 358.0, 1.0, 3.0]): + client.create_ephem( + sso_id=crosses_zero.sso_id, + MPC_id=crosses_zero.MPC_id, + name=crosses_zero.name, + time=t_min + i * u.h, + position=ICRS(ra * u.deg, 41.0 * u.deg), + flux=1.0 * u.mJy, + ) + + # Outside the RA range entirely. + not_in_box = client.create_sso(name="wrap-not-in-box", MPC_id=9003) + client.create_ephem( + sso_id=not_in_box.sso_id, + MPC_id=not_in_box.MPC_id, + name=not_in_box.name, + time=t_min, + position=ICRS(180.0 * u.deg, 41.0 * u.deg), + flux=1.0 * u.mJy, + ) + + lower_left = ICRS(350.0 * u.deg, 40.0 * u.deg) + upper_right = ICRS(10.0 * u.deg, 43.0 * u.deg) + + fixed_matches = client.get_box_fixed(lower_left=lower_left, upper_right=upper_right) + assert [s.name for s in fixed_matches] == ["wrap-fixed"] + + sso_matches = client.get_box_sso( + lower_left=lower_left, + upper_right=upper_right, + t_min=t_min, + t_max=t_max, + ) + assert sorted(s.name for s in sso_matches) == [ + "wrap-crosses-zero", + "wrap-right-only", + ] + + source_gens = client.get_box( + lower_left=lower_left, + upper_right=upper_right, + t_min=t_min, + t_max=t_max, + ) + assert sorted(g.source.name for g in source_gens) == [ + "wrap-crosses-zero", + "wrap-fixed", + "wrap-right-only", + ] + + client.delete_sso(sso_id=right_only.sso_id) + client.delete_sso(sso_id=crosses_zero.sso_id) + client.delete_sso(sso_id=not_in_box.sso_id) + client.delete_source(source_id=fixed_in_box.source_id) + client.delete_source(source_id=fixed_out_of_box.source_id) + + def test_not_found_behavior(db_client): client = db_client