From 8163ade4808ea1ca7d2998bd3f3996d72e2728f1 Mon Sep 17 00:00:00 2001 From: jakeross Date: Sun, 1 Feb 2026 00:14:23 +1100 Subject: [PATCH] feat: refactor transfer_all function and add transferable_wells utility --- transfers/transfer.py | 14 ++++---------- transfers/transferable_wells.py | 28 ++++++++++++++++++++++++++++ 2 files changed, 32 insertions(+), 10 deletions(-) create mode 100644 transfers/transferable_wells.py diff --git a/transfers/transfer.py b/transfers/transfer.py index 267334e1..67e6ba78 100644 --- a/transfers/transfer.py +++ b/transfers/transfer.py @@ -281,7 +281,7 @@ def _drop_and_rebuild_db() -> None: @timeit -def transfer_all(metrics, limit=100, profile_waterlevels: bool = True): +def transfer_all(metrics: Metrics) -> list[ProfileArtifact]: message("STARTING TRANSFER", new_line_at_top=False) if get_bool_env("DROP_AND_REBUILD_DB", False): logger.info("Dropping schema and rebuilding database from migrations") @@ -300,7 +300,7 @@ def transfer_all(metrics, limit=100, profile_waterlevels: bool = True): for field in transfer_options.__dataclass_fields__ }, ) - + limit = int(os.getenv("TRANSFER_LIMIT", 1000)) flags = {"TRANSFER_ALL_WELLS": True, "LIMIT": limit} message("TRANSFER_FLAGS") logger.info(flags) @@ -313,9 +313,7 @@ def transfer_all(metrics, limit=100, profile_waterlevels: bool = True): # ========================================================================= if continuous_water_levels_only: logger.info("CONTINUOUS_WATER_LEVELS set; running only continuous transfers") - _run_continuous_water_level_transfers( - metrics, flags, profile_waterlevels, profile_artifacts - ) + _run_continuous_water_level_transfers(metrics, flags) return profile_artifacts else: message("PHASE 1: FOUNDATIONAL TRANSFERS (PARALLEL)") @@ -631,13 +629,9 @@ def main(): "Set POSTGRES_DB=ocotilloapi_dev in .env file" ) - limit = int(os.getenv("TRANSFER_LIMIT", 1000)) - profile_waterlevels = get_bool_env("PROFILE_WATERLEVELS_CONTINUOUS", True) metrics = Metrics() - profile_artifacts = transfer_all( - metrics, limit=limit, profile_waterlevels=profile_waterlevels - ) + profile_artifacts = transfer_all(metrics) if get_bool_env("CLEANUP_LOCATIONS", True): message("CLEANING UP LOCATIONS") diff --git a/transfers/transferable_wells.py b/transfers/transferable_wells.py new file mode 100644 index 00000000..d27d1167 --- /dev/null +++ b/transfers/transferable_wells.py @@ -0,0 +1,28 @@ +# =============================================================================== +# Copyright 2026 ross +# +# Licensed under the Apache License, Version 2.0 (the "License"); +# you may not use this file except in compliance with the License. +# You may obtain a copy of the License at +# +# http://www.apache.org/licenses/LICENSE-2.0 +# +# Unless required by applicable law or agreed to in writing, software +# distributed under the License is distributed on an "AS IS" BASIS, +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +# See the License for the specific language governing permissions and +# limitations under the License. +# =============================================================================== +from transfers.util import read_csv, get_transferable_wells + + +def main(): + df = read_csv("WellData", dtype={"OSEWelltagID": str}) + df = get_transferable_wells(df) + df = df[["PointID", "DataSource"]] + df.to_csv("transferable_wells.csv", index=False, float_format="%.2f") + + +if __name__ == "__main__": + main() +# ============= EOF =============================================