data-collection-framework 0.1.0__py3-none-any.whl

dcf/gcp/bootstrap.py

@@ -0,0 +1,131 @@
+import base64
+import json
+import logging
+
+from google.api_core.exceptions import Conflict, NotFound
+from google.auth.credentials import Credentials
+from google.cloud import secretmanager, storage
+from googleapiclient import discovery
+from googleapiclient.errors import HttpError
+
+logger = logging.getLogger(__name__)
+
+_SA_ACCOUNT_ID = "dcf-lake"
+_SECRET_ID     = "dcf-lake-sa-key"
+
+
+def create_state_bucket(project_id: str, region: str, credentials: Credentials) -> str:
+    """Create the GCS bucket used for Terraform state. Returns bucket name."""
+    bucket_name = f"dcf-tf-state-{project_id}"
+    client = storage.Client(project=project_id, credentials=credentials)
+    try:
+        bucket = client.create_bucket(bucket_name, location=region)
+        bucket.versioning_enabled = True
+        bucket.patch()
+        logger.info("Created Terraform state bucket %s", bucket_name)
+    except Conflict:
+        logger.info("Terraform state bucket %s already exists", bucket_name)
+    except Exception as e:
+        from google.api_core.exceptions import Forbidden
+        if isinstance(e, Forbidden) or "billing" in str(e).lower():
+            raise RuntimeError(
+                f"Billing is not enabled for project '{project_id}'. "
+                f"Enable it at: https://console.cloud.google.com/billing"
+            ) from e
+        raise
+    return bucket_name
+
+
+def create_service_account(project_id: str, credentials: Credentials) -> str:
+    """Create the dcf-lake service account. Returns SA email."""
+    sa_email = f"{_SA_ACCOUNT_ID}@{project_id}.iam.gserviceaccount.com"
+    service  = discovery.build("iam", "v1", credentials=credentials, cache_discovery=False)
+    try:
+        service.projects().serviceAccounts().create(
+            name=f"projects/{project_id}",
+            body={
+                "accountId": _SA_ACCOUNT_ID,
+                "serviceAccount": {"displayName": "dcf Lake Service Account"},
+            },
+        ).execute()
+        logger.info("Created service account %s", sa_email)
+    except HttpError as e:
+        if e.resp.status == 409:
+            logger.info("Service account %s already exists", sa_email)
+        else:
+            raise
+    return sa_email
+
+
+def create_service_account_key(project_id: str, sa_email: str, credentials: Credentials) -> dict:
+    """Create a new JSON key for the SA. Returns decoded key dict."""
+    service = discovery.build("iam", "v1", credentials=credentials, cache_discovery=False)
+    result  = service.projects().serviceAccounts().keys().create(
+        name=f"projects/{project_id}/serviceAccounts/{sa_email}",
+        body={"privateKeyType": "TYPE_GOOGLE_CREDENTIALS_FILE"},
+    ).execute()
+    key_data = json.loads(base64.b64decode(result["privateKeyData"]).decode())
+    logger.info("Created SA key for %s", sa_email)
+    return key_data
+
+
+def store_key_in_secret_manager(project_id: str, key_data: dict, credentials: Credentials) -> str:
+    """
+    Store the SA key in Secret Manager as 'dcf-lake-sa-key'.
+    Creates the secret if it doesn't exist, then adds a new version.
+    Returns the full secret resource name.
+    """
+    client      = secretmanager.SecretManagerServiceClient(credentials=credentials)
+    parent      = f"projects/{project_id}"
+    secret_name = f"{parent}/secrets/{_SECRET_ID}"
+
+    try:
+        client.create_secret(request={
+            "parent":    parent,
+            "secret_id": _SECRET_ID,
+            "secret":    {"replication": {"automatic": {}}},
+        })
+        logger.info("Created Secret Manager secret %s", _SECRET_ID)
+    except Conflict:
+        logger.info("Secret %s already exists, adding new version", _SECRET_ID)
+
+    client.add_secret_version(request={
+        "parent":  secret_name,
+        "payload": {"data": json.dumps(key_data).encode()},
+    })
+    logger.info("Stored SA key in Secret Manager")
+    return secret_name
+
+
+def delete_secret(secret_name: str, credentials: Credentials) -> None:
+    """Delete a Secret Manager secret and all its versions."""
+    client = secretmanager.SecretManagerServiceClient(credentials=credentials)
+    try:
+        client.delete_secret(request={"name": secret_name})
+        logger.info("Deleted secret %s", secret_name)
+    except NotFound:
+        logger.info("Secret %s not found, skipping", secret_name)
+
+
+def delete_service_account(project_id: str, sa_email: str, credentials: Credentials) -> None:
+    """Delete the dcf-lake service account."""
+    service = discovery.build("iam", "v1", credentials=credentials, cache_discovery=False)
+    try:
+        service.projects().serviceAccounts().delete(
+            name=f"projects/{project_id}/serviceAccounts/{sa_email}",
+        ).execute()
+        logger.info("Deleted service account %s", sa_email)
+    except HttpError as e:
+        if e.resp.status == 404:
+            logger.info("Service account %s not found, skipping", sa_email)
+        else:
+            raise
+
+
+def fetch_service_account_key(project_id: str, secret_name: str) -> dict:
+    """Fetch the latest SA key from Secret Manager using ADC credentials."""
+    from .gcloud import get_credentials
+    credentials = get_credentials()
+    client      = secretmanager.SecretManagerServiceClient(credentials=credentials)
+    response    = client.access_secret_version(request={"name": f"{secret_name}/versions/latest"})
+    return json.loads(response.payload.data.decode())

