diff --git a/import-automation/executor/app/executor/import_executor.py b/import-automation/executor/app/executor/import_executor.py index 0b3e08a17c..3a798da952 100644 --- a/import-automation/executor/app/executor/import_executor.py +++ b/import-automation/executor/app/executor/import_executor.py @@ -1066,6 +1066,16 @@ def _upload_import_inputs(self, import_dir: str, output_dir: str, src=manifest_file, dest=dest, ) + # Copy import-specific validation config file if specified + validation_config_file = import_spec.get('validation_config_file') + if validation_config_file: + val_cfg_src = os.path.join(import_dir, validation_config_file) + if os.path.exists(val_cfg_src): + self._upload_file_helper( + src=val_cfg_src, + dest= + f'{output_dir}/{version}/{os.path.basename(validation_config_file)}', + ) import_inputs = import_spec.get('import_inputs', []) errors = [] data_size = 0 diff --git a/import-automation/validator/Dockerfile b/import-automation/validator/Dockerfile new file mode 100644 index 0000000000..421537eed5 --- /dev/null +++ b/import-automation/validator/Dockerfile @@ -0,0 +1,35 @@ +# Copyright 2026 Google LLC +# +# 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 python:3.12-slim + +WORKDIR /app + +# Install fast uv package installer and dependencies +COPY --from=ghcr.io/astral-sh/uv:latest /uv /uvx /bin/ +COPY import-automation/validator/requirements.txt /app/requirements.txt +RUN uv pip install --system --no-cache -r /app/requirements.txt + +# Copy required tools and utilities directly from the data repo +COPY util/ /app/util/ +COPY tools/statvar_importer/ /app/tools/statvar_importer/ +COPY tools/import_differ/ /app/tools/import_differ/ +COPY tools/import_validation/ /app/tools/import_validation/ +COPY import-automation/validator/main.py /app/main.py +COPY import-automation/validator/validator_test.py /app/validator_test.py + +ENV PYTHONPATH="/app:/app/util:/app/tools/import_differ:/app/tools/import_validation:/app/tools/statvar_importer" +ENV PYTHONUNBUFFERED=1 + +ENTRYPOINT ["python", "/app/main.py"] diff --git a/import-automation/validator/cloudbuild.yaml b/import-automation/validator/cloudbuild.yaml new file mode 100644 index 0000000000..c80cb7289d --- /dev/null +++ b/import-automation/validator/cloudbuild.yaml @@ -0,0 +1,67 @@ +# Copyright 2026 Google LLC +# +# 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. +# +# Builds the docker image of import validator, verifies using unit tests, +# and pushes it to artifact registry. +# +# Run it from the data repo root using: +# gcloud builds submit --config=import-automation/validator/cloudbuild.yaml \ +# --substitutions=_DOCKER_IMAGE="us-docker.pkg.dev/datcom-ci/gcr.io/dc-import-validator" . + +steps: + # Docker Build from data repo root + - name: 'gcr.io/cloud-builders/docker' + id: 'build' + entrypoint: 'bash' + args: + - '-c' + - | + docker build -f import-automation/validator/Dockerfile \ + -t ${_DOCKER_IMAGE}:${_TAG} \ + -t ${_DOCKER_IMAGE}:latest . + env: ['DOCKER_BUILDKIT=1'] + + # Docker push to Google Artifact Registry + - name: 'gcr.io/cloud-builders/docker' + id: 'push' + args: ['push', '${_DOCKER_IMAGE}', '--all-tags'] + + # Run unit tests inside built container + - name: '${_DOCKER_IMAGE}:${_TAG}' + id: 'unit-test' + entrypoint: 'pytest' + args: + - '/app/validator_test.py' + - '/app/tools/import_differ/bigquery_differ_test.py' + - '/app/tools/import_validation/runner_test.py' + waitFor: ['push'] + + # Tag image as stable and push + - name: 'gcr.io/cloud-builders/docker' + id: 'tag' + entrypoint: 'bash' + waitFor: ['unit-test'] + args: + - '-c' + - | + docker tag ${_DOCKER_IMAGE}:${_TAG} ${_DOCKER_IMAGE}:stable \ + && docker push ${_DOCKER_IMAGE}:stable + +substitutions: + _DOCKER_IMAGE: 'us-docker.pkg.dev/datcom-ci/gcr.io/dc-import-validator' + _TAG: 'latest' + +options: + logging: CLOUD_LOGGING_ONLY + dynamicSubstitutions: true diff --git a/import-automation/validator/main.py b/import-automation/validator/main.py new file mode 100644 index 0000000000..d837e6f03b --- /dev/null +++ b/import-automation/validator/main.py @@ -0,0 +1,349 @@ +# Copyright 2026 Google LLC +# +# 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. +"""Standalone Cloud Run Job entrypoint for Differ + Validation.""" + +import json +import os +import sys +import tempfile +import time +from typing import Dict, List + +from absl import app +from absl import flags +from absl import logging +from google.cloud import storage + +_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +# Support running both inside Docker (/app) and from data/import-automation/validator/ +_ROOT_DIR = (_SCRIPT_DIR if os.path.exists(os.path.join(_SCRIPT_DIR, 'tools')) + else os.path.dirname(os.path.dirname(_SCRIPT_DIR))) + +for _p in [ + _ROOT_DIR, + os.path.join(_ROOT_DIR, 'util'), + os.path.join(_ROOT_DIR, 'tools', 'import_differ'), + os.path.join(_ROOT_DIR, 'tools', 'import_validation'), + os.path.join(_ROOT_DIR, 'tools', 'statvar_importer'), +]: + if _p not in sys.path: + sys.path.insert(0, _p) + +from tools.import_differ import bigquery_differ +from tools.import_validation.runner import ValidationRunner +from tools.import_validation.validation_config import merge_and_save_config +import file_util + +FLAGS = flags.FLAGS +flags.DEFINE_string( + 'import_name', '', + 'Import name in format : (required).') +flags.DEFINE_string('import_config', '{}', + 'JSON string of import executor config overrides.') +flags.DEFINE_string('version', '', + 'Explicit candidate version override (optional).') +flags.DEFINE_string('gcs_bucket', '', + 'GCS bucket name (defaults to env GCS_BUCKET_ID).') +if 'bq_dataset' not in FLAGS: + flags.DEFINE_string( + 'bq_dataset', 'datcom_import_differ', + 'BigQuery dataset ID for temporary differ tables.') +if 'bq_table_ttl_hours' not in FLAGS: + flags.DEFINE_integer( + 'bq_table_ttl_hours', 12, + 'TTL in hours for temporary BigQuery differ tables.') + + +def _read_gcs_text(client: storage.Client, bucket_name: str, + blob_path: str) -> str: + """Reads text content from a GCS blob if it exists, else returns ''.""" + bucket = client.bucket(bucket_name) + blob = bucket.blob(blob_path) + if not blob.exists(): + return '' + return blob.download_as_text().strip() + + +def _write_gcs_text(client: storage.Client, bucket_name: str, blob_path: str, + content: str) -> None: + """Uploads text content to a GCS blob.""" + bucket = client.bucket(bucket_name) + blob = bucket.blob(blob_path) + blob.upload_from_string(content) + + +def _discover_input_prefixes(client: storage.Client, bucket_name: str, + output_dir: str, version: str) -> List[str]: + """Discovers input0, input1, ... prefixes under gs://///.""" + base_prefix = f'{output_dir}/{version}/' + bucket = client.bucket(bucket_name) + iterator = client.list_blobs(bucket, + prefix=f'{base_prefix}input', + delimiter='/') + for _ in iterator: + pass + prefixes = set() + for p in iterator.prefixes: + top_folder = p[len(base_prefix):].rstrip('/') + if top_folder.startswith('input'): + prefixes.add(top_folder) + if not prefixes: + return ['input0'] + return sorted(prefixes) + + +def _resolve_validation_config(client: storage.Client, bucket_name: str, + output_dir: str, relative_import_dir: str, + import_name: str, version: str, + default_val_config: str, tmpdir: str) -> str: + """Resolves validation_config.json by merging GCS/local override with default config.""" + manifest_text = _read_gcs_text(client, bucket_name, + f'{output_dir}/{version}/manifest.json') + if not manifest_text: + return default_val_config + + try: + manifest = json.loads(manifest_text) + custom_cfg_rel = manifest.get('validation_config_file') + for spec in manifest.get('import_specifications', []): + if spec.get('import_name') == import_name and spec.get( + 'validation_config_file'): + custom_cfg_rel = spec.get('validation_config_file') + break + if not custom_cfg_rel: + return default_val_config + + override_filename = os.path.basename(custom_cfg_rel) + gcs_override_cfg = ( + f'gs://{bucket_name}/{output_dir}/{version}/{override_filename}') + if file_util.file_get_matching(gcs_override_cfg): + local_override_cfg = os.path.join(tmpdir, override_filename) + file_util.file_copy(gcs_override_cfg, local_override_cfg) + logging.info('Downloaded override validation config from GCS: %s', + gcs_override_cfg) + return merge_and_save_config(default_val_config, + local_override_cfg, tmpdir) + + local_repo_cfg = os.path.join(_ROOT_DIR, relative_import_dir, + custom_cfg_rel) + if os.path.exists(local_repo_cfg): + return merge_and_save_config(default_val_config, local_repo_cfg, + tmpdir) + logging.warning( + 'Custom validation config %s not found in GCS or local repo; using default.', + custom_cfg_rel) + return default_val_config + except Exception as exc: + logging.warning( + 'Failed to resolve custom validation config (%s); using default.', + exc) + return default_val_config + + +def run_validation_job(absolute_import_name: str, + import_config_str: str, + version_override: str, + bucket_name: str, + bq_dataset: str, + bq_table_ttl_hours: int = 12) -> int: + """Executes BigQuery differ and ValidationRunner for an import version.""" + start_time = time.time() + if ':' not in absolute_import_name: + raise ValueError( + f'--import_name must be :, got: {absolute_import_name}') + + relative_import_dir, import_name = absolute_import_name.split(':', 1) + output_dir = f'{relative_import_dir}/{import_name}' + + user_config = json.loads(import_config_str) if import_config_str else {} + if not bucket_name: + bucket_name = (user_config.get('storage_prod_bucket_name') or + os.environ.get('GCS_BUCKET_ID') or + 'datcom-prod-imports') + project_id = (user_config.get('gcp_project_id') or + os.environ.get('PROJECT_ID') or + os.environ.get('GOOGLE_CLOUD_PROJECT', '')) + if bq_dataset == 'datcom_import_differ': + bq_dataset = (user_config.get('bq_dataset') or + os.environ.get('BQ_DATASET') or bq_dataset) + if bq_table_ttl_hours == 12: + env_ttl = os.environ.get('BQ_TABLE_TTL_HOURS') + bq_table_ttl_hours = int( + user_config.get('bq_table_ttl_hours') or env_ttl or + bq_table_ttl_hours) + + ignore_validation_status = user_config.get('ignore_validation_status', + False) + enable_skip_status = user_config.get('enable_skip_status', True) + + gcs_client = storage.Client(project=project_id or None) + + version = version_override + if not version: + version = _read_gcs_text(gcs_client, bucket_name, + f'{output_dir}/staging_version.txt') + if not version: + raise RuntimeError( + f'No candidate version found in gs://{bucket_name}/{output_dir}/staging_version.txt' + ) + + latest_version = _read_gcs_text(gcs_client, bucket_name, + f'{output_dir}/latest_version.txt') + latest_version_uri = (f'gs://{bucket_name}/{output_dir}/{latest_version}' + if latest_version else '') + + logging.info('Running validator for %s version=%s (latest=%s)', output_dir, + version, latest_version or 'None') + + summary_raw = _read_gcs_text( + gcs_client, bucket_name, + f'{output_dir}/{version}/import_summary.json') + import_summary: Dict = json.loads(summary_raw) if summary_raw else { + 'import_name': import_name, + 'latest_version': f'gs://{bucket_name}/{output_dir}/{version}', + 'import_stats': {}, + } + + input_prefixes = _discover_input_prefixes(gcs_client, bucket_name, + output_dir, version) + + validation_status = True + differ_status = False + validation_data_size = 0 + + default_val_config = os.path.join(_ROOT_DIR, 'tools', 'import_validation', + 'validation_config.json') + + with tempfile.TemporaryDirectory() as tmpdir: + for input_prefix in input_prefixes: + genmcf_local_dir = os.path.join(tmpdir, input_prefix, 'genmcf') + val_local_dir = os.path.join(tmpdir, input_prefix, 'validation') + os.makedirs(genmcf_local_dir, exist_ok=True) + os.makedirs(val_local_dir, exist_ok=True) + + current_mcf_pattern = ( + f'gs://{bucket_name}/{output_dir}/{version}/{input_prefix}/genmcf/*.mcf' + ) + previous_mcf_pattern = ( + f'{latest_version_uri}/{input_prefix}/genmcf/*.mcf' + if latest_version_uri else '') + + diff_found = True + differ_output_dir = '' + + # 1. Run BigQuery Differ if previous version exists + if previous_mcf_pattern and file_util.file_get_matching( + previous_mcf_pattern): + logging.info('Running BigQuery differ for %s vs %s', + current_mcf_pattern, previous_mcf_pattern) + differ_summary = bigquery_differ.run_bigquery_differ( + current_data=current_mcf_pattern, + previous_data=previous_mcf_pattern, + output_location=val_local_dir, + project_id=project_id, + job_name=f'differ_{import_name}_{input_prefix}', + dataset_id=bq_dataset, + expiration_hours=bq_table_ttl_hours, + ) + diff_found = (differ_summary.get('obs_diff_count', 1) != 0 or + differ_summary.get('schema_diff_count', 1) != 0) + differ_output_dir = val_local_dir + else: + logging.info( + 'No previous MCF files found at %s; skipping differ.', + previous_mcf_pattern) + + if not differ_status: + differ_status = diff_found + + # 2. Download summary_report.csv and report.json for ValidationRunner + summary_stats_local = os.path.join(genmcf_local_dir, + 'summary_report.csv') + report_json_local = os.path.join(genmcf_local_dir, 'report.json') + gcs_summary_stats = ( + f'gs://{bucket_name}/{output_dir}/{version}/{input_prefix}/genmcf/summary_report.csv' + ) + gcs_report_json = ( + f'gs://{bucket_name}/{output_dir}/{version}/{input_prefix}/genmcf/report.json' + ) + if file_util.file_get_matching(gcs_summary_stats): + file_util.file_copy(gcs_summary_stats, summary_stats_local) + if file_util.file_get_matching(gcs_report_json): + file_util.file_copy(gcs_report_json, report_json_local) + + val_config_path = _resolve_validation_config( + gcs_client, bucket_name, output_dir, relative_import_dir, + import_name, version, default_val_config, tmpdir) + + val_output_file = os.path.join(val_local_dir, + 'validation_output.csv') + runner = ValidationRunner( + validation_config_path=val_config_path, + differ_output=differ_output_dir, + stats_summary=summary_stats_local, + lint_report=report_json_local, + validation_output=val_output_file, + ) + overall_status, _ = runner.run_validations() + validation_status = validation_status and overall_status + + # 3. Upload validation artifacts to GCS + gcs_val_dest = f'{output_dir}/{version}/{input_prefix}/validation' + bucket = gcs_client.bucket(bucket_name) + for fname in os.listdir(val_local_dir): + fpath = os.path.join(val_local_dir, fname) + if os.path.isfile(fpath): + validation_data_size += os.path.getsize(fpath) + dest_blob_name = f'{gcs_val_dest}/{fname}' + bucket.blob(dest_blob_name).upload_from_filename(fpath) + + # 4. Update import_summary.json status in GCS + if validation_status or ignore_validation_status: + if not differ_status and enable_skip_status: + import_summary['status'] = 'SKIP' + else: + import_summary['status'] = 'STAGING' + exit_code = 0 + else: + import_summary['status'] = 'VALIDATION' + exit_code = 1 + + stats = import_summary.setdefault('import_stats', {}) + stats['validation_execution_time'] = int(time.time() - start_time) + stats['validation_data_size'] = validation_data_size + + _write_gcs_text(gcs_client, bucket_name, + f'{output_dir}/{version}/import_summary.json', + json.dumps(import_summary, indent=2)) + logging.info('Completed validator: status=%s, exit_code=%d', + import_summary['status'], exit_code) + return exit_code + + +def main(_): + if not FLAGS.import_name: + raise ValueError('--import_name is required.') + code = run_validation_job( + absolute_import_name=FLAGS.import_name, + import_config_str=FLAGS.import_config, + version_override=FLAGS.version, + bucket_name=FLAGS.gcs_bucket, + bq_dataset=FLAGS.bq_dataset, + bq_table_ttl_hours=FLAGS.bq_table_ttl_hours, + ) + sys.exit(code) + + +if __name__ == '__main__': + app.run(main) diff --git a/import-automation/validator/requirements.txt b/import-automation/validator/requirements.txt new file mode 100644 index 0000000000..d3c05a6cba --- /dev/null +++ b/import-automation/validator/requirements.txt @@ -0,0 +1,15 @@ +absl-py>=2.1.0 +chardet>=5.2.0 +db-dtypes>=1.2.0 +duckdb>=1.0.0 +google-cloud-bigquery>=3.25.0 +google-cloud-storage>=2.18.0 +gspread>=6.0.0 +numpy>=1.26.0 +omegaconf>=2.3.0 +pandas>=2.2.0 +prettytable>=3.10.0 +psutil>=5.9.0 +pytest>=8.0.0 +requests>=2.31.0 +retry>=0.9.2 diff --git a/import-automation/validator/validator_test.py b/import-automation/validator/validator_test.py new file mode 100644 index 0000000000..2551613434 --- /dev/null +++ b/import-automation/validator/validator_test.py @@ -0,0 +1,137 @@ +# Copyright 2026 Google LLC +# +# 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. +"""Unit tests for import-automation/validator/main.py.""" + +import json +import os +import sys +import unittest +from unittest import mock + +_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +if _SCRIPT_DIR not in sys.path: + sys.path.insert(0, _SCRIPT_DIR) + +import main as validator_main + + +class ValidatorTest(unittest.TestCase): + + @mock.patch('main.ValidationRunner') + @mock.patch('main.bigquery_differ.run_bigquery_differ') + @mock.patch('main.file_util.file_copy') + @mock.patch('main.file_util.file_get_matching') + @mock.patch('main.storage.Client') + def test_run_validation_job_staging(self, mock_storage_cls, mock_matching, + mock_copy, mock_bq_differ, + mock_runner_cls): + mock_gcs = mock_storage_cls.return_value + mock_bucket = mock.MagicMock() + mock_gcs.bucket.return_value = mock_bucket + + blob_store = { + 'scripts/us_fed/treasury/staging_version.txt': + '2026_09_17', + 'scripts/us_fed/treasury/latest_version.txt': + '2026_09_10', + 'scripts/us_fed/treasury/2026_09_17/import_summary.json': + json.dumps({ + 'import_name': 'treasury', + 'status': 'PENDING' + }), + } + uploaded_blobs = {} + + def _get_blob(path): + b = mock.MagicMock() + b.name = path + b.exists.return_value = path in blob_store + b.download_as_text.return_value = blob_store.get(path, '') + + def _upload_str(content): + uploaded_blobs[path] = content + + b.upload_from_string.side_effect = _upload_str + return b + + mock_bucket.blob.side_effect = _get_blob + mock_iterator = mock.MagicMock() + mock_iterator.prefixes = [] + mock_gcs.list_blobs.return_value = mock_iterator + + mock_matching.return_value = ['gs://bucket/file.mcf'] + mock_bq_differ.return_value = { + 'obs_diff_count': 10, + 'schema_diff_count': 0, + } + + mock_runner = mock_runner_cls.return_value + mock_runner.run_validations.return_value = (True, []) + + exit_code = validator_main.run_validation_job( + absolute_import_name='scripts/us_fed:treasury', + import_config_str='{"gcp_project_id": "test-proj"}', + version_override='', + bucket_name='test-bucket', + bq_dataset='test_dataset', + ) + + self.assertEqual(exit_code, 0) + summary_blob_key = 'scripts/us_fed/treasury/2026_09_17/import_summary.json' + self.assertIn(summary_blob_key, uploaded_blobs) + updated_summary = json.loads(uploaded_blobs[summary_blob_key]) + self.assertEqual(updated_summary['status'], 'STAGING') + + @mock.patch('main.merge_and_save_config') + @mock.patch('main._read_gcs_text') + @mock.patch('main.file_util.file_copy') + @mock.patch('main.file_util.file_get_matching') + def test_resolve_validation_config_from_gcs(self, mock_matching, mock_copy, + mock_read_gcs, mock_merge): + mock_client = mock.MagicMock() + mock_read_gcs.return_value = json.dumps({ + 'import_specifications': [{ + 'import_name': 'treasury', + 'validation_config_file': 'custom_val.json', + }] + }) + mock_matching.return_value = [ + 'gs://test-bucket/scripts/us_fed/treasury/2026_09_17/custom_val.json' + ] + mock_merge.return_value = '/tmp/test_val/merged_validation_config.json' + resolved = validator_main._resolve_validation_config( + client=mock_client, + bucket_name='test-bucket', + output_dir='scripts/us_fed/treasury', + relative_import_dir='scripts/us_fed', + import_name='treasury', + version='2026_09_17', + default_val_config='/default/validation_config.json', + tmpdir='/tmp/test_val', + ) + self.assertEqual(resolved, + '/tmp/test_val/merged_validation_config.json') + mock_copy.assert_called_once_with( + 'gs://test-bucket/scripts/us_fed/treasury/2026_09_17/custom_val.json', + '/tmp/test_val/custom_val.json', + ) + mock_merge.assert_called_once_with( + '/default/validation_config.json', + '/tmp/test_val/custom_val.json', + '/tmp/test_val', + ) + + +if __name__ == '__main__': + unittest.main() diff --git a/tools/import_differ/bigquery_differ.py b/tools/import_differ/bigquery_differ.py new file mode 100644 index 0000000000..d154054d40 --- /dev/null +++ b/tools/import_differ/bigquery_differ.py @@ -0,0 +1,461 @@ +# Copyright 2026 Google LLC +# +# 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. +"""BigQuery-based MCF differ for low-memory ($O(1)$ RAM) execution. + +Streams current and previous `.mcf` files (local or `gs://` wildcard paths) into +normalized observation and schema CSVs in constant memory, loads them into +ephemeral BigQuery tables, executes SQL FULL OUTER JOIN queries to compute +per-StatVar ADDED/DELETED/MODIFIED counts, writes `differ_summary.json` and +`differ_summary.csv` to `--output_location`, and drops the temporary tables. + +To view all CLI flags, defaults, and usage details (recommended for agents): + $ python3 tools/import_differ/bigquery_differ.py --help + +Standalone CLI Usage: + $ python3 tools/import_differ/bigquery_differ.py \\ + --current_data="gs://bucket/import/2026_09_17/input0/genmcf/*.mcf" \\ + --previous_data="gs://bucket/import/2026_09_10/input0/genmcf/*.mcf" \\ + --output_location="/tmp/diff_results" \\ + --project_id="datcom-import-automation-prod" \\ + --bq_dataset="datcom_import_differ" \\ + --bq_table_ttl_hours=12 + +Usage via import_differ.py: + $ python3 tools/import_differ/import_differ.py \\ + --runner_mode=bigquery \\ + --current_data="gs://bucket/import/2026_09_17/input0/genmcf/*.mcf" \\ + --previous_data="gs://bucket/import/2026_09_10/input0/genmcf/*.mcf" \\ + --output_location="/tmp/diff_results" \\ + --project_id="datcom-import-automation-prod" +""" + +import csv +from datetime import datetime, timedelta, timezone +import os +import re +import sys +import tempfile +from typing import Any, Dict, List, Optional, Tuple +import uuid + +from absl import app +from absl import flags +from absl import logging +from google.cloud import bigquery + +_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +_DATA_DIR = os.path.dirname(os.path.dirname(_SCRIPT_DIR)) +sys.path.append(_SCRIPT_DIR) +sys.path.append(os.path.join(_DATA_DIR, 'util')) +sys.path.append(os.path.join(_DATA_DIR, 'tools', 'statvar_importer')) + +import differ_utils +from file_util import FileIO, file_get_matching +from mcf_file_util import normalize_value + +OBSERVATION_KEY_PROPERTIES = [ + 'variableMeasured', + 'observationAbout', + 'observationDate', + 'observationPeriod', + 'measurementMethod', + 'unit', + 'scalingFactor', +] + +OBS_BQ_SCHEMA = [ + bigquery.SchemaField('key_combined', 'STRING'), + bigquery.SchemaField('variableMeasured', 'STRING'), + bigquery.SchemaField('value', 'STRING'), +] + +SCHEMA_NODE_BQ_SCHEMA = [ + bigquery.SchemaField('dcid', 'STRING'), + bigquery.SchemaField('value_combined', 'STRING'), +] + + +def ensure_bq_dataset(bq_client: bigquery.Client, + dataset_ref: str, + expiration_hours: int = 12) -> None: + """Creates or updates the BigQuery dataset with the specified default table TTL.""" + expected_ttl_ms = expiration_hours * 3600 * 1000 + try: + ds = bq_client.get_dataset(dataset_ref) + if ds.default_table_expiration_ms != expected_ttl_ms: + ds.default_table_expiration_ms = expected_ttl_ms + bq_client.update_dataset(ds, ['default_table_expiration_ms']) + except Exception: + logging.info('Creating BigQuery dataset %s', dataset_ref) + ds = bigquery.Dataset(dataset_ref) + ds.default_table_expiration_ms = expected_ttl_ms + bq_client.create_dataset(ds, exists_ok=True) + + +def _normalize_prop_val(value) -> str: + """Normalizes a property value string identically to ImportDiffer.""" + if isinstance(value, list): + return ', '.join(sorted([_normalize_prop_val(v) for v in value])) + if isinstance(value, str): + return str(normalize_value(value)) + return str(value) if value is not None else '' + + +def _flush_node_to_csv(node: Dict[str, Any], obs_writer: csv.writer, + schema_writer: csv.writer) -> Tuple[int, int]: + """Writes a single parsed MCF node to either the observation or schema CSV writer.""" + type_of = node.get('typeOf', '') + type_of_list = type_of if isinstance(type_of, list) else [type_of] + if any('StatVarObservation' in str(t) for t in type_of_list): + key_parts = [ + _normalize_prop_val(node.get(prop, '')) + for prop in OBSERVATION_KEY_PROPERTIES + ] + key_combined = ';'.join(key_parts) + var_measured = _normalize_prop_val(node.get('variableMeasured', '')) + val = _normalize_prop_val(node.get('value', '')) + obs_writer.writerow([key_combined, var_measured, val]) + return 1, 0 + else: + raw_id = node.get('dcid') or node.get('Node', '') + dcid = _normalize_prop_val(raw_id) + if dcid and not dcid.startswith('dcid:'): + dcid = f'dcid:{dcid}' + props = [] + for k in sorted(node.keys()): + if k not in ('Node', 'dcid'): + props.append(f'{k}:{_normalize_prop_val(node[k])}') + value_combined = ';'.join(props) + schema_writer.writerow([dcid, value_combined]) + return 0, 1 + + +def stream_mcf_to_csv(mcf_pattern: str, obs_csv_path: str, + schema_csv_path: str) -> Tuple[int, int]: + """Streams MCF files node-by-node into observation and schema CSVs in O(1) memory.""" + mcf_files = file_get_matching(mcf_pattern) + total_obs = 0 + total_schema = 0 + + with FileIO(obs_csv_path, mode='w', encoding='utf-8') as obs_file, \ + FileIO(schema_csv_path, mode='w', encoding='utf-8') as schema_file: + obs_writer = csv.writer(obs_file) + schema_writer = csv.writer(schema_file) + obs_writer.writerow(['key_combined', 'variableMeasured', 'value']) + schema_writer.writerow(['dcid', 'value_combined']) + + for mcf_file in mcf_files: + logging.info('Streaming MCF file to CSV: %s', mcf_file) + with FileIO(mcf_file, mode='r', encoding='utf-8') as in_f: + current_node: Dict[str, Any] = {} + for raw_line in in_f: + line = raw_line.strip() + if not line or line.startswith('#'): + if current_node: + o_inc, s_inc = _flush_node_to_csv( + current_node, obs_writer, schema_writer) + total_obs += o_inc + total_schema += s_inc + current_node = {} + continue + if ':' in line: + k, v = line.split(':', 1) + k = k.strip() + v = v.strip() + if k == 'Node' and current_node: + o_inc, s_inc = _flush_node_to_csv( + current_node, obs_writer, schema_writer) + total_obs += o_inc + total_schema += s_inc + current_node = {} + if k in current_node: + if isinstance(current_node[k], list): + current_node[k].append(v) + else: + current_node[k] = [current_node[k], v] + else: + current_node[k] = v + if current_node: + o_inc, s_inc = _flush_node_to_csv(current_node, obs_writer, + schema_writer) + total_obs += o_inc + total_schema += s_inc + + logging.info('Completed streaming MCF to CSV: %d obs, %d schema nodes', + total_obs, total_schema) + return total_obs, total_schema + + +def load_csv_to_bq_table(bq_client: bigquery.Client, + csv_path: str, + table_ref: str, + schema: List[bigquery.SchemaField], + expiration_hours: int = 12) -> None: + """Loads a local or GCS CSV file into a BigQuery table with a specified TTL.""" + job_config = bigquery.LoadJobConfig( + source_format=bigquery.SourceFormat.CSV, + skip_leading_rows=1, + schema=schema, + write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE, + allow_quoted_newlines=True, + ) + if csv_path.startswith('gs://'): + load_job = bq_client.load_table_from_uri(csv_path, + table_ref, + job_config=job_config) + else: + with open(csv_path, 'rb') as f: + load_job = bq_client.load_table_from_file(f, + table_ref, + job_config=job_config) + load_job.result() + + table = bq_client.get_table(table_ref) + table.expires = datetime.now( + timezone.utc) + timedelta(hours=expiration_hours) + bq_client.update_table(table, ['expires']) + + +def load_mcf_to_bq_tables(bq_client: bigquery.Client, + mcf_pattern: str, + obs_table_ref: str, + schema_table_ref: str, + temp_dir: Optional[str] = None, + suffix: str = '', + expiration_hours: int = 12) -> Tuple[int, int]: + """Streams MCF files to CSVs and loads them into BigQuery observation and schema tables.""" + with tempfile.TemporaryDirectory() as local_tmpdir: + base_dir = temp_dir if temp_dir else local_tmpdir + tag = f'_{suffix}' if suffix else '' + obs_csv = os.path.join(base_dir, f'obs{tag}.csv') + schema_csv = os.path.join(base_dir, f'schema{tag}.csv') + + obs_count, schema_count = stream_mcf_to_csv(mcf_pattern, obs_csv, + schema_csv) + load_csv_to_bq_table(bq_client, obs_csv, obs_table_ref, OBS_BQ_SCHEMA, + expiration_hours) + load_csv_to_bq_table(bq_client, schema_csv, schema_table_ref, + SCHEMA_NODE_BQ_SCHEMA, expiration_hours) + return obs_count, schema_count + + +def _sanitize_job_suffix(job_name: str) -> str: + """Converts a job name into a safe BigQuery table suffix.""" + cleaned = re.sub(r'[^a-zA-Z0-9_]', '_', job_name) + return f'{cleaned}_{uuid.uuid4().hex[:8]}' + + +def run_bigquery_differ( + current_data: str, + previous_data: str, + output_location: str, + project_id: str, + job_name: str = 'differ', + dataset_id: str = 'datcom_import_differ', + gcs_temp_dir: Optional[str] = None, + expiration_hours: int = 12, +) -> Dict: + """Executes dataset diff using streaming MCF-to-CSV and BigQuery FULL OUTER JOIN.""" + if not project_id: + project_id = os.environ.get('GOOGLE_CLOUD_PROJECT') or os.environ.get( + 'PROJECT_ID', '') + + bq_client = bigquery.Client(project=project_id) + dataset_ref = f'{project_id}.{dataset_id}' + ensure_bq_dataset(bq_client, dataset_ref, expiration_hours=expiration_hours) + + suffix = _sanitize_job_suffix(job_name) + curr_obs_table = f'{dataset_ref}.curr_obs_{suffix}' + prev_obs_table = f'{dataset_ref}.prev_obs_{suffix}' + curr_schema_table = f'{dataset_ref}.curr_schema_{suffix}' + prev_schema_table = f'{dataset_ref}.prev_schema_{suffix}' + + try: + logging.info('Step 1/3: Loading current MCF data into BigQuery...') + curr_obs_count, curr_schema_count = load_mcf_to_bq_tables( + bq_client, + current_data, + curr_obs_table, + curr_schema_table, + temp_dir=gcs_temp_dir, + suffix=f'curr_{suffix}', + expiration_hours=expiration_hours) + + logging.info('Step 2/3: Loading previous MCF data into BigQuery...') + prev_obs_count, prev_schema_count = load_mcf_to_bq_tables( + bq_client, + previous_data, + prev_obs_table, + prev_schema_table, + temp_dir=gcs_temp_dir, + suffix=f'prev_{suffix}', + expiration_hours=expiration_hours) + + logging.info('Step 3/3: Running BigQuery FULL OUTER JOIN diff...') + obs_diff_sql = f""" + WITH diff AS ( + SELECT + COALESCE(c.variableMeasured, p.variableMeasured) AS variableMeasured, + CASE + WHEN p.key_combined IS NULL THEN 'ADDED' + WHEN c.key_combined IS NULL THEN 'DELETED' + WHEN c.value != p.value THEN 'MODIFIED' + ELSE 'UNMODIFIED' + END AS diff_type + FROM `{curr_obs_table}` c + FULL OUTER JOIN `{prev_obs_table}` p + ON c.key_combined = p.key_combined + WHERE c.value IS DISTINCT FROM p.value + ) + SELECT + REGEXP_REPLACE(variableMeasured, '^dcid:', '') AS StatVar, + COUNTIF(diff_type = 'ADDED') AS ADDED, + COUNTIF(diff_type = 'DELETED') AS DELETED, + COUNTIF(diff_type = 'MODIFIED') AS MODIFIED + FROM diff + GROUP BY StatVar + ORDER BY StatVar + """ + obs_diff_df = bq_client.query(obs_diff_sql).to_dataframe() + + schema_diff_sql = f""" + WITH diff AS ( + SELECT + CASE + WHEN p.dcid IS NULL THEN 'ADDED' + WHEN c.dcid IS NULL THEN 'DELETED' + WHEN c.value_combined != p.value_combined THEN 'MODIFIED' + ELSE 'UNMODIFIED' + END AS diff_type + FROM `{curr_schema_table}` c + FULL OUTER JOIN `{prev_schema_table}` p + ON c.dcid = p.dcid + WHERE c.value_combined IS DISTINCT FROM p.value_combined + ) + SELECT + COUNTIF(diff_type = 'ADDED') AS added_schema_count, + COUNTIF(diff_type = 'DELETED') AS deleted_schema_count, + COUNTIF(diff_type = 'MODIFIED') AS modified_schema_count + FROM diff + """ + schema_rows = list(bq_client.query(schema_diff_sql).result()) + if schema_rows: + added_schema = int(schema_rows[0].added_schema_count or 0) + deleted_schema = int(schema_rows[0].deleted_schema_count or 0) + modified_schema = int(schema_rows[0].modified_schema_count or 0) + else: + added_schema = deleted_schema = modified_schema = 0 + + finally: + for table_id in (curr_obs_table, prev_obs_table, curr_schema_table, + prev_schema_table): + bq_client.delete_table(table_id, not_found_ok=True) + + added_obs = int(obs_diff_df['ADDED'].sum()) if not obs_diff_df.empty else 0 + deleted_obs = int( + obs_diff_df['DELETED'].sum()) if not obs_diff_df.empty else 0 + modified_obs = int( + obs_diff_df['MODIFIED'].sum()) if not obs_diff_df.empty else 0 + obs_diff_total = added_obs + deleted_obs + modified_obs + schema_diff_total = added_schema + deleted_schema + modified_schema + + differ_summary = { + 'current_version': current_data, + 'previous_version': previous_data, + 'current_obs_count': curr_obs_count, + 'previous_obs_count': prev_obs_count, + 'current_schema_count': curr_schema_count, + 'previous_schema_count': prev_schema_count, + 'added_obs_count': added_obs, + 'deleted_obs_count': deleted_obs, + 'modified_obs_count': modified_obs, + 'added_schema_count': added_schema, + 'deleted_schema_count': deleted_schema, + 'modified_schema_count': modified_schema, + 'obs_diff_count': obs_diff_total, + 'schema_diff_count': schema_diff_total, + } + + with tempfile.TemporaryDirectory() as tmp_dir: + differ_utils.write_json_data(differ_summary, output_location, + 'differ_summary.json', tmp_dir) + differ_utils.write_csv_data(obs_diff_df, output_location, + 'differ_summary.csv', tmp_dir) + + logging.info('BigQuery Differ summary: %s', differ_summary) + return differ_summary + + +_FLAGS = flags.FLAGS + + +def _define_cli_flags(): + """Defines CLI flags when bigquery_differ.py is executed as a script.""" + if 'current_data' not in _FLAGS: + flags.DEFINE_string( + 'current_data', '', + 'Path to the current MCF data (local or gs:// wildcard supported).') + if 'previous_data' not in _FLAGS: + flags.DEFINE_string( + 'previous_data', '', + 'Path to the previous MCF data (local or gs:// wildcard supported).' + ) + if 'output_location' not in _FLAGS: + flags.DEFINE_string( + 'output_location', 'results', + 'Output directory (local or gs://) for differ_summary.json and differ_summary.csv.' + ) + if 'project_id' not in _FLAGS: + flags.DEFINE_string( + 'project_id', '', + 'GCP project ID for BigQuery jobs and dataset (defaults to GOOGLE_CLOUD_PROJECT).' + ) + if 'job_name' not in _FLAGS: + flags.DEFINE_string( + 'job_name', 'differ', + 'Prefix name for temporary BigQuery differ tables.') + if 'bq_dataset' not in _FLAGS: + flags.DEFINE_string('bq_dataset', 'datcom_import_differ', + 'BigQuery dataset ID for temporary differ tables.') + if 'gcs_temp_dir' not in _FLAGS: + flags.DEFINE_string( + 'gcs_temp_dir', None, + 'Optional temporary directory path for intermediate CSV files.') + if 'bq_table_ttl_hours' not in _FLAGS: + flags.DEFINE_integer( + 'bq_table_ttl_hours', 12, + 'TTL in hours for temporary BigQuery differ tables.') + + +def main(_): + """CLI entrypoint for running the BigQuery differ directly.""" + if not _FLAGS.current_data or not _FLAGS.previous_data: + raise ValueError( + 'Both --current_data and --previous_data are required. Run with --help for usage.' + ) + run_bigquery_differ( + current_data=_FLAGS.current_data, + previous_data=_FLAGS.previous_data, + output_location=_FLAGS.output_location, + project_id=_FLAGS.project_id, + job_name=_FLAGS.job_name, + dataset_id=_FLAGS.bq_dataset, + gcs_temp_dir=_FLAGS.gcs_temp_dir, + expiration_hours=_FLAGS.bq_table_ttl_hours, + ) + + +if __name__ == '__main__': + _define_cli_flags() + app.run(main) diff --git a/tools/import_differ/bigquery_differ_test.py b/tools/import_differ/bigquery_differ_test.py new file mode 100644 index 0000000000..26f1a91fe6 --- /dev/null +++ b/tools/import_differ/bigquery_differ_test.py @@ -0,0 +1,163 @@ +# Copyright 2026 Google LLC +# +# 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. +"""Unit tests for bigquery_differ.py.""" + +import csv +import json +import os +import sys +import tempfile +import unittest +from unittest import mock + +import pandas as pd + +_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) +_DATA_DIR = os.path.dirname(os.path.dirname(_SCRIPT_DIR)) +sys.path.append(_DATA_DIR) +sys.path.append(_SCRIPT_DIR) + +import bigquery_differ + +_SAMPLE_CURR_MCF = """ +Node: dcid:obs1 +typeOf: dcs:StatVarObservation +variableMeasured: dcid:Count_Person +observationAbout: dcid:country/USA +observationDate: "2024" +value: 335000000 + +Node: dcid:obs2 +typeOf: dcs:StatVarObservation +variableMeasured: dcid:Count_Person +observationAbout: dcid:country/CAN +observationDate: "2024" +value: 40000000 + +Node: dcid:Count_Person +typeOf: dcs:StatisticalVariable +populationType: dcs:Person +measuredProperty: dcs:count +""" + +_SAMPLE_PREV_MCF = """ +Node: dcid:obs1 +typeOf: dcs:StatVarObservation +variableMeasured: dcid:Count_Person +observationAbout: dcid:country/USA +observationDate: "2024" +value: 330000000 + +Node: dcid:Count_Person +typeOf: dcs:StatisticalVariable +populationType: dcs:Person +measuredProperty: dcs:count +""" + + +class BigQueryDifferTest(unittest.TestCase): + + def test_stream_mcf_to_csv(self): + with tempfile.TemporaryDirectory() as tmpdir: + mcf_path = os.path.join(tmpdir, 'test.mcf') + obs_csv = os.path.join(tmpdir, 'obs.csv') + schema_csv = os.path.join(tmpdir, 'schema.csv') + with open(mcf_path, 'w', encoding='utf-8') as f: + f.write(_SAMPLE_CURR_MCF) + + obs_count, schema_count = bigquery_differ.stream_mcf_to_csv( + mcf_path, obs_csv, schema_csv) + + self.assertEqual(obs_count, 2) + self.assertEqual(schema_count, 1) + + with open(obs_csv, 'r', encoding='utf-8') as f: + rows = list(csv.DictReader(f)) + self.assertEqual(len(rows), 2) + self.assertEqual(rows[0]['variableMeasured'], + 'dcid:Count_Person') + self.assertEqual(rows[0]['value'], '335000000') + self.assertIn('dcid:country/USA', rows[0]['key_combined']) + + with open(schema_csv, 'r', encoding='utf-8') as f: + rows = list(csv.DictReader(f)) + self.assertEqual(len(rows), 1) + self.assertEqual(rows[0]['dcid'], 'dcid:Count_Person') + self.assertIn('populationType:dcid:Person', + rows[0]['value_combined']) + + @mock.patch('bigquery_differ.bigquery.Client') + def test_run_bigquery_differ_mocked(self, mock_bq_client_cls): + mock_client = mock_bq_client_cls.return_value + + mock_table = mock.MagicMock() + mock_client.get_table.return_value = mock_table + + mock_obs_query_job = mock.MagicMock() + mock_obs_query_job.to_dataframe.return_value = pd.DataFrame([{ + 'StatVar': 'Count_Person', + 'ADDED': 1, + 'DELETED': 0, + 'MODIFIED': 1 + }]) + + mock_schema_row = mock.MagicMock() + mock_schema_row.added_schema_count = 0 + mock_schema_row.deleted_schema_count = 0 + mock_schema_row.modified_schema_count = 0 + mock_schema_query_job = mock.MagicMock() + mock_schema_query_job.result.return_value = [mock_schema_row] + + mock_client.query.side_effect = [ + mock_obs_query_job, mock_schema_query_job + ] + + with tempfile.TemporaryDirectory() as tmpdir: + curr_mcf = os.path.join(tmpdir, 'curr.mcf') + prev_mcf = os.path.join(tmpdir, 'prev.mcf') + out_dir = os.path.join(tmpdir, 'output') + with open(curr_mcf, 'w', encoding='utf-8') as f: + f.write(_SAMPLE_CURR_MCF) + with open(prev_mcf, 'w', encoding='utf-8') as f: + f.write(_SAMPLE_PREV_MCF) + + summary = bigquery_differ.run_bigquery_differ( + current_data=curr_mcf, + previous_data=prev_mcf, + output_location=out_dir, + project_id='test-project', + job_name='test_job', + dataset_id='test_dataset') + + self.assertEqual(summary['current_obs_count'], 2) + self.assertEqual(summary['previous_obs_count'], 1) + self.assertEqual(summary['added_obs_count'], 1) + self.assertEqual(summary['modified_obs_count'], 1) + self.assertEqual(summary['obs_diff_count'], 2) + self.assertEqual(summary['schema_diff_count'], 0) + + summary_json_path = os.path.join(out_dir, 'differ_summary.json') + summary_csv_path = os.path.join(out_dir, 'differ_summary.csv') + self.assertTrue(os.path.exists(summary_json_path)) + self.assertTrue(os.path.exists(summary_csv_path)) + + with open(summary_json_path, 'r', encoding='utf-8') as f: + saved_json = json.load(f) + self.assertEqual(saved_json['obs_diff_count'], 2) + + self.assertEqual(mock_client.delete_table.call_count, 4) + + +if __name__ == '__main__': + unittest.main() diff --git a/tools/import_differ/import_differ.py b/tools/import_differ/import_differ.py index 7fde461701..b005dd0fda 100644 --- a/tools/import_differ/import_differ.py +++ b/tools/import_differ/import_differ.py @@ -28,7 +28,10 @@ from absl import app from absl import flags from absl import logging -from googleapiclient.discovery import build +try: + from googleapiclient.discovery import build +except ImportError: + build = None _SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__)) _DATA_DIR = os.path.dirname(os.path.dirname(_SCRIPT_DIR)) @@ -75,9 +78,15 @@ flags.DEFINE_string('file_format', 'mcf', 'Format of the input data (mcf,tfrecord)') flags.DEFINE_string('runner_mode', 'native', - 'Runner mode (native/direct/cloud)') + 'Runner mode (native/direct/cloud/bigquery)') flags.DEFINE_string('job_name', 'differ', 'Name of the differ job.') -flags.DEFINE_string('project_id', '', 'GCP project id for the dataflow job.') +flags.DEFINE_string('project_id', '', 'GCP project id for the dataflow/bq job.') +if 'bq_dataset' not in _FLAGS: + flags.DEFINE_string('bq_dataset', 'datcom_import_differ', + 'BigQuery dataset ID for temporary differ tables.') +if 'bq_table_ttl_hours' not in _FLAGS: + flags.DEFINE_integer('bq_table_ttl_hours', 12, + 'TTL in hours for temporary BigQuery differ tables.') def val_str(value) -> str: @@ -127,7 +136,9 @@ def __init__(self, project_id='', job_name='differ', file_format='mcf', - runner_mode='native'): + runner_mode='native', + bq_dataset='datcom_import_differ', + bq_table_ttl_hours=12): self.current_data = current_data self.previous_data = previous_data self.output_path = output_location @@ -135,6 +146,8 @@ def __init__(self, self.job_name = job_name self.file_format = file_format self.runner_mode = runner_mode + self.bq_dataset = bq_dataset + self.bq_table_ttl_hours = bq_table_ttl_hours def _cleanup_data(self, df: pd.DataFrame): for column in [Diff.ADDED, Diff.DELETED, Diff.MODIFIED]: @@ -419,6 +432,18 @@ def run_differ(self): raise RuntimeError(f'Direct job {self.job_name} failed.') return + elif self.runner_mode == 'bigquery': + import bigquery_differ + logging.info("Invoking BigQuery mode for differ") + return bigquery_differ.run_bigquery_differ( + current_data=self.current_data, + previous_data=self.previous_data, + output_location=self.output_path, + project_id=self.project_id, + job_name=self.job_name, + dataset_id=self.bq_dataset, + expiration_hours=self.bq_table_ttl_hours, + ) else: # Runs native Python differ. current_dir = os.path.join(tmp_path, 'current') @@ -507,7 +532,8 @@ def main(_): differ = ImportDiffer(_FLAGS.current_data, _FLAGS.previous_data, _FLAGS.output_location, _FLAGS.project_id, _FLAGS.job_name, _FLAGS.file_format, - _FLAGS.runner_mode) + _FLAGS.runner_mode, _FLAGS.bq_dataset, + _FLAGS.bq_table_ttl_hours) differ.run_differ() diff --git a/tools/import_validation/runner.py b/tools/import_validation/runner.py index a8fbf556d6..a534989cd1 100644 --- a/tools/import_validation/runner.py +++ b/tools/import_validation/runner.py @@ -154,7 +154,12 @@ def _initialize_data_sources(self, stats_summary: str, lint_report: str, lint_report) def _load_differ_df_from_mcf(self, input_dir: str) -> pd.DataFrame: - """Parses MCF diff files and returns a summary DataFrame.""" + """Parses differ summary CSV or MCF diff files and returns a summary DataFrame.""" + summary_csv_path = os.path.join(input_dir, 'differ_summary.csv') + if os.path.exists(summary_csv_path) and os.path.getsize( + summary_csv_path) > 0: + return pd.read_csv(summary_csv_path) + import glob from collections import defaultdict stats = defaultdict(lambda: {'ADDED': 0, 'DELETED': 0, 'MODIFIED': 0})