Merge pull request #455 from DataIntegrationGroup/BDMS-520-1-1-Cleanup-2.0

feat: refactor transfer_all function and add transferable_wells utility

Author: jirhiker

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")

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 =============================================