From 7a4ca8a946efb18b4871be06a9307c3f93fc4e0f Mon Sep 17 00:00:00 2001 From: jakeross Date: Thu, 15 Jan 2026 19:02:53 -0700 Subject: [PATCH 01/10] feat: add NMA_Stratigraphy table and transfer functionality --- ...c3b4a5e67_create_nma_stratigraphy_table.py | 71 ++++++++ db/nma_legacy.py | 35 +++- db/thing.py | 18 +- tests/test_stratigraphy_legacy.py | 93 ++++++++++ transfers/metrics.py | 13 +- transfers/stratigraphy_legacy.py | 163 ++++++++++++++++++ transfers/transfer.py | 28 ++- 7 files changed, 403 insertions(+), 18 deletions(-) create mode 100644 alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py create mode 100644 tests/test_stratigraphy_legacy.py create mode 100644 transfers/stratigraphy_legacy.py diff --git a/alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py b/alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py new file mode 100644 index 000000000..231a67dd4 --- /dev/null +++ b/alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py @@ -0,0 +1,71 @@ +"""Create legacy NMA_Stratigraphy table. + +Revision ID: 1d2c3b4a5e67 +Revises: a7b8c9d0e1f2 +Create Date: 2026-01-15 00:00:00.000000 +""" + +from typing import Sequence, Union + +import sqlalchemy as sa +from alembic import op +from sqlalchemy import inspect +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision: str = "1d2c3b4a5e67" +down_revision: Union[str, Sequence[str], None] = "a7b8c9d0e1f2" +branch_labels: Union[str, Sequence[str], None] = None +depends_on: Union[str, Sequence[str], None] = None + + +def upgrade() -> None: + """Create the legacy stratigraphy table.""" + bind = op.get_bind() + inspector = inspect(bind) + if inspector.has_table("NMA_Stratigraphy"): + return + + op.create_table( + "NMA_Stratigraphy", + sa.Column( + "GlobalID", + postgresql.UUID(as_uuid=True), + primary_key=True, + nullable=False, + ), + sa.Column("WellID", postgresql.UUID(as_uuid=True), nullable=True), + sa.Column("PointID", sa.String(length=10), nullable=False), + sa.Column( + "thing_id", + sa.Integer(), + sa.ForeignKey("thing.id", ondelete="CASCADE"), + nullable=False, + ), + sa.Column("StratTop", sa.Float(), nullable=True), + sa.Column("StratBottom", sa.Float(), nullable=True), + sa.Column("UnitIdentifier", sa.String(length=50), nullable=True), + sa.Column("Lithology", sa.String(length=100), nullable=True), + sa.Column("LithologicModifier", sa.String(length=100), nullable=True), + sa.Column("ContributingUnit", sa.String(length=10), nullable=True), + sa.Column("StratSource", sa.Text(), nullable=True), + sa.Column("StratNotes", sa.Text(), nullable=True), + sa.Column("OBJECTID", sa.Integer(), nullable=True, unique=True), + sa.ForeignKeyConstraint(["thing_id"], ["thing.id"], ondelete="CASCADE"), + ) + op.create_index( + "ix_nma_stratigraphy_point_id", + "NMA_Stratigraphy", + ["PointID"], + ) + op.create_index( + "ix_nma_stratigraphy_thing_id", + "NMA_Stratigraphy", + ["thing_id"], + ) + + +def downgrade() -> None: + op.drop_index("ix_nma_stratigraphy_thing_id", table_name="NMA_Stratigraphy") + op.drop_index("ix_nma_stratigraphy_point_id", table_name="NMA_Stratigraphy") + op.drop_table("NMA_Stratigraphy") diff --git a/db/nma_legacy.py b/db/nma_legacy.py index d36ee0f28..b76fad1e4 100644 --- a/db/nma_legacy.py +++ b/db/nma_legacy.py @@ -17,10 +17,10 @@ """Legacy NM Aquifer models copied from AMPAPI.""" import uuid - from datetime import date, datetime from typing import TYPE_CHECKING, List, Optional +from db.base import Base from sqlalchemy import ( Boolean, Date, @@ -37,8 +37,6 @@ from sqlalchemy.dialects.postgresql import UUID from sqlalchemy.orm import Mapped, mapped_column, relationship, validates -from db.base import Base - if TYPE_CHECKING: from db.thing import Thing @@ -206,6 +204,37 @@ class NMAHydraulicsData(Base): thing: Mapped["Thing"] = relationship("Thing") +class Stratigraphy(Base): + """Legacy stratigraphy (lithology log) data from AMPAPI.""" + + __tablename__ = "NMA_Stratigraphy" + + global_id: Mapped[uuid.UUID] = mapped_column( + "GlobalID", UUID(as_uuid=True), primary_key=True + ) + well_id: Mapped[Optional[uuid.UUID]] = mapped_column("WellID", UUID(as_uuid=True)) + point_id: Mapped[str] = mapped_column("PointID", String(10), nullable=False) + thing_id: Mapped[int] = mapped_column( + Integer, ForeignKey("thing.id", ondelete="CASCADE"), nullable=False + ) + + strat_top: Mapped[Optional[float]] = mapped_column("StratTop", Float) + strat_bottom: Mapped[Optional[float]] = mapped_column("StratBottom", Float) + unit_identifier: Mapped[Optional[str]] = mapped_column("UnitIdentifier", String(50)) + lithology: Mapped[Optional[str]] = mapped_column("Lithology", String(100)) + lithologic_modifier: Mapped[Optional[str]] = mapped_column( + "LithologicModifier", String(100) + ) + contributing_unit: Mapped[Optional[str]] = mapped_column( + "ContributingUnit", String(10) + ) + strat_source: Mapped[Optional[str]] = mapped_column("StratSource", Text) + strat_notes: Mapped[Optional[str]] = mapped_column("StratNotes", Text) + object_id: Mapped[Optional[int]] = mapped_column("OBJECTID", Integer, unique=True) + + thing: Mapped["Thing"] = relationship("Thing", back_populates="stratigraphy_logs") + + class ChemistrySampleInfo(Base): """ Legacy Chemistry SampleInfo table from AMPAPI. diff --git a/db/thing.py b/db/thing.py index 8283bdc4f..0ea31abfc 100644 --- a/db/thing.py +++ b/db/thing.py @@ -16,11 +16,6 @@ from datetime import date from typing import List, TYPE_CHECKING -from sqlalchemy import Integer, ForeignKey, String, Column, Float, Date -from sqlalchemy.ext.associationproxy import association_proxy, AssociationProxy -from sqlalchemy.orm import relationship, mapped_column, Mapped -from sqlalchemy_utils import TSVectorType - from db import lexicon_term, NotesMixin from db.asset import Asset from db.base import ( @@ -33,6 +28,10 @@ from db.permission_history import PermissionHistoryMixin from db.status_history import StatusHistoryMixin from services.util import retrieve_latest_polymorphic_history_table_record +from sqlalchemy import Integer, ForeignKey, String, Column, Float, Date +from sqlalchemy.ext.associationproxy import association_proxy, AssociationProxy +from sqlalchemy.orm import relationship, mapped_column, Mapped +from sqlalchemy_utils import TSVectorType if TYPE_CHECKING: from db.location import Location @@ -47,7 +46,7 @@ from db.thing_geologic_formation_association import ( ThingGeologicFormationAssociation, ) - from db.nma_legacy import ChemistrySampleInfo + from db.nma_legacy import ChemistrySampleInfo, Stratigraphy class Thing( @@ -312,6 +311,13 @@ class Thing( passive_deletes=True, ) + stratigraphy_logs: Mapped[List["Stratigraphy"]] = relationship( + "Stratigraphy", + back_populates="thing", + cascade="all, delete-orphan", + passive_deletes=True, + ) + # --- Association Proxies --- assets: AssociationProxy[list["Asset"]] = association_proxy( "asset_associations", "asset" diff --git a/tests/test_stratigraphy_legacy.py b/tests/test_stratigraphy_legacy.py new file mode 100644 index 000000000..4693eac64 --- /dev/null +++ b/tests/test_stratigraphy_legacy.py @@ -0,0 +1,93 @@ +import uuid + +import pandas as pd +from db import Stratigraphy, Thing +from db.engine import session_ctx +from transfers.stratigraphy_legacy import StratigraphyLegacyTransferer + + +def _create_test_thing(name: str = "ST-0001") -> Thing: + with session_ctx() as session: + thing = Thing(name=name, thing_type="water well", release_status="draft") + session.add(thing) + session.commit() + session.refresh(thing) + return thing + + +def test_nma_stratigraphy_model_relationship(): + thing = _create_test_thing("ST-REL-001") + record_id = uuid.uuid4() + with session_ctx() as session: + record = Stratigraphy( + global_id=record_id, + well_id=uuid.uuid4(), + point_id=thing.name, + thing_id=thing.id, + strat_top=0.0, + strat_bottom=10.5, + unit_identifier="110ALVM", + object_id=123, + ) + session.add(record) + session.commit() + + fetched = session.get(Stratigraphy, record_id) + assert fetched is not None + assert fetched.thing_id == thing.id + assert fetched.thing.name == thing.name + assert fetched.strat_bottom == 10.5 + + +def test_stratigraphy_transfer_inserts_rows(monkeypatch): + thing = _create_test_thing("ST-TR-001") + + data = pd.DataFrame( + [ + { + "GlobalID": str(uuid.uuid4()), + "WellID": str(uuid.uuid4()), + "PointID": thing.name, + "StratTop": 0, + "StratBottom": 15.25, + "UnitIdentifier": "110ALVM", + "Lithology": "Alluvium", + "LithologicModifier": "sandy", + "ContributingUnit": "S", + "StratSource": "Test source", + "StratNotes": "Test note", + "OBJECTID": 999, + }, + { + "GlobalID": "not-a-uuid", + "WellID": None, + "PointID": thing.name, + "StratTop": 5, + "StratBottom": 10, + "UnitIdentifier": None, + "Lithology": None, + "LithologicModifier": None, + "ContributingUnit": None, + "StratSource": None, + "StratNotes": None, + "OBJECTID": 1000, + }, + ] + ) + + monkeypatch.setattr( + "transfers.stratigraphy_legacy.read_csv", lambda _: data.copy(deep=True) + ) + monkeypatch.setattr("transfers.stratigraphy_legacy.replace_nans", lambda df: df) + + transferer = StratigraphyLegacyTransferer(batch_size=10) + transferer.transfer() + + with session_ctx() as session: + rows = session.query(Stratigraphy).filter_by(point_id=thing.name).all() + assert len(rows) == 1 + record = rows[0] + assert record.point_id == thing.name + assert record.thing_id == thing.id + assert record.strat_bottom == 15.25 + assert record.object_id == 999 diff --git a/transfers/metrics.py b/transfers/metrics.py index cf50644a1..cdd1b93d4 100644 --- a/transfers/metrics.py +++ b/transfers/metrics.py @@ -18,11 +18,6 @@ from datetime import datetime from pathlib import Path -from pandas import DataFrame -from pydantic import ValidationError -from sqlalchemy import select, func -from sqlalchemy.exc import ProgrammingError - from db import ( Thing, WellScreen, @@ -40,6 +35,7 @@ NMAHydraulicsData, NMARadionuclides, NMAMajorChemistry, + Stratigraphy, SurfaceWaterData, NMAWaterLevelsContinuousPressureDaily, ViewNGWMNWellConstruction, @@ -49,7 +45,11 @@ NMAMinorTraceChemistry, ) from db.engine import session_ctx +from pandas import DataFrame +from pydantic import ValidationError from services.gcs_helper import get_storage_bucket +from sqlalchemy import select, func +from sqlalchemy.exc import ProgrammingError class Metrics: @@ -155,6 +155,9 @@ def permissions_metrics(self, *args, **kw) -> None: def stratigraphy_metrics(self, *args, **kw) -> None: self._handle_metrics(ThingGeologicFormationAssociation, *args, **kw) + def nma_stratigraphy_metrics(self, *args, **kw) -> None: + self._handle_metrics(Stratigraphy, name="NMA_Stratigraphy", *args, **kw) + def minor_trace_chemistry_metrics(self, *args, **kw) -> None: self._handle_metrics( NMAMinorTraceChemistry, name="MinorTraceChemistry", *args, **kw diff --git a/transfers/stratigraphy_legacy.py b/transfers/stratigraphy_legacy.py new file mode 100644 index 000000000..029879456 --- /dev/null +++ b/transfers/stratigraphy_legacy.py @@ -0,0 +1,163 @@ +"""Transfer Stratigraphy.csv into the NMA_Stratigraphy legacy table.""" + +from __future__ import annotations + +from typing import Any, Dict, List +from uuid import UUID + +import pandas as pd +from db import Stratigraphy, Thing +from sqlalchemy.dialects.postgresql import insert +from sqlalchemy.orm import Session +from transfers.logger import logger +from transfers.transferer import Transferer +from transfers.util import ( + filter_to_valid_point_ids, + read_csv, + replace_nans, +) + + +class StratigraphyLegacyTransferer(Transferer): + """Imports Stratigraphy.csv rows into NMA_Stratigraphy.""" + + source_table = "NMA_Stratigraphy" + + def __init__(self, batch_size: int = 1000, *args, **kwargs) -> None: + super().__init__(*args, **kwargs) + self.batch_size = batch_size + self._thing_id_cache: dict[str, int] = {} + + def _get_dfs(self): # type: ignore[override] + df = read_csv("Stratigraphy") + cleaned = replace_nans(df) + cleaned = filter_to_valid_point_ids(cleaned, self.pointids) + return df, cleaned + + def _transfer_hook(self, session: Session) -> None: # type: ignore[override] + if self.cleaned_df is None or self.cleaned_df.empty: + logger.info("No Stratigraphy rows to import after cleaning") + return + + self._prime_thing_cache(session) + rows: List[Dict[str, Any]] = [] + for row in self.cleaned_df.itertuples(): + record = self._row_dict(row) + if record is None: + continue + rows.append(record) + + if not rows: + logger.warning("All Stratigraphy rows were skipped during processing") + return + + insert_stmt = insert(Stratigraphy) + excluded = insert_stmt.excluded + + for start in range(0, len(rows), self.batch_size): + chunk = rows[start : start + self.batch_size] + logger.info( + "Upserting Stratigraphy batch %s-%s (%s rows)", + start, + start + len(chunk) - 1, + len(chunk), + ) + stmt = insert_stmt.values(chunk).on_conflict_do_update( + index_elements=["GlobalID"], + set_={ + "WellID": excluded.WellID, + "PointID": excluded.PointID, + "thing_id": excluded.thing_id, + "StratTop": excluded.StratTop, + "StratBottom": excluded.StratBottom, + "UnitIdentifier": excluded.UnitIdentifier, + "Lithology": excluded.Lithology, + "LithologicModifier": excluded.LithologicModifier, + "ContributingUnit": excluded.ContributingUnit, + "StratSource": excluded.StratSource, + "StratNotes": excluded.StratNotes, + "OBJECTID": excluded.OBJECTID, + }, + ) + session.execute(stmt) + session.commit() + session.expunge_all() + + def _prime_thing_cache(self, session: Session) -> None: + point_ids = set(self.cleaned_df["PointID"].dropna()) # type: ignore[index] + if not point_ids: + self._thing_id_cache = {} + return + results = ( + session.query(Thing.id, Thing.name).filter(Thing.name.in_(point_ids)).all() + ) + self._thing_id_cache = {name: thing_id for thing_id, name in results} + + def _row_dict(self, row: pd.Series) -> Dict[str, Any] | None: + point_id = getattr(row, "PointID", None) + if not point_id: + self._capture_error("", "Missing PointID", "PointID") + return None + thing_id = self._thing_id_cache.get(point_id) + if not thing_id: + self._capture_error(point_id, "No Thing found for PointID", "thing_id") + return None + + global_id = self._uuid_value(getattr(row, "GlobalID", None)) + if global_id is None: + self._capture_error(point_id, "Invalid GlobalID", "GlobalID") + return None + + return { + "GlobalID": global_id, + "WellID": self._uuid_value(getattr(row, "WellID", None)), + "PointID": point_id, + "thing_id": thing_id, + "StratTop": self._float_value(getattr(row, "StratTop", None)), + "StratBottom": self._float_value(getattr(row, "StratBottom", None)), + "UnitIdentifier": self._string_value(getattr(row, "UnitIdentifier", None)), + "Lithology": self._string_value(getattr(row, "Lithology", None)), + "LithologicModifier": self._string_value( + getattr(row, "LithologicModifier", None) + ), + "ContributingUnit": self._string_value( + getattr(row, "ContributingUnit", None) + ), + "StratSource": self._string_value(getattr(row, "StratSource", None)), + "StratNotes": self._string_value(getattr(row, "StratNotes", None)), + "OBJECTID": self._int_value(getattr(row, "OBJECTID", None)), + } + + def _uuid_value(self, value: Any) -> UUID | None: + if value in (None, ""): + return None + if isinstance(value, UUID): + return value + try: + return UUID(str(value)) + except (ValueError, TypeError): + return None + + def _float_value(self, value: Any) -> float | None: + if value in (None, ""): + return None + try: + return float(value) + except (TypeError, ValueError): + return None + + def _int_value(self, value: Any) -> int | None: + if value in (None, ""): + return None + try: + return int(value) + except (TypeError, ValueError): + return None + + def _string_value(self, value: Any) -> str | None: + if value is None: + return None + if isinstance(value, str): + trimmed = value.strip() + return trimmed or None + return str(value) diff --git a/transfers/transfer.py b/transfers/transfer.py index 0586e62bc..8342ad266 100644 --- a/transfers/transfer.py +++ b/transfers/transfer.py @@ -17,17 +17,16 @@ import time from concurrent.futures import ThreadPoolExecutor, as_completed -from dotenv import load_dotenv - from alembic import command from alembic.config import Config - from db.engine import session_ctx -from sqlalchemy import text +from dotenv import load_dotenv from services.util import get_bool_env +from sqlalchemy import text from transfers.aquifer_system_transfer import transfer_aquifer_systems from transfers.geologic_formation_transfer import transfer_geologic_formations from transfers.permissions_transfer import transfer_permissions +from transfers.stratigraphy_legacy import StratigraphyLegacyTransferer from transfers.stratigraphy_transfer import transfer_stratigraphy load_dotenv() @@ -238,6 +237,7 @@ def transfer_all(metrics, limit=100): transfer_minor_trace_chemistry = get_bool_env( "TRANSFER_MINOR_TRACE_CHEMISTRY", True ) + transfer_nma_stratigraphy = get_bool_env("TRANSFER_NMA_STRATIGRAPHY", True) use_parallel = get_bool_env("TRANSFER_PARALLEL", True) if use_parallel: @@ -263,6 +263,7 @@ def transfer_all(metrics, limit=100): transfer_pressure_daily, transfer_weather_data, transfer_minor_trace_chemistry, + transfer_nma_stratigraphy, ) else: _transfer_sequential( @@ -287,6 +288,7 @@ def transfer_all(metrics, limit=100): transfer_pressure_daily, transfer_weather_data, transfer_minor_trace_chemistry, + transfer_nma_stratigraphy, ) @@ -312,6 +314,7 @@ def _transfer_parallel( transfer_pressure_daily, transfer_weather_data, transfer_minor_trace_chemistry, + transfer_nma_stratigraphy, ): """Execute transfers in parallel where possible.""" message("PARALLEL TRANSFER GROUP 1") @@ -376,6 +379,15 @@ def _transfer_parallel( futures[future] = name # Submit session-based transfers + if transfer_nma_stratigraphy: + future = executor.submit( + _execute_transfer_with_timing, + "NMAStratigraphy", + StratigraphyLegacyTransferer, + flags, + ) + futures[future] = "NMAStratigraphy" + future = executor.submit( _execute_session_transfer_with_timing, "Stratigraphy", @@ -404,6 +416,8 @@ def _transfer_parallel( metrics.contact_metrics(*results_map["Contacts"]) if "Stratigraphy" in results_map and results_map["Stratigraphy"]: metrics.stratigraphy_metrics(*results_map["Stratigraphy"]) + if "NMAStratigraphy" in results_map and results_map["NMAStratigraphy"]: + metrics.nma_stratigraphy_metrics(*results_map["NMAStratigraphy"]) if "WaterLevels" in results_map and results_map["WaterLevels"]: metrics.water_level_metrics(*results_map["WaterLevels"]) if "LinkIdsWellData" in results_map and results_map["LinkIdsWellData"]: @@ -521,6 +535,7 @@ def _transfer_sequential( transfer_pressure_daily, transfer_weather_data, transfer_minor_trace_chemistry, + transfer_nma_stratigraphy, ): """Original sequential transfer logic.""" if transfer_screens: @@ -542,6 +557,11 @@ def _transfer_sequential( with session_ctx() as session: transfer_permissions(session) + if transfer_nma_stratigraphy: + message("TRANSFERRING NMA STRATIGRAPHY") + results = _execute_transfer(StratigraphyLegacyTransferer, flags=flags) + metrics.nma_stratigraphy_metrics(*results) + message("TRANSFERRING STRATIGRAPHY") with session_ctx() as session: results = transfer_stratigraphy(session, limit=limit) From e834aadd958ff2cdb235778f6f310452f173e8ad Mon Sep 17 00:00:00 2001 From: Jake Ross Date: Thu, 15 Jan 2026 21:52:52 -0700 Subject: [PATCH 02/10] Update alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py | 1 - 1 file changed, 1 deletion(-) diff --git a/alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py b/alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py index 231a67dd4..ca742a372 100644 --- a/alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py +++ b/alembic/versions/1d2c3b4a5e67_create_nma_stratigraphy_table.py @@ -51,7 +51,6 @@ def upgrade() -> None: sa.Column("StratSource", sa.Text(), nullable=True), sa.Column("StratNotes", sa.Text(), nullable=True), sa.Column("OBJECTID", sa.Integer(), nullable=True, unique=True), - sa.ForeignKeyConstraint(["thing_id"], ["thing.id"], ondelete="CASCADE"), ) op.create_index( "ix_nma_stratigraphy_point_id", From 47bb8af08a11371f12f2860a9a87d0633b3242ec Mon Sep 17 00:00:00 2001 From: Jake Ross Date: Thu, 15 Jan 2026 21:53:54 -0700 Subject: [PATCH 03/10] Update transfers/stratigraphy_legacy.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- transfers/stratigraphy_legacy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transfers/stratigraphy_legacy.py b/transfers/stratigraphy_legacy.py index 029879456..63ec13ca0 100644 --- a/transfers/stratigraphy_legacy.py +++ b/transfers/stratigraphy_legacy.py @@ -93,7 +93,7 @@ def _prime_thing_cache(self, session: Session) -> None: ) self._thing_id_cache = {name: thing_id for thing_id, name in results} - def _row_dict(self, row: pd.Series) -> Dict[str, Any] | None: + def _row_dict(self, row: Any) -> Dict[str, Any] | None: point_id = getattr(row, "PointID", None) if not point_id: self._capture_error("", "Missing PointID", "PointID") From 06432fb0a356a617e281acbd38648a749d174227 Mon Sep 17 00:00:00 2001 From: jakeross Date: Thu, 15 Jan 2026 22:00:57 -0700 Subject: [PATCH 04/10] fix: handle None data case in populate_fields method --- schemas/location.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/schemas/location.py b/schemas/location.py index 1a61e257a..e0f38e2fb 100644 --- a/schemas/location.py +++ b/schemas/location.py @@ -17,12 +17,11 @@ from typing import Any from typing import List +from core.constants import SRID_WGS84, SRID_UTM_ZONE_13N +from core.enums import ElevationMethod, CoordinateMethod from geoalchemy2 import WKBElement from geoalchemy2.shape import to_shape from pydantic import BaseModel, model_validator, field_validator, Field, ConfigDict - -from core.constants import SRID_WGS84, SRID_UTM_ZONE_13N -from core.enums import ElevationMethod, CoordinateMethod from schemas import BaseCreateModel, BaseUpdateModel, BaseResponseModel from schemas.notes import NoteResponse, CreateNote, UpdateNote from services.util import convert_m_to_ft, transform_srid @@ -130,6 +129,8 @@ class LocationGeoJSONResponse(BaseModel): @model_validator(mode="before") @classmethod def populate_fields(cls, data: Any) -> Any: + if data is None: + return None # convert row to dictionary if not isinstance(data, dict): data_dict = {c.name: getattr(data, c.name) for c in data.__table__.columns} From adcdeb00a7eeedea44912a4231ff6850c7d5c88f Mon Sep 17 00:00:00 2001 From: jakeross Date: Thu, 15 Jan 2026 22:47:38 -0700 Subject: [PATCH 05/10] feat: enhance test setup with location and measuring point creation for wells --- db/thing.py | 4 ++ tests/conftest.py | 55 +++++++++++++-- tests/test_geospatial.py | 101 ++++++++++++++-------------- tests/test_location.py | 27 ++++---- tests/test_nma_chemistry_lineage.py | 77 +++++++++++---------- tests/test_stratigraphy_legacy.py | 34 +++++++++- tests/test_thing.py | 9 ++- 7 files changed, 196 insertions(+), 111 deletions(-) diff --git a/db/thing.py b/db/thing.py index 0ea31abfc..7a3b620ba 100644 --- a/db/thing.py +++ b/db/thing.py @@ -423,6 +423,8 @@ def measuring_point_height(self) -> int | None: Since measuring_point_history is eagerly loaded, this should not introduce N+1 query issues. """ if self.thing_type == "water well": + if not self.measuring_points: + return None sorted_measuring_point_history = sorted( self.measuring_points, key=lambda x: x.start_date, reverse=True ) @@ -439,6 +441,8 @@ def measuring_point_description(self) -> str | None: Since measuring_point_history is eagerly loaded, this should not introduce N+1 query issues. """ if self.thing_type == "water well": + if not self.measuring_points: + return None sorted_measuring_point_history = sorted( self.measuring_points, key=lambda x: x.start_date, reverse=True ) diff --git a/tests/conftest.py b/tests/conftest.py index 979778cb3..d1aa28842 100644 --- a/tests/conftest.py +++ b/tests/conftest.py @@ -1,23 +1,19 @@ import os -import uuid import pytest from alembic import command from alembic.config import Config -from dotenv import load_dotenv -from sqlalchemy import text -from sqlalchemy_utils import TSVectorType -from sqlalchemy_searchable import sync_trigger - from core.initializers import init_lexicon, init_parameter from db import * from db.engine import session_ctx +from dotenv import load_dotenv +from sqlalchemy_searchable import sync_trigger from tests import get_parameter_id def pytest_configure(): load_dotenv(override=True) - os.environ.setdefault("POSTGRES_PORT", "54321") + os.environ["POSTGRES_PORT"] = "54321" def _alembic_config() -> Config: @@ -72,9 +68,54 @@ def _setup_test_db(): _sync_search_vectors() init_lexicon() init_parameter() + _ensure_locations_for_things() + _ensure_measuring_points_for_wells() yield +def _ensure_locations_for_things() -> None: + with session_ctx() as session: + things = session.query(Thing).all() + for thing in things: + if thing.location_associations: + continue + location = Location( + point="POINT(-107.0 33.0)", + elevation=0, + release_status=thing.release_status, + ) + session.add(location) + session.flush() + session.add( + LocationThingAssociation( + location_id=location.id, + thing_id=thing.id, + effective_start=date.today().isoformat(), + ) + ) + session.commit() + + +def _ensure_measuring_points_for_wells() -> None: + with session_ctx() as session: + wells = session.query(Thing).filter(Thing.thing_type == "water well").all() + missing = [well for well in wells if not well.measuring_points] + for well in missing: + session.add( + MeasuringPointHistory( + thing_id=well.id, + measuring_point_height=well.well_casing_depth + or well.well_depth + or 0.0, + measuring_point_description="autogenerated for tests", + start_date=date.today(), + release_status=well.release_status, + ) + ) + if missing: + session.commit() + + @pytest.fixture() def location(): with session_ctx() as session: diff --git a/tests/test_geospatial.py b/tests/test_geospatial.py index a25ce24cb..bcb1482d7 100644 --- a/tests/test_geospatial.py +++ b/tests/test_geospatial.py @@ -13,11 +13,10 @@ # See the License for the specific language governing permissions and # limitations under the License. # =============================================================================== +from datetime import date from pathlib import Path import pytest -from geoalchemy2 import functions as geofunc - from core.constants import SRID_WGS84 from core.dependencies import ( admin_function, @@ -29,6 +28,7 @@ ) from db import Thing, Location, LocationThingAssociation, Group, MeasuringPointHistory from db.engine import session_ctx +from geoalchemy2 import functions as geofunc from main import app from tests import client, override_authentication @@ -68,46 +68,16 @@ def override_authentication_dependency_fixture(): @pytest.fixture(autouse=True, scope="module") def populate(): with session_ctx() as session: - # Create some sample data - thing1 = Thing(name="Thing 1", thing_type="water well") - thing2 = Thing(name="Thing 2", thing_type="water well") - session.add(thing1) - session.add(thing2) - - session.commit() - - mp_history_1 = MeasuringPointHistory( - thing_id=thing1.id, - measuring_point_height=5.0, - measuring_point_description="MP for Thing 1", - start_date="2023-01-01", - reason="Initial entry", - ) - mp_history_2 = MeasuringPointHistory( - thing_id=thing2.id, - measuring_point_height=10.0, - measuring_point_description="MP for Thing 2", - start_date="2023-01-01", - reason="Initial entry", - ) - session.add(mp_history_1) - session.add(mp_history_2) - - loc1 = Location( - # name="Test Location 1", - point=geofunc.ST_GeomFromText("POINT(10.1 10.1)", srid=SRID_WGS84), - elevation=0, + thing1 = _create_well_with_location( + session, + name="Thing 1", + point_wkt="POINT(10.1 10.1)", ) - loc2 = Location( - # name="Test Location 2", - point=geofunc.ST_GeomFromText("POINT(20 20)", srid=SRID_WGS84), - elevation=0, + thing2 = _create_well_with_location( + session, + name="Thing 2", + point_wkt="POINT(20 20)", ) - session.add(loc1) - session.add(loc2) - - session.add(LocationThingAssociation(location=loc1, thing=thing1)) - session.add(LocationThingAssociation(location=loc2, thing=thing2)) group = Group( name="Test Group Foo", @@ -117,15 +87,48 @@ def populate(): session.add(group) session.commit() - yield - - # Cleanup - session.delete(loc1) - session.delete(loc2) - session.delete(group) - session.delete(thing1) - session.delete(thing2) - session.commit() + try: + yield + finally: + session.delete(group) + session.delete(thing1) + session.delete(thing2) + session.commit() + + +def _create_well_with_location(session, name: str, point_wkt: str) -> Thing: + thing = Thing(name=name, thing_type="water well", release_status="draft") + session.add(thing) + session.commit() + session.refresh(thing) + + location = Location( + point=geofunc.ST_GeomFromText(point_wkt, srid=SRID_WGS84), + elevation=0, + release_status="draft", + ) + session.add(location) + session.flush() + session.add( + LocationThingAssociation( + location_id=location.id, + thing_id=thing.id, + effective_start=date.today().isoformat(), + ) + ) + session.add( + MeasuringPointHistory( + thing_id=thing.id, + measuring_point_height=5.0, + measuring_point_description=f"MP for {name}", + start_date=date.today(), + release_status="draft", + ) + ) + session.commit() + return thing + + # Cleanup def test_get_project_area(): diff --git a/tests/test_location.py b/tests/test_location.py index 31ab8d3c6..021fc7c8d 100644 --- a/tests/test_location.py +++ b/tests/test_location.py @@ -16,10 +16,9 @@ from datetime import timezone import pytest -from geoalchemy2.shape import to_shape - from core.dependencies import admin_function, editor_function, viewer_function from db import Location +from geoalchemy2.shape import to_shape from main import app from schemas import DT_FMT from tests import ( @@ -159,28 +158,28 @@ def test_get_locations(location): response = client.get("/location") assert response.status_code == 200 data = response.json() - assert data["total"] == 1 - assert data["items"][0]["id"] == location.id - assert data["items"][0]["created_at"] == location.created_at.astimezone( - timezone.utc - ).strftime(DT_FMT) + assert data["total"] >= 1 + item = next(item for item in data["items"] if item["id"] == location.id) + assert item["created_at"] == location.created_at.astimezone(timezone.utc).strftime( + DT_FMT + ) # assert data["items"][0]["name"] == location.name - assert isinstance(data["items"][0]["notes"], list) + assert isinstance(item["notes"], list) # If you know the exact number of notes expected: # assert len(data["items"][0]["notes"]) == expected_count # If you want to check content of a specific note: # if data["items"][0]["notes"]: # assert data["items"][0]["notes"][0]["content"] == expected_content - assert data["items"][0]["point"] == to_shape(location.point).wkt - assert data["items"][0]["elevation"] == location.elevation - assert data["items"][0]["release_status"] == location.release_status + assert item["point"] == to_shape(location.point).wkt + assert item["elevation"] == location.elevation + assert item["release_status"] == location.release_status # assert data["items"][0]["elevation_accuracy"] == location.elevation_accuracy # assert data["items"][0]["elevation_method"] == location.elevation_method # assert data["items"][0]["coordinate_accuracy"] == location.coordinate_accuracy # assert data["items"][0]["coordinate_method"] == location.coordinate_method - assert data["items"][0]["state"] == location.state - assert data["items"][0]["county"] == location.county - assert data["items"][0]["quad_name"] == location.quad_name + assert item["state"] == location.state + assert item["county"] == location.county + assert item["quad_name"] == location.quad_name def test_get_location_by_id(location): diff --git a/tests/test_nma_chemistry_lineage.py b/tests/test_nma_chemistry_lineage.py index b58edb911..47f28a900 100644 --- a/tests/test_nma_chemistry_lineage.py +++ b/tests/test_nma_chemistry_lineage.py @@ -27,10 +27,11 @@ - mtc.chemistry_sample_info.thing (full chain) """ +from datetime import date from uuid import uuid4 import pytest - +from db import Location, LocationThingAssociation, MeasuringPointHistory, Thing from db.engine import session_ctx @@ -57,14 +58,7 @@ def shared_well(): from db import Thing with session_ctx() as session: - thing = Thing( - name=f"Shared-Well-{uuid4().hex[:8]}", - thing_type="water well", - release_status="draft", - ) - session.add(thing) - session.commit() - session.refresh(thing) + thing = _create_well_with_defaults(session, f"Shared-Well-{uuid4().hex[:8]}") thing_id = thing.id yield thing_id @@ -195,18 +189,9 @@ def test_thing_has_chemistry_sample_infos_attribute(shared_well): def test_thing_chemistry_sample_infos_empty_by_default(): """New Thing should have empty chemistry_sample_infos.""" - from db import Thing with session_ctx() as session: - # Create a fresh Thing for this test - new_thing = Thing( - name=f"Empty-Test-{uuid4().hex[:8]}", - thing_type="water well", - release_status="draft", - ) - session.add(new_thing) - session.commit() - session.refresh(new_thing) + new_thing = _create_well_with_defaults(session, f"Empty-Test-{uuid4().hex[:8]}") assert new_thing.chemistry_sample_infos == [] @@ -549,17 +534,11 @@ def test_cascade_delete_sample_info_deletes_mtc(shared_well): def test_cascade_delete_thing_deletes_sample_infos(): """Deleting Thing should cascade delete its ChemistrySampleInfos.""" from db.nma_legacy import ChemistrySampleInfo - from db import Thing with session_ctx() as session: - # Create a separate thing for this test - test_thing = Thing( - name=f"Cascade-Test-{uuid4().hex[:8]}", - thing_type="water well", - release_status="draft", + test_thing = _create_well_with_defaults( + session, f"Cascade-Test-{uuid4().hex[:8]}" ) - session.add(test_thing) - session.commit() sample_info = ChemistrySampleInfo( object_id=_next_object_id(), @@ -588,17 +567,11 @@ def test_cascade_delete_thing_deletes_sample_infos(): def test_multiple_sample_infos_per_thing(): """Thing can have multiple ChemistrySampleInfos.""" from db.nma_legacy import ChemistrySampleInfo - from db import Thing with session_ctx() as session: - # Create a dedicated thing for this test - test_thing = Thing( - name=f"Multi-SI-Test-{uuid4().hex[:8]}", - thing_type="water well", - release_status="draft", + test_thing = _create_well_with_defaults( + session, f"Multi-SI-Test-{uuid4().hex[:8]}" ) - session.add(test_thing) - session.commit() for i in range(3): sample_info = ChemistrySampleInfo( @@ -654,4 +627,38 @@ def test_multiple_mtc_per_sample_info(shared_well): session.commit() +def _create_well_with_defaults(session, name: str) -> Thing: + thing = Thing(name=name, thing_type="water well", release_status="draft") + session.add(thing) + session.commit() + session.refresh(thing) + + location = Location( + point="POINT(-107 33)", + elevation=0, + release_status="draft", + ) + session.add(location) + session.flush() + session.add( + LocationThingAssociation( + location_id=location.id, + thing_id=thing.id, + effective_start=date.today().isoformat(), + ) + ) + session.add( + MeasuringPointHistory( + thing_id=thing.id, + measuring_point_height=1.0, + measuring_point_description="autogenerated", + start_date=date.today(), + release_status="draft", + ) + ) + session.commit() + session.refresh(thing) + return thing + + # ============= EOF ============================================= diff --git a/tests/test_stratigraphy_legacy.py b/tests/test_stratigraphy_legacy.py index 4693eac64..144c99a29 100644 --- a/tests/test_stratigraphy_legacy.py +++ b/tests/test_stratigraphy_legacy.py @@ -1,7 +1,13 @@ import uuid import pandas as pd -from db import Stratigraphy, Thing +from db import ( + Stratigraphy, + Thing, + Location, + LocationThingAssociation, + MeasuringPointHistory, +) from db.engine import session_ctx from transfers.stratigraphy_legacy import StratigraphyLegacyTransferer @@ -12,6 +18,32 @@ def _create_test_thing(name: str = "ST-0001") -> Thing: session.add(thing) session.commit() session.refresh(thing) + + location = Location( + point="POINT(-107 33)", + elevation=0, + release_status="draft", + ) + session.add(location) + session.flush() + session.add( + LocationThingAssociation( + location_id=location.id, + thing_id=thing.id, + effective_start="2024-01-01", + ) + ) + session.add( + MeasuringPointHistory( + thing_id=thing.id, + measuring_point_height=1.0, + measuring_point_description="autogenerated", + start_date="2024-01-01", + release_status="draft", + ) + ) + session.commit() + session.refresh(thing) return thing diff --git a/tests/test_thing.py b/tests/test_thing.py index f60a32f7b..50c2c6da7 100644 --- a/tests/test_thing.py +++ b/tests/test_thing.py @@ -16,7 +16,6 @@ from datetime import timezone import pytest - from core.dependencies import ( admin_function, editor_function, @@ -750,11 +749,11 @@ def test_get_things(water_well_thing, spring_thing, location): response = client.get("/thing") assert response.status_code == 200 - expected_location = LocationResponse.model_validate(location).model_dump() - # created_at is already serialized to UTC format by UTCAwareDatetime - data = response.json() - assert data["total"] == 2 + assert data["total"] >= 2 + returned_ids = {item["id"] for item in data["items"]} + assert water_well_thing.id in returned_ids + assert spring_thing.id in returned_ids @pytest.mark.skip("Needs to be updated per changes made from feature files") From 35b5d0032bdd4e8c323fdd26f7c07e2c68c32adc Mon Sep 17 00:00:00 2001 From: Jake Ross Date: Fri, 16 Jan 2026 09:57:49 -0700 Subject: [PATCH 06/10] Update transfers/stratigraphy_legacy.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- transfers/stratigraphy_legacy.py | 8 ++++++-- 1 file changed, 6 insertions(+), 2 deletions(-) diff --git a/transfers/stratigraphy_legacy.py b/transfers/stratigraphy_legacy.py index 63ec13ca0..311aecfca 100644 --- a/transfers/stratigraphy_legacy.py +++ b/transfers/stratigraphy_legacy.py @@ -103,9 +103,13 @@ def _row_dict(self, row: Any) -> Dict[str, Any] | None: self._capture_error(point_id, "No Thing found for PointID", "thing_id") return None - global_id = self._uuid_value(getattr(row, "GlobalID", None)) + raw_global_id = getattr(row, "GlobalID", None) + if raw_global_id in (None, ""): + self._capture_error(point_id, "Missing GlobalID", "GlobalID") + return None + global_id = self._uuid_value(raw_global_id) if global_id is None: - self._capture_error(point_id, "Invalid GlobalID", "GlobalID") + self._capture_error(point_id, "Malformed GlobalID", "GlobalID") return None return { From fa53d476d772951190c5a5e3c27ca90dd54d3d63 Mon Sep 17 00:00:00 2001 From: Jake Ross Date: Fri, 16 Jan 2026 09:59:46 -0700 Subject: [PATCH 07/10] Update tests/test_geospatial.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- tests/test_geospatial.py | 3 --- 1 file changed, 3 deletions(-) diff --git a/tests/test_geospatial.py b/tests/test_geospatial.py index bcb1482d7..27252fe40 100644 --- a/tests/test_geospatial.py +++ b/tests/test_geospatial.py @@ -128,9 +128,6 @@ def _create_well_with_location(session, name: str, point_wkt: str) -> Thing: session.commit() return thing - # Cleanup - - def test_get_project_area(): response = client.get("/geospatial/project-area/1") assert response.status_code == 200 From 9eab1c98526e18cd4718c51832378e1357f29feb Mon Sep 17 00:00:00 2001 From: jirhiker Date: Fri, 16 Jan 2026 17:00:04 +0000 Subject: [PATCH 08/10] Formatting changes --- tests/test_geospatial.py | 1 + 1 file changed, 1 insertion(+) diff --git a/tests/test_geospatial.py b/tests/test_geospatial.py index 27252fe40..c2b993148 100644 --- a/tests/test_geospatial.py +++ b/tests/test_geospatial.py @@ -128,6 +128,7 @@ def _create_well_with_location(session, name: str, point_wkt: str) -> Thing: session.commit() return thing + def test_get_project_area(): response = client.get("/geospatial/project-area/1") assert response.status_code == 200 From 76dcc29e936d189be1372ebd62f68425f55458e3 Mon Sep 17 00:00:00 2001 From: Jake Ross Date: Fri, 16 Jan 2026 11:58:32 -0700 Subject: [PATCH 09/10] Update transfers/stratigraphy_legacy.py Co-authored-by: Copilot <175728472+Copilot@users.noreply.github.com> --- transfers/stratigraphy_legacy.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transfers/stratigraphy_legacy.py b/transfers/stratigraphy_legacy.py index 311aecfca..2a97ffc4d 100644 --- a/transfers/stratigraphy_legacy.py +++ b/transfers/stratigraphy_legacy.py @@ -28,7 +28,7 @@ def __init__(self, batch_size: int = 1000, *args, **kwargs) -> None: self.batch_size = batch_size self._thing_id_cache: dict[str, int] = {} - def _get_dfs(self): # type: ignore[override] + def _get_dfs(self) -> tuple[pd.DataFrame, pd.DataFrame]: df = read_csv("Stratigraphy") cleaned = replace_nans(df) cleaned = filter_to_valid_point_ids(cleaned, self.pointids) From b898e7492087e74ff140f9eeedd098df4bab69af Mon Sep 17 00:00:00 2001 From: jross Date: Fri, 16 Jan 2026 16:11:28 -0700 Subject: [PATCH 10/10] fix: remove dtype argument from read_csv call in transferer.py --- transfers/transferer.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/transfers/transferer.py b/transfers/transferer.py index 55f9324b6..c28304219 100644 --- a/transfers/transferer.py +++ b/transfers/transferer.py @@ -138,7 +138,7 @@ def _read_csv(self, name: str, dtype: dict | None = None, **kw) -> pd.DataFrame: csv_path = csv_paths.get(name) if csv_path: return pd.read_csv(csv_path, **kw) - return read_csv(name, dtype=dtype, **kw) + return read_csv(name, **kw) class ChunkTransferer(Transferer):