#299 fix export of variation_identity table. Parameterize output bucket. Add some additional error handling and log details. (#301)

Author: theferrit32

clinvar_ingest/.dev.env

@@ -3,3 +3,4 @@ BQ_DEST_PROJECT="clingen-dev"
 CLINVAR_INGEST_BQ_META_DATASET="clinvar_ingest"
 CLINVAR_INGEST_BUCKET="clinvar-ingest-dev"
 CLINVAR_INGEST_RELEASE_TAG="v1_2_3"         # note the underscore separators due to BQ naming limitations
+CLINVAR_GKS_BUCKET="clinvar-gks"

clinvar_ingest/cloud/bigquery/processing_history.py

@@ -650,6 +650,7 @@ def processed_entries_ready_for_sp_processing(
     query = f"""
     SELECT
         release_date,
+        final_dataset_id,
         vcv_file_type,
         vcv_pipeline_version,
         vcv_schema_version,

clinvar_ingest/config.py

@@ -1,17 +1,15 @@
 import os
 import pathlib
-from typing import Literal
+from typing import Annotated, Literal
 
 from dotenv import dotenv_values
-from pydantic import BaseModel, field_validator
+from pydantic import BaseModel, Field, field_validator
 
 _dotenv_env = os.environ.get("DOTENV_ENV", "dev")
 _dotenv_values = dotenv_values(pathlib.Path(__file__).parent / f".{_dotenv_env}.env")
 
 
-def env_or_dotenv_or(
-    key_name: str, default: str | None = None, throw: bool = False
-) -> str:
+def env_or_dotenv_or(key_name: str, default: str | None = None, throw: bool = False) -> str:
     """
     Retrieves a value from the environment.
     If not set, retrieve it from the dotenv file.
@@ -43,14 +41,10 @@ class BaseEnv(Env):
 def _get_base_env() -> BaseEnv:
     return BaseEnv(
         bq_dest_project=env_or_dotenv_or("BQ_DEST_PROJECT", throw=True),
-        bq_meta_dataset=env_or_dotenv_or(
-            "CLINVAR_INGEST_BQ_META_DATASET", default="clinvar_ingest"
-        ),
+        bq_meta_dataset=env_or_dotenv_or("CLINVAR_INGEST_BQ_META_DATASET", default="clinvar_ingest"),
         slack_token=env_or_dotenv_or("CLINVAR_INGEST_SLACK_TOKEN"),
         # defaults to test "clinvar-message-test"
-        slack_channel=env_or_dotenv_or(
-            "CLINVAR_INGEST_SLACK_CHANNEL", default="C06QFR0278D"
-        ),
+        slack_channel=env_or_dotenv_or("CLINVAR_INGEST_SLACK_CHANNEL", default="C06QFR0278D"),
         release_tag=env_or_dotenv_or("CLINVAR_INGEST_RELEASE_TAG", throw=True),
         schema_version=env_or_dotenv_or("CLINVAR_INGEST_SCHEMA_VERSION", default="v2"),
         location=env_or_dotenv_or("CLINVAR_INGEST_LOCATION", default="us-east1"),
@@ -77,28 +71,23 @@ def get_file_ingest_env() -> ClinVarEnv:
     env = ClinVarEnv(
         **_base_env.model_dump(),
         bucket_name=env_or_dotenv_or("CLINVAR_INGEST_BUCKET", throw=True),
-        bucket_staging_prefix=env_or_dotenv_or(
-            "CLINVAR_INGEST_STAGING_PREFIX", default="clinvar_xml"
-        ),
-        parse_output_prefix=env_or_dotenv_or(
-            "CLINVAR_INGEST_PARSED_PREFIX", default="clinvar_parsed"
-        ),
-        executions_output_prefix=env_or_dotenv_or(
-            "CLINVAR_INGEST_EXECUTIONS_PREFIX", default="executions"
-        ),
+        bucket_staging_prefix=env_or_dotenv_or("CLINVAR_INGEST_STAGING_PREFIX", default="clinvar_xml"),
+        parse_output_prefix=env_or_dotenv_or("CLINVAR_INGEST_PARSED_PREFIX", default="clinvar_parsed"),
+        executions_output_prefix=env_or_dotenv_or("CLINVAR_INGEST_EXECUTIONS_PREFIX", default="executions"),
     )
     _set_env(env)
     return env
 
 
 class StoredProceduresEnv(BaseEnv):
-    pass
+    clinvar_gks_bucket: Annotated[str, Field(description="The GKS bucket where the VI files are stored.")]
 
 
 def get_stored_procedures_env() -> StoredProceduresEnv:
     _base_env = _get_base_env()
     env = StoredProceduresEnv(
         **_base_env.model_dump(),
+        clinvar_gks_bucket=env_or_dotenv_or("CLINVAR_GKS_BUCKET", throw=True),
     )
     _set_env(env)
     return env

misc/bin/stored-procedures-workflow.py

@@ -4,7 +4,6 @@
 # stored procedures against on or more datasets in the ingestion workflow.
 
 import logging
-import subprocess
 import sys
 import traceback
 
@@ -73,6 +72,11 @@ def sp_rollback_exception_handler(exc_type, exc_value, exc_traceback):
 env = get_stored_procedures_env()
 _logger.info(f"Stored procedures execution environment: {env}")
 
+if env.file_format_mode != ClinVarIngestFileFormat.SP.value:
+    msg = f"stored-procedure workflow got unexpected file_format_mode: {env.file_format_mode}"
+    _logger.warning(msg)
+    raise ValueError(msg)
+
 ################################################################
 #
 processing_history_table = processing_history.ensure_initialized(client=_get_bq_client())
@@ -85,17 +89,20 @@ def sp_rollback_exception_handler(exc_type, exc_value, exc_traceback):
 processed_entries_needing_sp_run = processing_history.processed_entries_ready_for_sp_processing(
     processing_history_view, client=_get_bq_client()
 )
-msg = f"Found {processed_entries_needing_sp_run.total_rows} datasets to run stored procedures on."
+total_rows = processed_entries_needing_sp_run.total_rows
+rows_needing_sp_run = list(processed_entries_needing_sp_run)
+release_dates_str = ", ".join(r.get("release_date").isoformat() for r in rows_needing_sp_run)
+msg = f"Found {total_rows} datasets to run stored procedures on. ({release_dates_str})"
 _logger.info(msg)
 
-if not processed_entries_needing_sp_run.total_rows:
+if not total_rows:
     sys.exit(0)
 
 send_slack_message(msg)
 
 # update processing_history.bq_ingest_started for ALL processing_history_view
 rows_to_ingest = []
-for row in processed_entries_needing_sp_run:
+for row in rows_needing_sp_run:
     rows_to_ingest.append(row)
     vcv_pipeline_version = row.get("vcv_pipeline_version", None)
     vcv_xml_release_date = row.get("vcv_xml_release_date", None)
@@ -131,12 +138,13 @@ def sp_rollback_exception_handler(exc_type, exc_value, exc_traceback):
 # Now process individual rows
 for row in rows_to_ingest:
     _logger.info(row)
-    release_date = row.get("release_date", None)
-    vcv_pipeline_version = row.get("vcv_pipeline_version", None)
-    vcv_release_date = row.get("vcv_release_date", None)
-    vcv_xml_release_date = row.get("vcv_xml_release_date", None)
-    vcv_bucket_dir = row.get("vcv_bucket_dir", None)
-    schema_version = row.get("vcv_schema_version", None)
+    # required
+    release_date = row["release_date"]
+    vcv_pipeline_version = row["vcv_pipeline_version"]
+    vcv_xml_release_date = row["vcv_xml_release_date"]
+    vcv_bucket_dir = row["vcv_bucket_dir"]
+    # optional
+    schema_version = row.get("vcv_schema_version")
 
     msg = f"Executing stored procedures on dataset dated {release_date}"
     _logger.info(msg)
@@ -154,13 +162,13 @@ def sp_rollback_exception_handler(exc_type, exc_value, exc_traceback):
             client=_get_bq_client(),
         )
         msg = f"""
-                Stored procedure execution successful for release dated {vcv_xml_release_date=} {vcv_pipeline_version=} release_tag={env.release_tag}.
+                Stored procedure execution successful for release dated vcv_xml_release_date={vcv_xml_release_date.isoformat()} {vcv_pipeline_version=} release_tag={env.release_tag}.
             """
         _logger.info(msg)
         send_slack_message(msg)
     except Exception as e:
         msg = f"""
-              Stored procedure execution failed for release dated {vcv_xml_release_date=} {vcv_pipeline_version=} release_tag={env.release_tag}.
+              Stored procedure execution failed for release dated vcv_xml_release_date={vcv_xml_release_date.isoformat()} {vcv_pipeline_version=} release_tag={env.release_tag}.
               """
         _logger.error(msg)
         send_slack_message(msg)
@@ -173,22 +181,22 @@ def sp_rollback_exception_handler(exc_type, exc_value, exc_traceback):
         if row["xml_release_date"] != str(vcv_xml_release_date) or row["release_tag"] != env.release_tag
     ]
 