dcf/gcp/gcloud.py

@@ -0,0 +1,42 @@
+import shutil
+import subprocess
+import logging
+
+import google.auth
+from google.auth.exceptions import DefaultCredentialsError
+
+logger = logging.getLogger(__name__)
+
+_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
+_INSTALL_URL = "https://cloud.google.com/sdk/docs/install"
+
+
+def get_credentials():
+    """
+    Return ADC credentials scoped to cloud-platform.
+
+    Resolution order:
+      1. Existing ADC (GOOGLE_APPLICATION_CREDENTIALS env var or gcloud ADC file)
+      2. If not configured but gcloud is installed, run `gcloud auth application-default login`
+         (opens a browser on the local machine) then retry.
+      3. If gcloud is not installed, raise RuntimeError with install instructions.
+    """
+    try:
+        creds, _ = google.auth.default(scopes=_SCOPES)
+        return creds
+    except DefaultCredentialsError:
+        pass
+
+    gcloud = shutil.which("gcloud")
+    if not gcloud:
+        raise RuntimeError(
+            "No Google credentials found and gcloud CLI is not installed.\n"
+            f"Install it at: {_INSTALL_URL}\n"
+            "Then run: gcloud auth application-default login"
+        )
+
+    logger.info("ADC not configured — running gcloud auth application-default login")
+    subprocess.run(["gcloud", "auth", "application-default", "login"], check=True)
+
+    creds, _ = google.auth.default(scopes=_SCOPES)
+    return creds

dcf/gcp/terraform.py

