Add a TS Data loading script

Author: krowvin

cwmscli/load/timeseries/timeseries.py

@@ -1,3 +1,4 @@
+from datetime import datetime, timedelta
 from typing import Optional
 
 import click
@@ -53,3 +54,94 @@ def load_timeseries_ids_all(
         timeseries_id_regex=timeseries_id_regex,
         dry_run=dry_run,
     )
+
+
+@timeseries.command(
+    "data",
+    help="Copy timeseries data for a single location in a target CDA from a source CDA.",
+)
+@shared_source_target_options
+@click.option(
+    "--ts-id",
+    "ts_id",
+    default=None,
+    type=str,
+    help="ID of the timeseries to copy (e.g. 'LocID.*').",
+)
+@click.option(
+    "--ts-group",
+    "ts_group",
+    default=None,
+    type=str,
+    help="ID of the timeseries group to copy (e.g. 'GroupID.*').",
+)
+@click.option(
+    "--ts-group-category-id",
+    "ts_group_category_id",
+    default=None,
+    type=str,
+    help="ID of the timeseries group category to copy.",
+)
+@click.option(
+    "--ts-group-category-office-id",
+    "ts_group_category_office_id",
+    default=None,
+    type=str,
+    help="ID of the timeseries group category office to copy.",
+)
+@click.option(
+    "--begin",
+    "begin",
+    default=(datetime.now().astimezone() - timedelta(days=1)).replace(microsecond=0),
+    type=click.DateTime(formats=["%Y-%m-%dT%H:%M:%S%z", "%Y-%m-%dT%H:%M:%S %Z"]),
+    help="Start date for timeseries data (e.g. '2022-01-01T00:00:00-06:00').",
+)
+@click.option(
+    "--end",
+    "end",
+    default=datetime.now().astimezone().replace(microsecond=0),
+    type=click.DateTime(formats=["%Y-%m-%dT%H:%M:%S%z", "%Y-%m-%dT%H:%M:%S %Z"]),
+    help="End date for timeseries data (e.g. '2022-01-31T00:00:00-0600').",
+)
+@requires(reqs.cwms)
+@validate_cda_targets
+def load_timeseries_data(
+    source_cda: str,
+    source_office: str,
+    target_cda: str,
+    target_api_key: Optional[str],
+    verbose: int,
+    ts_id: Optional[str],
+    dry_run: bool,
+    ts_group: Optional[str],
+    ts_group_category_id: Optional[str],
+    ts_group_category_office_id: Optional[str],
+    begin: Optional[datetime] = None,
+    end: Optional[datetime] = None,
+):
+    if (ts_id is None) == (ts_group is None):
+        raise click.UsageError("Exactly one of --ts-id or --ts-group must be provided.")
+    if ts_group and not ts_group_category_id:
+        raise click.UsageError(
+            "When specifying --ts-group, --ts-group-category-id must also be provided."
+        )
+    if ts_group and not ts_group_category_office_id:
+        raise click.UsageError(
+            "When specifying --ts-group, --ts-group-category-office-id must also be provided."
+        )
+    from cwmscli.load.timeseries.timeseries_data import _load_timeseries_data
+
+    _load_timeseries_data(
+        source_cda=source_cda,
+        source_office=source_office,
+        target_cda=target_cda,
+        target_api_key=target_api_key,
+        verbose=verbose,
+        dry_run=dry_run,
+        ts_id=ts_id,
+        ts_group=ts_group,
+        ts_group_category_id=ts_group_category_id,
+        ts_group_category_office_id=ts_group_category_office_id,
+        begin=begin,
+        end=end,
+    )

cwmscli/load/timeseries/timeseries_data.py

@@ -0,0 +1,109 @@
+# cwmscli/load/timeseries_ids.py
+import json
+import logging
+from datetime import datetime
+from typing import Optional
+
+import click
+import pandas as pd
+from cwms import JSON
+
+
+def _load_timeseries_data(
+    source_cda: str,
+    source_office: str,
+    target_cda: str,
+    target_api_key: Optional[str],
+    verbose: int,
+    dry_run: bool,
+    ts_id: str,
+    ts_group: str,
+    ts_group_category_id: str,
+    ts_group_category_office_id: str,
+    begin: Optional[datetime] = None,
+    end: Optional[datetime] = None,
+):
+    import cwms
+
+    if verbose:
+        click.echo(
+            f"Loading timeseries data from source CDA '{source_cda}' (office '{source_office}') "
+            f"to target CDA '{target_cda}'."
+        )
+    cwms.init_session(api_root=source_cda, api_key=None)
+    # User has a ts_id
+    if ts_id and not ts_group:
+        ts_data = [
+            cwms.get_timeseries(
+                ts_id=ts_id, office_id=source_office, begin=begin, end=end
+            )
+        ]
+        # only grab time_ids for locations that are in the target database
+        try:
+            if dry_run:
+                click.echo("Dry run enabled. No changes will be made.")
+                logging.debug(f"Would store {ts_data} for {ts_id}({source_office})")
+            else:
+                cwms.init_session(api_root=target_cda, api_key=target_api_key)
+                cwms.store_timeseries(
+                    data=ts_data.json,
+                    store_rule="REPLACE_ALL",
+                    override_protection=False,
+                )
+        except Exception as e:
+            click.echo(f"Error storing timeseries ({ts_id}) data: {e}", err=True)
+    # User did not have a ts_id but has a ts_group
+    if ts_group:
+        ts_ids = cwms.get_timeseries_group(
+            group_id=ts_group,
+            category_id=ts_group_category_id,
+            office_id=source_office,
+            category_office_id=ts_group_category_office_id,
+        )
+        logging.info(
+            f"Found {len(ts_ids.json.get('assigned-time-series', []))} timeseries in group {ts_group}."
+        )
+        logging.info(f"Storing TSID from begin: {begin} to end: {end}")
+        for ts in ts_ids.json.get("assigned-time-series", []):
+            try:
+                if dry_run:
+                    click.echo(
+                        f"Would store timeseries data for {ts['timeseries-id']}({ts['office-id']})"
+                    )
+                else:
+                    cwms.init_session(api_root=source_cda, api_key=None)
+                    ts_data = cwms.get_timeseries(
+                        ts_id=ts["timeseries-id"],
+                        office_id=ts["office-id"],
+                        begin=begin,
+                        end=end,
+                    )
+                    cwms.init_session(api_root=target_cda, api_key=target_api_key)
+                    # Convert the TS Values to a format CDA expects
+                    ts_cda_format_values = [
+                        [
+                            pd.to_datetime(value["date-time"]).isoformat(),
+                            value["value"],
+                            value["quality-code"],
+                        ]
+                        for value in ts_data.json.get("values", [])
+                    ]
+                    # Build a payload dict with formatted values instead of copying the response object
+                    ts_data_copy = ts_data.json.copy()
+                    ts_data_copy["values"] = ts_cda_format_values
+                    cwms.store_timeseries(
+                        data=ts_data_copy,
+                        store_rule="REPLACE_ALL",
+                        override_protection=False,
+                    )
+                    click.echo(
+                        f"Wrote {len(ts_data_copy.get('values', []))} values to {ts['timeseries-id']}."
+                    )
+            except Exception as e:
+                logging.warning(
+                    f"Error storing timeseries ({ts['timeseries-id']}) data: {e}",
+                    exc_info=True,
+                )
+
+    if verbose:
+        click.echo("Timeseries ID copy operation completed.")