diff --git a/packages/datacommons-admin/datacommons_admin/admin_cli.py b/packages/datacommons-admin/datacommons_admin/admin_cli.py index d2dd327e..9d7051c5 100644 --- a/packages/datacommons-admin/datacommons_admin/admin_cli.py +++ b/packages/datacommons-admin/datacommons_admin/admin_cli.py @@ -12,23 +12,23 @@ # See the License for the specific language governing permissions and # limitations under the License. -from pathlib import Path import re import sys import urllib.request +from pathlib import Path from typing import Any, Tuple import click from google.api_core import exceptions from google.cloud import storage -from . import __version__ from datacommons_admin.infra_templates import ( BACKEND_TF_TEMPLATE, README_TEMPLATE, REMOTE_STATE_TEMPLATE, ) +from . import __version__ DEFAULT_BUCKET_LOCATION = "US" GITHUB_RAW_BASE_URL = "https://raw.githubusercontent.com/datacommonsorg/datacommons" @@ -95,7 +95,7 @@ def _create_and_configure_bucket( new_bucket.iam_configuration.uniform_bucket_level_access_enabled = True new_bucket.versioning_enabled = True new_bucket.patch() - click.secho(f" Enabling versioning...", fg="bright_black") + click.secho(" Enabling versioning...", fg="bright_black") click.secho(" Configuring bucket IAM policy...", fg="bright_black") policy = new_bucket.get_iam_policy(requested_policy_version=3) policy["roles/storage.objectAdmin"].add(f"projectEditor:{project_id}") @@ -145,8 +145,7 @@ def _ensure_bucket_ready( storage_client, bucket_name, project_id, location ): return True - else: - _abort_bucket_setup(is_default) + _abort_bucket_setup(is_default) def _configure_remote_state( @@ -247,7 +246,7 @@ def _validate_namespace(ns: str) -> Tuple[bool, str]: def _resolve_project_config( project_id: str, namespace: str, force: bool -) -> Tuple[str, str, Path]: +) -> tuple[str, str, Path]: """Resolves project ID and namespace, and determines target directory.""" if project_id: _log_resolved_value("Project ID", project_id, is_default=False) @@ -503,19 +502,19 @@ def init( ) -def _setup_ingestion_client() -> Tuple[Any, str, str]: +def _setup_ingestion_client() -> tuple[Any, str, str]: click.secho( "Fetching ingestion service URL, workflow service account, and Spanner details from Terraform outputs...", fg="bright_black", ) + from datacommons_admin.ingestion_helper_client import IngestionHelperClient from datacommons_admin.tf_utils import ( get_ingestion_service_url, get_ingestion_workflow_service_account_email, - get_spanner_instance_id, get_spanner_database_id, + get_spanner_instance_id, ) - from datacommons_admin.ingestion_helper_client import IngestionHelperClient url = get_ingestion_service_url() sa_email = get_ingestion_workflow_service_account_email() diff --git a/packages/datacommons-admin/datacommons_admin/ingest_cli.py b/packages/datacommons-admin/datacommons_admin/ingest_cli.py index 08c04a24..dbed9eb0 100644 --- a/packages/datacommons-admin/datacommons_admin/ingest_cli.py +++ b/packages/datacommons-admin/datacommons_admin/ingest_cli.py @@ -12,16 +12,18 @@ # See the License for the specific language governing permissions and # limitations under the License. -import click import re -from datacommons_admin.ingestion_job_client import IngestionJobClient +import click + +from datacommons_admin.ingestion_helper_client import IngestionHelperClient from datacommons_admin.tf_utils import ( get_ingestion_prep_job_name, + get_ingestion_service_url, + get_ingestion_workflow_name, get_ingestion_workflow_service_account_email, get_project_id, get_region, - get_ingestion_workflow_name, ) @@ -40,36 +42,41 @@ def ingest() -> None: def start(imports: str | None = None) -> None: """Start a data ingestion job execution.""" click.secho("Datacommons Admin Ingest Start", fg="cyan", bold=True) + click.secho( - "Fetching data job name and workflow service account from Terraform outputs...", + "Fetching Ingestion parameters from Terraform outputs...", fg="bright_black", ) + helper_url = get_ingestion_service_url() job_name = get_ingestion_prep_job_name() sa_email = get_ingestion_workflow_service_account_email() project_id = get_project_id() region = get_region() workflow_name = get_ingestion_workflow_name() - click.secho(f"Found data job: {job_name}", fg="green") - click.secho(f"Found workflow service account: {sa_email}", fg="green") - click.secho(f"Found GCP project ID: {project_id}", fg="green") - click.secho(f"Found GCP region: {region}", fg="green") + click.secho(f"Found Ingestion Helper URL: {helper_url}", fg="green") + click.secho(f"Found Data Job Name: {job_name}", fg="green") + click.secho(f"Found Workflow Service Account: {sa_email}", fg="green") + click.secho(f"Found GCP Project ID: {project_id}", fg="green") + click.secho(f"Found GCP Region: {region}", fg="green") click.secho( - f"Starting Cloud Run job '{job_name}' via Admin API (this may take a few moments)...", + "Triggering data ingestion via Ingestion Helper API...", fg="bright_black", ) - client = IngestionJobClient( - job_name, + client = IngestionHelperClient( + helper_url, service_account_email=sa_email, - project_id=project_id, - location=region, ) - result = client.start_job(imports=imports) + result = client.start_ingestion(imports=imports) click.secho("Successfully started ingestion job!", fg="green", bold=True) - res_name = result.get("name") or result.get("metadata", {}).get("name") + res_name = ( + result.get("operationName") + or result.get("name") + or result.get("metadata", {}).get("name") + ) if res_name: op_pattern = r"projects/([^/]+)/locations/([^/]+)/operations/([^/]+)" @@ -110,27 +117,28 @@ def show_config() -> None: fg="bright_black", ) + helper_url = get_ingestion_service_url() job_name = get_ingestion_prep_job_name() sa_email = get_ingestion_workflow_service_account_email() project_id = get_project_id() region = get_region() - click.secho(f"Found data job: {job_name}", fg="green") - click.secho(f"Found workflow service account: {sa_email}", fg="green") - click.secho(f"Found GCP project ID: {project_id}", fg="green") - click.secho(f"Found GCP region: {region}", fg="green") + click.secho(f"Found Ingestion Helper URL: {helper_url}", fg="green") + click.secho(f"Found Data Job Name: {job_name}", fg="green") + click.secho(f"Found Workflow Service Account: {sa_email}", fg="green") + click.secho(f"Found GCP Project ID: {project_id}", fg="green") + click.secho(f"Found GCP Region: {region}", fg="green") click.secho( - f"Fetching configuration for Cloud Run job '{job_name}'...", + f"Fetching configuration for Cloud Run job '{job_name}' via Ingestion Helper...", fg="bright_black", ) - client = IngestionJobClient( - job_name, + client = IngestionHelperClient( + helper_url, service_account_email=sa_email, - project_id=project_id, - location=region, ) - env_vars = client.get_config() + result = client.get_job_config() + env_vars = result.get("config", []) click.secho("\nCurrent ingestion job configuration:", fg="cyan", bold=True) if not env_vars: diff --git a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py index 9be2a852..b1bb2241 100644 --- a/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py +++ b/packages/datacommons-admin/datacommons_admin/ingestion_helper_client.py @@ -120,3 +120,21 @@ def initialize_database(self) -> dict: def seed_database(self) -> dict: """Calls the seed_database endpoint on the ingestion helper service.""" return self._call_endpoint("database/seed") + + def start_ingestion( + self, job_name: str | None = None, imports: str | None = None + ) -> dict: + """Calls the ingestion/start endpoint on the ingestion helper service.""" + payload = {} + if job_name: + payload["jobName"] = job_name + if imports: + payload["imports"] = imports + return self._call_endpoint("ingestion/start", payload=payload) + + def get_job_config(self, job_name: str | None = None) -> dict: + """Calls the ingestion/config endpoint on the ingestion helper service to get environment variables.""" + payload = {} + if job_name: + payload["jobName"] = job_name + return self._call_endpoint("ingestion/config", payload=payload) diff --git a/packages/datacommons-admin/datacommons_admin/ingestion_job_client.py b/packages/datacommons-admin/datacommons_admin/ingestion_job_client.py deleted file mode 100644 index 93f9c4d7..00000000 --- a/packages/datacommons-admin/datacommons_admin/ingestion_job_client.py +++ /dev/null @@ -1,152 +0,0 @@ -# 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. - -import click -import google.auth -from google.auth.transport.requests import AuthorizedSession - - -class IngestionJobClient: - """Client for interacting with Cloud Run Admin API to manage CDC data ingestion jobs.""" - - def __init__( - self, - job_name: str, - service_account_email: str = None, - project_id: str = None, - location: str = None, - ) -> None: - self.service_account_email = service_account_email - base_credentials, _ = google.auth.default() - - if not job_name.startswith("projects/"): - if not project_id: - raise click.ClickException( - "Project ID must be provided via Terraform outputs or as an argument when job name is not a full resource name." - ) - if not location: - raise click.ClickException( - "Location must be provided via Terraform outputs or as an argument when job name is not a full resource name." - ) - self.full_job_name = ( - f"projects/{project_id}/locations/{location}/jobs/{job_name}" - ) - else: - self.full_job_name = job_name - - if service_account_email: - from google.auth import impersonated_credentials - - creds = impersonated_credentials.Credentials( - source_credentials=base_credentials, - target_principal=service_account_email, - target_scopes=["https://www.googleapis.com/auth/cloud-platform"], - ) - else: - creds = base_credentials - - self.session = AuthorizedSession(creds) - - def start_job(self, imports: str | None = None) -> dict: - """Starts an execution of the Cloud Run job.""" - url = f"https://run.googleapis.com/v2/{self.full_job_name}:run" - - container_override = {"env": [{"name": "DATA_RUN_MODE", "value": "dcpbridge"}]} - if imports: - container_override["args"] = [f"--imports={imports}"] - - json_payload = {"overrides": {"containerOverrides": [container_override]}} - - try: - response = self.session.post(url, json=json_payload, timeout=300) - except Exception as e: - msg = f"Network or authentication error connecting to Cloud Run Admin API at {url}: {e}" - if self.service_account_email: - msg += f"\nFailed to impersonate {self.service_account_email}. Please ensure your GCP user account has the 'Service Account Token Creator' (roles/iam.serviceAccountTokenCreator) IAM role." - raise click.ClickException(msg) - - if response.status_code == 401: - raise click.ClickException( - f"HTTP 401 Unauthorized when calling Cloud Run Admin API at {url}.\n" - "Your GCP credentials were rejected. Please verify your authentication.\n" - "To re-authenticate, run:\n" - " gcloud auth application-default login" - ) - - if not response.ok: - try: - error_data = response.json() - error_msg = ( - error_data.get("message") - or error_data.get("error", {}).get("message") - or response.text - ) - except Exception: - error_msg = response.text - - raise click.ClickException( - f"Cloud Run Admin API returned HTTP {response.status_code}: {error_msg}" - ) - - try: - return response.json() - except Exception: - return {"status": "success", "message": response.text} - - def get_config(self) -> list: - """Retrieves the environment variables configuration of the Cloud Run job.""" - url = f"https://run.googleapis.com/v2/{self.full_job_name}" - try: - response = self.session.get(url, timeout=300) - except Exception as e: - msg = f"Network or authentication error connecting to Cloud Run Admin API at {url}: {e}" - if self.service_account_email: - msg += f"\nFailed to impersonate {self.service_account_email}. Please ensure your GCP user account has the 'Service Account Token Creator' (roles/iam.serviceAccountTokenCreator) IAM role." - raise click.ClickException(msg) - - if response.status_code == 401: - raise click.ClickException( - f"HTTP 401 Unauthorized when calling Cloud Run Admin API at {url}.\n" - "Your GCP credentials were rejected. Please verify your authentication.\n" - "To re-authenticate, run:\n" - " gcloud auth application-default login" - ) - - if not response.ok: - try: - error_data = response.json() - error_msg = ( - error_data.get("message") - or error_data.get("error", {}).get("message") - or response.text - ) - except Exception: - error_msg = response.text - - raise click.ClickException( - f"Cloud Run Admin API returned HTTP {response.status_code}: {error_msg}" - ) - - try: - job_data = response.json() - except Exception as e: - raise click.ClickException(f"Failed to parse Cloud Run job response: {e}") - - containers = ( - job_data.get("template", {}).get("template", {}).get("containers", []) - ) - if not containers: - return [] - - return containers[0].get("env", []) diff --git a/packages/datacommons-admin/datacommons_admin/tf_utils.py b/packages/datacommons-admin/datacommons_admin/tf_utils.py index ed6ed9c2..96172f83 100644 --- a/packages/datacommons-admin/datacommons_admin/tf_utils.py +++ b/packages/datacommons-admin/datacommons_admin/tf_utils.py @@ -18,7 +18,6 @@ import click - TF_OUTPUT_INGESTION_SERVICE_URL = "ingestion_service_url" TF_OUTPUT_INGESTION_WORKFLOW_SERVICE_ACCOUNT_EMAIL = ( "ingestion_workflow_service_account_email" @@ -40,7 +39,7 @@ def get_terraform_output(key: str) -> str: try: result = subprocess.run( - ["terraform", "output", "-json"], + ["terraform", "output", "-json"], # noqa: S607 capture_output=True, text=True, check=True, @@ -49,14 +48,14 @@ def get_terraform_output(key: str) -> str: raise click.ClickException( f"Failed to run 'terraform output'. Are you in an initialized Terraform deployment directory?\n" f"Error details: {e.stderr.strip() or e.stdout.strip()}" - ) + ) from e try: outputs = json.loads(result.stdout) - except json.JSONDecodeError: + except json.JSONDecodeError as e: raise click.ClickException( "Failed to parse 'terraform output -json'. The output was not valid JSON." - ) + ) from e if not outputs: from pathlib import Path @@ -73,11 +72,10 @@ def get_terraform_output(key: str) -> str: f"No Terraform outputs found in '{cwd}'.\n" "Please navigate to your initialized DCP Terraform directory (e.g., 'cd my-namespace') and ensure 'terraform apply' has been run." ) - else: - raise click.ClickException( - f"No Terraform outputs found in '{cwd}'.\n" - "Please ensure you have successfully run 'terraform apply' to generate the deployment state." - ) + raise click.ClickException( + f"No Terraform outputs found in '{cwd}'.\n" + "Please ensure you have successfully run 'terraform apply' to generate the deployment state." + ) if key not in outputs: from pathlib import Path diff --git a/packages/datacommons-admin/tests/test_admin_cli.py b/packages/datacommons-admin/tests/test_admin_cli.py index 946426fc..fb465259 100644 --- a/packages/datacommons-admin/tests/test_admin_cli.py +++ b/packages/datacommons-admin/tests/test_admin_cli.py @@ -15,10 +15,9 @@ from pathlib import Path from unittest.mock import patch -from click.testing import CliRunner import pytest - -from datacommons_admin.admin_cli import admin, init +from click.testing import CliRunner +from datacommons_admin.admin_cli import admin @pytest.fixture @@ -325,8 +324,8 @@ def test_seed_db_success( @patch("datacommons_admin.tf_utils.shutil.which") @patch("datacommons_admin.tf_utils.subprocess.run") -@patch("datacommons_admin.ingestion_job_client.AuthorizedSession") -@patch("datacommons_admin.ingestion_job_client.google.auth.default") +@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession") +@patch("datacommons_admin.ingestion_helper_client.google.auth.default") def test_ingest_start_success( mock_auth_default: patch, mock_session: patch, @@ -338,7 +337,7 @@ def test_ingest_start_success( from unittest.mock import MagicMock mock_proc = MagicMock() - mock_proc.stdout = '{"ingestion_prep_job_name": {"value": "projects/mock-proj/locations/us-central1/jobs/mock-job"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "project_id": {"value": "mock-proj"}, "region": {"value": "us-central1"}, "ingestion_workflow_name": {"value": "mock-workflow"}}' + mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_prep_job_name": {"value": "projects/mock-proj/locations/us-central1/jobs/mock-job"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "project_id": {"value": "mock-proj"}, "region": {"value": "us-central1"}, "ingestion_workflow_name": {"value": "mock-workflow"}}' mock_run.return_value = mock_proc mock_creds = MagicMock() @@ -348,7 +347,7 @@ def test_ingest_start_success( mock_resp = MagicMock() mock_resp.ok = True mock_resp.json.return_value = { - "name": "projects/mock-proj/locations/us-central1/operations/op-123" + "operationName": "projects/mock-proj/locations/us-central1/operations/op-123" } mock_session_inst.post.return_value = mock_resp mock_session.return_value = mock_session_inst @@ -369,8 +368,8 @@ def test_ingest_start_success( @patch("datacommons_admin.tf_utils.shutil.which") @patch("datacommons_admin.tf_utils.subprocess.run") -@patch("datacommons_admin.ingestion_job_client.AuthorizedSession") -@patch("datacommons_admin.ingestion_job_client.google.auth.default") +@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession") +@patch("datacommons_admin.ingestion_helper_client.google.auth.default") def test_ingest_show_config_success( mock_auth_default: patch, mock_session: patch, @@ -382,7 +381,7 @@ def test_ingest_show_config_success( from unittest.mock import MagicMock mock_proc = MagicMock() - mock_proc.stdout = '{"ingestion_prep_job_name": {"value": "mock-job"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "project_id": {"value": "mock-proj"}, "region": {"value": "us-central1"}}' + mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_prep_job_name": {"value": "mock-job"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "project_id": {"value": "mock-proj"}, "region": {"value": "us-central1"}}' mock_run.return_value = mock_proc mock_creds = MagicMock() @@ -392,20 +391,12 @@ def test_ingest_show_config_success( mock_resp = MagicMock() mock_resp.ok = True mock_resp.json.return_value = { - "template": { - "template": { - "containers": [ - { - "env": [ - {"name": "GCS_BUCKET", "value": "my-test-bucket"}, - {"name": "API_KEY", "valueSource": "secret-api-key"}, - ] - } - ] - } - } + "config": [ + {"name": "GCS_BUCKET", "value": "my-test-bucket"}, + {"name": "API_KEY", "valueSource": "secret-api-key"}, + ] } - mock_session_inst.get.return_value = mock_resp + mock_session_inst.post.return_value = mock_resp mock_session.return_value = mock_session_inst result = runner.invoke(admin, ["ingest", "show-config"]) @@ -446,8 +437,8 @@ def test_init_uses_default_ref_v_prefixed( @patch("datacommons_admin.tf_utils.shutil.which") @patch("datacommons_admin.tf_utils.subprocess.run") -@patch("datacommons_admin.ingestion_job_client.AuthorizedSession") -@patch("datacommons_admin.ingestion_job_client.google.auth.default") +@patch("datacommons_admin.ingestion_helper_client.AuthorizedSession") +@patch("datacommons_admin.ingestion_helper_client.google.auth.default") def test_ingest_start_with_imports_success( mock_auth_default: patch, mock_session: patch, @@ -459,7 +450,7 @@ def test_ingest_start_with_imports_success( from unittest.mock import MagicMock mock_proc = MagicMock() - mock_proc.stdout = '{"ingestion_prep_job_name": {"value": "projects/mock-proj/locations/us-central1/jobs/mock-job"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "project_id": {"value": "mock-proj"}, "region": {"value": "us-central1"}, "ingestion_workflow_name": {"value": "mock-workflow"}}' + mock_proc.stdout = '{"ingestion_service_url": {"value": "https://mock-helper"}, "ingestion_prep_job_name": {"value": "projects/mock-proj/locations/us-central1/jobs/mock-job"}, "ingestion_workflow_service_account_email": {"value": "mock-orch-sa@mock.com"}, "project_id": {"value": "mock-proj"}, "region": {"value": "us-central1"}, "ingestion_workflow_name": {"value": "mock-workflow"}}' mock_run.return_value = mock_proc mock_creds = MagicMock() @@ -469,7 +460,7 @@ def test_ingest_start_with_imports_success( mock_resp = MagicMock() mock_resp.ok = True mock_resp.json.return_value = { - "name": "projects/mock-proj/locations/us-central1/operations/op-123" + "operationName": "projects/mock-proj/locations/us-central1/operations/op-123" } mock_session_inst.post.return_value = mock_resp mock_session.return_value = mock_session_inst @@ -478,16 +469,9 @@ def test_ingest_start_with_imports_success( assert result.exit_code == 0 assert "Successfully started ingestion job!" in result.output mock_session_inst.post.assert_called_once_with( - "https://run.googleapis.com/v2/projects/mock-proj/locations/us-central1/jobs/mock-job:run", + "https://mock-helper/ingestion/start", json={ - "overrides": { - "containerOverrides": [ - { - "env": [{"name": "DATA_RUN_MODE", "value": "dcpbridge"}], - "args": ["--imports=oecd,doubleup"], - } - ] - } + "imports": "oecd,doubleup", }, timeout=300, )