@@ -0,0 +1,151 @@
+import json
+import logging
+import os
+import shutil
+import subprocess
+from pathlib import Path
+
+from google.api_core.exceptions import NotFound, Forbidden
+from google.cloud import storage
+
+logger = logging.getLogger(__name__)
+
+_MODULE_DIR = Path(__file__).parent.parent / "infra" / "modules" / "gcp"
+_WORK_DIR   = Path.home() / ".dcf" / "terraform"
+
+
+def provision(
+    project_id: str,
+    region: str,
+    sa_email: str,
+    tf_state_bucket: str,
+) -> str:
+    """
+    Run terraform init + apply.
+    Auth is handled automatically via ADC (gcloud application-default credentials).
+    Returns the warehouse_bucket name from terraform output.
+    Raises RuntimeError on non-zero exit.
+    """
+    _WORK_DIR.mkdir(parents=True, exist_ok=True)
+
+    for tf_file in _MODULE_DIR.glob("*.tf"):
+        shutil.copy2(tf_file, _WORK_DIR / tf_file.name)
+
+    env = {
+        **os.environ,
+        "TF_INPUT":            "0",
+        "TF_PLUGIN_CACHE_DIR": str(_WORK_DIR / ".plugin-cache"),
+    }
+    (_WORK_DIR / ".plugin-cache").mkdir(exist_ok=True)
+
+    _run(
+        [
+            "terraform", "init", "-reconfigure",
+            f"-backend-config=bucket={tf_state_bucket}",
+            "-backend-config=prefix=terraform/state",
+        ],
+        _WORK_DIR, env,
+    )
+
+    # Import any already-existing resources so apply doesn't fail with 409
+    _import_existing_resources(project_id, _WORK_DIR, env)
+
+    _run(
+        [
+            "terraform", "apply", "-auto-approve",
+            f"-var=project_id={project_id}",
+            f"-var=region={region}",
+            f"-var=sa_email={sa_email}",
+        ],
+        _WORK_DIR, env,
+    )
+
+    return _read_output(_WORK_DIR, env, "warehouse_bucket")
+
+
+def destroy(
+    project_id: str,
+    region: str,
+    sa_email: str,
+    tf_state_bucket: str,
+) -> None:
+    """Run terraform destroy to remove all provisioned GCP resources."""
+    _WORK_DIR.mkdir(parents=True, exist_ok=True)
+
+    for tf_file in _MODULE_DIR.glob("*.tf"):
+        shutil.copy2(tf_file, _WORK_DIR / tf_file.name)
+
+    env = {
+        **os.environ,
+        "TF_INPUT":            "0",
+        "TF_PLUGIN_CACHE_DIR": str(_WORK_DIR / ".plugin-cache"),
+    }
+    (_WORK_DIR / ".plugin-cache").mkdir(exist_ok=True)
+
+    _run(
+        [
+            "terraform", "init", "-reconfigure",
+            f"-backend-config=bucket={tf_state_bucket}",
+            "-backend-config=prefix=terraform/state",
+        ],
+        _WORK_DIR, env,
+    )
+
+    _run(
+        [
+            "terraform", "destroy", "-auto-approve",
+            f"-var=project_id={project_id}",
+            f"-var=region={region}",
+            f"-var=sa_email={sa_email}",
+        ],
+        _WORK_DIR, env,
+    )
+
+
+def _import_existing_resources(project_id: str, work_dir: Path, env: dict) -> None:
+    """Import already-existing GCP resources into Terraform state to avoid 409 on apply."""
+    warehouse_bucket = f"dcf-warehouse-{project_id}"
+    client = storage.Client(project=project_id)
+    try:
+        client.get_bucket(warehouse_bucket)
+    except (NotFound, Forbidden):
+        return  # bucket doesn't exist yet — apply will create it
+
+    # Bucket exists; import it so terraform apply doesn't try to create it again
+    result = subprocess.run(
+        ["terraform", "import", "google_storage_bucket.warehouse", warehouse_bucket],
+        cwd=str(work_dir), env=env, capture_output=True, text=True,
+    )
+    if result.returncode == 0:
+        logger.info("Imported existing warehouse bucket '%s' into Terraform state", warehouse_bucket)
+    elif "already managed by Terraform" in result.stdout + result.stderr:
+        logger.info("Warehouse bucket '%s' already in Terraform state", warehouse_bucket)
+    else:
+        # Import failed for an unexpected reason — log and continue; apply may still succeed
+        logger.warning("terraform import returned non-zero: %s", result.stderr[-500:])
+
+
+def _run(cmd: list[str], cwd: Path, env: dict) -> None:
+    result = subprocess.run(cmd, cwd=str(cwd), env=env, capture_output=True, text=True)
+    if result.returncode != 0:
+        logger.error(
+            "Terraform command failed: %s\nSTDOUT: %s\nSTDERR: %s",
+            " ".join(cmd), result.stdout, result.stderr,
+        )
+        raise RuntimeError(
+            f"terraform {cmd[1]} failed (exit {result.returncode}): {result.stderr[-2000:]}"
+        )
+    logger.info("terraform %s OK", cmd[1])
+
+
+def _read_output(work_dir: Path, env: dict, key: str) -> str:
+    result = subprocess.run(
+        ["terraform", "output", "-json"],
+        cwd=str(work_dir), env=env, capture_output=True, text=True,
+    )
+    if result.returncode != 0:
+        raise RuntimeError(f"terraform output failed: {result.stderr}")
+    outputs = json.loads(result.stdout)
+    if key not in outputs:
+        raise RuntimeError(f"'{key}' not in terraform output. Got: {list(outputs)}")
+    return outputs[key]["value"]

dcf/infra/modules/batch_collector/gcp/airflow/main.tf

