Skip to content

Commit 8ae5bbb

Browse files
committed
refactor: Enhance transfer of waterlevelscontinuous_pressure_daily
- Cache Thing IDs and map PointID to thing_id to satisfy new FK - Filter orphan rows to prevent invalid inserts - Add focused transfer unit test to validate mapping and filtering
1 parent d6fb0fa commit 8ae5bbb

2 files changed

Lines changed: 77 additions & 3 deletions

File tree

Lines changed: 47 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,47 @@
1+
# ===============================================================================
2+
# Copyright 2026 ross
3+
#
4+
# Licensed under the Apache License, Version 2.0 (the "License");
5+
# you may not use this file except in compliance with the License.
6+
# You may obtain a copy of the License at
7+
#
8+
# http://www.apache.org/licenses/LICENSE-2.0
9+
#
10+
# Unless required by applicable law or agreed to in writing, software
11+
# distributed under the License is distributed on an "AS IS" BASIS,
12+
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
13+
# See the License for the specific language governing permissions and
14+
# limitations under the License.
15+
# ===============================================================================
16+
17+
import pandas as pd
18+
19+
from transfers.waterlevelscontinuous_pressure_daily import (
20+
NMA_WaterLevelsContinuous_Pressure_DailyTransferer,
21+
)
22+
23+
24+
def test_pressure_daily_transfer_filters_orphans(water_well_thing):
25+
transferer = NMA_WaterLevelsContinuous_Pressure_DailyTransferer(batch_size=1)
26+
df = pd.DataFrame(
27+
[
28+
{"PointID": water_well_thing.name, "GlobalID": "gid-1"},
29+
{"PointID": "MISSING-THING", "GlobalID": "gid-2"},
30+
]
31+
)
32+
33+
filtered = transferer._filter_to_valid_things(df)
34+
35+
assert list(filtered["PointID"]) == [water_well_thing.name]
36+
37+
38+
def test_pressure_daily_row_dict_sets_thing_id(water_well_thing):
39+
transferer = NMA_WaterLevelsContinuous_Pressure_DailyTransferer(batch_size=1)
40+
row = {"PointID": water_well_thing.name, "GlobalID": "gid-3"}
41+
42+
mapped = transferer._row_dict(row)
43+
44+
assert mapped["thing_id"] == water_well_thing.id
45+
46+
47+
# ============= EOF =============================================

transfers/waterlevelscontinuous_pressure_daily.py

Lines changed: 30 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -22,7 +22,8 @@
2222
from sqlalchemy.dialects.postgresql import insert
2323
from sqlalchemy.orm import Session
2424

25-
from db import NMA_WaterLevelsContinuous_Pressure_Daily
25+
from db import NMA_WaterLevelsContinuous_Pressure_Daily, Thing
26+
from db.engine import session_ctx
2627
from transfers.logger import logger
2728
from transfers.transferer import Transferer
2829
from transfers.util import read_csv
@@ -41,15 +42,39 @@ class NMA_WaterLevelsContinuous_Pressure_DailyTransferer(Transferer):
4142
def __init__(self, *args, batch_size: int = 1000, **kwargs):
4243
super().__init__(*args, **kwargs)
4344
self.batch_size = batch_size
45+
self._thing_id_cache: dict[str, int] = {}
46+
self._build_thing_id_cache()
47+
48+
def _build_thing_id_cache(self) -> None:
49+
with session_ctx() as session:
50+
things = session.query(Thing.name, Thing.id).all()
51+
self._thing_id_cache = {name: thing_id for name, thing_id in things}
52+
logger.info(f"Built Thing ID cache with {len(self._thing_id_cache)} entries")
53+
54+
def _filter_to_valid_things(self, df: pd.DataFrame) -> pd.DataFrame:
55+
valid_point_ids = set(self._thing_id_cache.keys())
56+
before_count = len(df)
57+
filtered_df = df[df["PointID"].isin(valid_point_ids)].copy()
58+
after_count = len(filtered_df)
59+
if before_count > after_count:
60+
skipped = before_count - after_count
61+
logger.warning(
62+
"Filtered out %s WaterLevelsContinuous_Pressure_Daily records without matching Things "
63+
"(%s valid, %s orphan records prevented)",
64+
skipped,
65+
after_count,
66+
skipped,
67+
)
68+
return filtered_df
4469

4570
def _get_dfs(self) -> tuple[pd.DataFrame, pd.DataFrame]:
4671
# Parse key datetime columns eagerly to avoid per-row parsing later.
4772
input_df = read_csv(
4873
self.source_table,
4974
parse_dates=["DateMeasured", "Created", "Updated"],
5075
)
51-
# No special cleaning/validation beyond raw import; keep identical copy.
52-
return input_df, input_df
76+
cleaned_df = self._filter_to_valid_things(input_df)
77+
return input_df, cleaned_df
5378

5479
def _transfer_hook(self, session: Session) -> None:
5580
rows = self._dedupe_rows(
@@ -71,6 +96,7 @@ def _transfer_hook(self, session: Session) -> None:
7196
"OBJECTID": excluded.OBJECTID,
7297
"WellID": excluded.WellID,
7398
"PointID": excluded.PointID,
99+
"thing_id": excluded.thing_id,
74100
"DateMeasured": excluded.DateMeasured,
75101
"TemperatureWater": excluded.TemperatureWater,
76102
"WaterHead": excluded.WaterHead,
@@ -104,6 +130,7 @@ def val(key: str) -> Optional[Any]:
104130
"OBJECTID": val("OBJECTID"),
105131
"WellID": val("WellID"),
106132
"PointID": val("PointID"),
133+
"thing_id": self._thing_id_cache.get(val("PointID")),
107134
"DateMeasured": val("DateMeasured"),
108135
"TemperatureWater": val("TemperatureWater"),
109136
"WaterHead": val("WaterHead"),

0 commit comments

Comments
 (0)