-    dataset_id = row.get("final_dataset_id")
+    dataset_id = row["final_dataset_id"]
 
-    vi_gs_url = f"gs://clinvar-gks/{release_date}/dev/vi.jsonl.gz"
-    cmd = f"""
-    bq extract \
-        --destination_format NEWLINE_DELIMITED_JSON \
-        --compression GZIP \
-        '{dataset_id}.variation_identity' \
-        {vi_gs_url}
-    """
+    vi_gs_url = f"gs://{env.clinvar_gks_bucket}/{release_date}/dev/vi.jsonl.gz"
     try:
-        subprocess.run(cmd, shell=True, check=True, capture_output=True, text=True)
+        client = _get_bq_client()
+        table_id = f"{dataset_id}.variation_identity"
+        job_config = bigquery.ExtractJobConfig(
+            destination_format=bigquery.DestinationFormat.NEWLINE_DELIMITED_JSON, compression=bigquery.Compression.GZIP
+        )
+        extract_job = client.extract_table(table_id, vi_gs_url, job_config=job_config)
+        extract_job.result(timeout=1800)  # Wait for the job to complete (30 minute timeout)
         msg = f"Successfully exported variation_identity file to {vi_gs_url}"
         _logger.info(msg)
         send_slack_message(msg)
-    except subprocess.CalledProcessError as e:
-        raise RuntimeError(
-            f"Command failed: {e.cmd}\nReturn code: {e.returncode}\nStdout: {e.stdout}\nStderr: {e.stderr}"
-        )
+    except Exception as e:
+        error_msg = f"BigQuery extract job to {vi_gs_url} failed: {e}"
+        _logger.error(error_msg)
+        send_slack_message(error_msg)
+        raise e