@@ -0,0 +1,194 @@
+terraform {
+  required_version = ">= 1.0"
+  required_providers {
+    google = {
+      source  = "hashicorp/google"
+      version = "~> 5.0"
+    }
+    local = {
+      source  = "hashicorp/local"
+      version = "~> 2.0"
+    }
+    null = {
+      source  = "hashicorp/null"
+      version = "~> 3.0"
+    }
+  }
+}
+
+provider "google" {
+  project = var.project_id
+  region  = var.region
+}
+
+resource "local_file" "dockerfile" {
+  content  = templatefile("${path.module}/templates/airflow.Dockerfile.tftpl", {
+    target = "gcp"
+  })
+  filename = "${var.build_context}/Dockerfile"
+}
+
+resource "null_resource" "build" {
+  depends_on = [local_file.dockerfile]
+
+  triggers = {
+    content_hash = var.content_hash
+  }
+
+  provisioner "local-exec" {
+    command = "gcloud builds submit --project ${var.project_id} --region ${var.region} --tag ${var.image_uri} --timeout 600s ${var.build_context}"
+  }
+}
+
+resource "google_sql_database_instance" "airflow_db" {
+  name             = "dcf-airflow-db"
+  database_version = "POSTGRES_15"
+  region           = var.region
+
+  deletion_protection = false
+
+  settings {
+    tier = "db-f1-micro"
+
+    backup_configuration {
+      enabled = false
+    }
+  }
+}
+
+resource "google_sql_database" "airflow" {
+  name     = "airflow"
+  instance = google_sql_database_instance.airflow_db.name
+}
+
+resource "google_sql_user" "airflow" {
+  name     = "airflow"
+  instance = google_sql_database_instance.airflow_db.name
+  password = var.db_password
+}
+
+resource "google_project_iam_member" "cloudsql_client" {
+  project = var.project_id
+  role    = "roles/cloudsql.client"
+  member  = "serviceAccount:${var.sa_email}"
+}
+
+resource "google_project_iam_member" "storage_viewer" {
+  project = var.project_id
+  role    = "roles/storage.objectViewer"
+  member  = "serviceAccount:${var.sa_email}"
+}
+
+locals {
+  db_conn_name = google_sql_database_instance.airflow_db.connection_name
+  db_url       = "postgresql+psycopg2://airflow:${var.db_password}@/airflow?host=/cloudsql/${local.db_conn_name}"
+}
+
+resource "google_cloud_run_v2_service" "airflow" {
+  depends_on = [null_resource.build, google_sql_database_instance.airflow_db]
+
+  name     = "dcf-airflow"
+  location = var.region
+
+  template {
+    service_account = var.sa_email
+
+    scaling {
+      min_instance_count = 1
+      max_instance_count = 1
+    }
+
+    volumes {
+      name = "dags"
+      gcs {
+        bucket    = var.warehouse_bucket
+        read_only = true
+      }
+    }
+
+    volumes {
+      name = "cloudsql"
+      cloud_sql_instance {
+        instances = [local.db_conn_name]
+      }
+    }
+
+    containers {
+      image = var.image_uri
+
+      ports {
+        container_port = 8080
+      }
+
+      volume_mounts {
+        name       = "dags"
+        mount_path = "/opt/airflow/dags"
+      }
+
+      volume_mounts {
+        name       = "cloudsql"
+        mount_path = "/cloudsql"
+      }
+
+      env {
+        name  = "AIRFLOW__DATABASE__SQL_ALCHEMY_CONN"
+        value = local.db_url
+      }
+
+      env {
+        name  = "AIRFLOW__CORE__EXECUTOR"
+        value = "LocalExecutor"
+      }
+
+      env {
+        name  = "AIRFLOW__CORE__FERNET_KEY"
+        value = var.fernet_key
+      }
+
+      env {
+        name  = "AIRFLOW__WEBSERVER__SECRET_KEY"
+        value = var.fernet_key
+      }
+
+      env {
+        name  = "AIRFLOW__CORE__LOAD_EXAMPLES"
+        value = "false"
+      }
+
+      env {
+        name  = "AIRFLOW__SCHEDULER__DAG_DIR_LIST_INTERVAL"
+        value = "30"
+      }
+
+      env {
+        name  = "_AIRFLOW_WWW_USER_CREATE"
+        value = "true"
+      }
+
+      env {
+        name  = "_AIRFLOW_WWW_USER_USERNAME"
+        value = "admin"
+      }
+
+      env {
+        name  = "_AIRFLOW_WWW_USER_PASSWORD"
+        value = var.admin_password
+      }
+
+      resources {
+        limits = {
+          memory = "2Gi"
+          cpu    = "1"
+        }
+      }
+
+      command = ["airflow", "standalone"]
+    }
+  }
+
+  lifecycle {
+    ignore_changes = [
+      template[0].containers[0].image,
+    ]
+  }
+}

dcf/infra/modules/batch_collector/gcp/airflow/outputs.tf

@@ -0,0 +1,9 @@
+output "webserver_url" {
+  description = "HTTPS URL of the Cloud Run Airflow service"
+  value       = google_cloud_run_v2_service.airflow.uri
+}
+
+output "service_name" {
+  description = "Name of the Cloud Run Airflow service"
+  value       = google_cloud_run_v2_service.airflow.name
+}

dcf/infra/modules/batch_collector/gcp/airflow/variables.tf

@@ -0,0 +1,52 @@
+variable "image_uri" {
+  type        = string
+  description = "Artifact Registry URI for the Airflow image"
+}
+
+variable "build_context" {
+  type        = string
+  description = "Absolute host path to the Airflow build context directory"
+}
+
+variable "content_hash" {
+  type        = string
+  description = "SHA256 of Airflow Dockerfile template — triggers Cloud Build rebuild"
+}
+
+variable "project_id" {
+  type        = string
+  description = "GCP project ID"
+}
+
+variable "region" {
+  type        = string
+  description = "GCP region"
+}
+
+variable "sa_email" {
+  type        = string
+  description = "Service account email for Cloud Run Airflow service"
+}
+
+variable "warehouse_bucket" {
+  type        = string
+  description = "GCS bucket where DAGs are stored at airflow/dags/"
+}
+
+variable "db_password" {
+  type        = string
+  sensitive   = true
+  description = "PostgreSQL password for Cloud SQL Airflow database"
+}
+
+variable "admin_password" {
+  type        = string
+  sensitive   = true
+  description = "Airflow webserver admin password"
+}
+
+variable "fernet_key" {
+  type        = string
+  sensitive   = true
+  description = "Airflow fernet key for encrypting connection passwords"
+}

dcf/infra/modules/batch_collector/gcp/main.tf

@@ -0,0 +1,70 @@
+terraform {
+  required_version = ">= 1.0"
+  required_providers {
+    google = {
+      source  = "hashicorp/google"
+      version = "~> 5.0"
+    }
+    local = {
+      source  = "hashicorp/local"
+      version = "~> 2.0"
+    }
+    null = {
+      source  = "hashicorp/null"
+      version = "~> 3.0"
+    }
+  }
+}
+
+provider "google" {
+  project = var.project_id
+  region  = var.region
+}
+
+resource "local_file" "dockerfile" {
+  content  = templatefile("${path.module}/templates/batch_collector.Dockerfile.tftpl", {
+    java_enabled = var.java_enabled
+  })
+  filename = "${var.build_context}/Dockerfile"
+}
+
+resource "null_resource" "build" {
+  depends_on = [local_file.dockerfile]
+
+  triggers = {
+    content_hash = var.content_hash
+  }
+
+  provisioner "local-exec" {
+    command = "gcloud builds submit --project ${var.project_id} --region ${var.region} --tag ${var.image_uri} --timeout 600s ${var.build_context}"
+  }
+}
+
+resource "google_cloud_run_v2_job" "collector" {
+  depends_on = [null_resource.build]
+
+  name     = "dcf-job-${replace(var.collector_name, "_", "-")}"
+  location = var.region
+
+  template {
+    template {
+      service_account = var.sa_email
+      max_retries     = 0
+
+      containers {
+        image = var.image_uri
+
+        env {
+          name  = "COLLECTOR_NAME"
+          value = var.collector_name
+        }
+
+        resources {
+          limits = {
+            memory = "512Mi"
+          }
+        }
+      }
+    }
+  }
+}

dcf/infra/modules/batch_collector/gcp/outputs.tf

@@ -0,0 +1,4 @@
+output "job_name" {
+  description = "Name of the provisioned Cloud Run job"
+  value       = google_cloud_run_v2_job.collector.name
+}

dcf/infra/modules/batch_collector/gcp/variables.tf

@@ -0,0 +1,40 @@
+variable "project_id" {
+  type        = string
+  description = "GCP project ID"
+}
+
+variable "region" {
+  type        = string
+  description = "GCP region"
+}
+
+variable "collector_name" {
+  type        = string
+  description = "dcf collector name (e.g. github_repos)"
+}
+
+variable "image_uri" {
+  type        = string
+  description = "Container image URI for the Cloud Run job"
+}
+
+variable "sa_email" {
+  type        = string
+  description = "Service account email for the Cloud Run job"
+}
+
+variable "build_context" {
+  type        = string
+  description = "Absolute path to the stable build context directory"
+}
+
+variable "content_hash" {
+  type        = string
+  description = "SHA256 of build context files — triggers Cloud Build rebuild when changed"
+}
+
+variable "java_enabled" {
+  type        = bool
+  default     = false
+  description = "Install OpenJDK in the container (false for GCP — uses PyArrow direct write)"
+}