#!/usr/bin/env python3 from __future__ import annotations import argparse import csv import gzip import hashlib import io import json import re import shutil import subprocess import time import urllib.parse import urllib.request from dataclasses import dataclass from datetime import datetime, timezone from pathlib import Path from typing import Any from file_history_contract import MINIMUM_FULL_CONTENT_DOCUMENT_COUNT DEFAULT_LOGIN_DOMAIN = "login.salesforce.com" DEFAULT_API_VERSION = "v63.0" DEFAULT_EXPORT_MODE = "daily_full" DEFAULT_S3_PREFIX = "salesforce/full" DEFAULT_POLL_SECONDS = 2 DEFAULT_TIMEOUT_SECONDS = 900 DEFAULT_MANIFEST_NAME = "manifest.json" def utc_now() -> datetime: return datetime.now(timezone.utc) def utc_today() -> str: return utc_now().date().isoformat() def sha256_file(path: Path) -> str: digest = hashlib.sha256() with path.open("rb") as fh: for chunk in iter(lambda: fh.read(1024 * 1024), b""): digest.update(chunk) return digest.hexdigest() def sha256_json(value: Any) -> str: payload = json.dumps(value, sort_keys=True, separators=(",", ":")).encode("utf-8") return hashlib.sha256(payload).hexdigest() def ensure_dir(path: Path) -> None: path.mkdir(parents=True, exist_ok=True) def sanitize_path_component(value: str) -> str: cleaned = re.sub(r"[^A-Za-z0-9._-]+", "-", value.strip()) return cleaned.strip("-") or "snapshot" def selected_reference_fields(describe: dict[str, Any], selected_fields: list[str]) -> list[str]: selected_field_names = set(selected_fields) return sorted( str(field.get("name") or "").strip() for field in describe.get("fields", []) if isinstance(field, dict) and str(field.get("name") or "").strip() in selected_field_names and str(field.get("type") or "").strip().lower() == "reference" ) def load_kv_config(path: Path) -> dict[str, str]: data: dict[str, str] = {} for raw in path.read_text(encoding="utf-8").splitlines(): line = raw.strip() if not line or line.startswith("#") or ":" not in line: continue key, value = line.split(":", 1) data[key.strip()] = value.strip() return data def load_export_config(path: Path) -> dict[str, Any]: payload = json.loads(path.read_text(encoding="utf-8")) if not isinstance(payload, dict): raise ValueError("export config must be a JSON object") if not isinstance(payload.get("objects"), list) or not payload["objects"]: raise ValueError("export config must include a non-empty objects array") if "profiles" in payload and not isinstance(payload["profiles"], dict): raise ValueError("profiles must be a JSON object when provided") return payload @dataclass class SalesforceAuth: instance_url: str access_token: str issued_at: str | None @dataclass class ExportTarget: snapshot_date: str run_label: str local_dir: Path manifest_path: Path s3_uri: str | None class SalesforceClient: def __init__(self, config: dict[str, str], api_version: str) -> None: self.login_domain = config.get("tokenDomain") or DEFAULT_LOGIN_DOMAIN self.client_id = config["clientId"] self.client_secret = config["clientSecret"] self.username = config["username"] self.password = config["password"] self.api_version = api_version self.auth: SalesforceAuth | None = None def login(self) -> SalesforceAuth: token_url = f"https://{self.login_domain}/services/oauth2/token" form = urllib.parse.urlencode( { "grant_type": "password", "client_id": self.client_id, "client_secret": self.client_secret, "username": self.username, "password": self.password, } ).encode("utf-8") req = urllib.request.Request(token_url, data=form, method="POST") with urllib.request.urlopen(req, timeout=30) as resp: data = json.loads(resp.read().decode("utf-8")) self.auth = SalesforceAuth( instance_url=data["instance_url"], access_token=data["access_token"], issued_at=data.get("issued_at"), ) return self.auth def _headers(self) -> dict[str, str]: if not self.auth: self.login() assert self.auth is not None return {"Authorization": f"Bearer {self.auth.access_token}"} def _api(self, path: str) -> str: if not self.auth: self.login() assert self.auth is not None return f"{self.auth.instance_url}/services/data/{self.api_version}{path}" def _request_json(self, req: urllib.request.Request) -> dict[str, Any]: with urllib.request.urlopen(req, timeout=30) as resp: return json.loads(resp.read().decode("utf-8")) def describe(self, object_name: str) -> dict[str, Any]: req = urllib.request.Request( self._api(f"/sobjects/{object_name}/describe"), headers=self._headers(), ) return self._request_json(req) def organization_id(self) -> str: query = urllib.parse.urlencode( {"q": "SELECT Id FROM Organization LIMIT 2"} ) req = urllib.request.Request( self._api(f"/query?{query}"), headers=self._headers(), ) payload = self._request_json(req) records = payload.get("records") if not isinstance(records, list) or len(records) != 1: raise RuntimeError( "Salesforce organization admission requires exactly one Organization row" ) organization_id = str(records[0].get("Id") or "").strip() if not re.fullmatch(r"00D[A-Za-z0-9]{15}", organization_id): raise RuntimeError( "Salesforce organization admission requires a canonical 18-character Organization.Id" ) return organization_id def create_bulk_query_job(self, soql: str) -> str: req = urllib.request.Request( self._api("/jobs/query"), data=json.dumps({"operation": "query", "query": soql}).encode("utf-8"), headers={**self._headers(), "Content-Type": "application/json"}, method="POST", ) data = self._request_json(req) return str(data["id"]) def get_bulk_job(self, job_id: str) -> dict[str, Any]: req = urllib.request.Request( self._api(f"/jobs/query/{job_id}"), headers=self._headers(), ) return self._request_json(req) def _download_bulk_query_result_page(self, job_id: str, locator: str | None = None) -> tuple[bytes, str | None]: path = f"/jobs/query/{job_id}/results" if locator: encoded = urllib.parse.urlencode({"locator": locator}) path = f"{path}?{encoded}" req = urllib.request.Request( self._api(path), headers=self._headers(), ) with urllib.request.urlopen(req, timeout=300) as resp: content = resp.read() next_locator = resp.headers.get("Sforce-Locator") normalized_locator = (next_locator or "").strip() if not normalized_locator or normalized_locator.lower() == "null": normalized_locator = None return content, normalized_locator def wait_for_bulk_query_job( self, job_id: str, poll_seconds: int, timeout_seconds: int, ) -> dict[str, Any]: start = time.time() while True: payload = self.get_bulk_job(job_id) state = payload.get("state") if state == "JobComplete": return payload if state in {"Aborted", "Failed"}: raise RuntimeError(f"bulk job {job_id} ended in state={state}: {payload}") if time.time() - start > timeout_seconds: raise TimeoutError(f"bulk job {job_id} timed out after {timeout_seconds}s") time.sleep(poll_seconds) def download_bulk_query_results(self, job_id: str) -> bytes: pages: list[bytes] = [] locator: str | None = None while True: content, locator = self._download_bulk_query_result_page(job_id, locator=locator) if not pages: pages.append(content) else: pages.append(strip_csv_header_row(content)) if locator is None: break return b"".join(pages) def strip_csv_header_row(content: bytes) -> bytes: if not content: return content newline_index = content.find(b"\n") if newline_index == -1: return b"" body = content[newline_index + 1 :] return body.lstrip(b"\r") def parse_args() -> argparse.Namespace: script_dir = Path(__file__).resolve().parent parser = argparse.ArgumentParser(description="Export a Salesforce snapshot using Bulk API 2.0.") parser.add_argument("--config", required=True, help="Path to readonly Salesforce config file.") parser.add_argument( "--export-config", default=str(script_dir / "configs" / "daily_full_export.json"), help="Path to export JSON config.", ) parser.add_argument( "--output-root", required=True, help="Root directory under which dated snapshot folders are created.", ) parser.add_argument("--profile", help="Profile name from export config.") parser.add_argument("--objects", help="Comma-separated object override.") parser.add_argument("--limit", type=int, default=None, help="Optional LIMIT appended to every object query.") parser.add_argument("--field-limit", type=int, default=None, help="Override field limit for every object.") parser.add_argument("--api-version", help="Salesforce API version override.") parser.add_argument("--manifest-name", default=DEFAULT_MANIFEST_NAME, help="Manifest file name.") parser.add_argument("--snapshot-date", default=utc_today(), help="Snapshot date in YYYY-MM-DD.") parser.add_argument("--run-label", help="Optional path label under snapshot date.") parser.add_argument("--s3-bucket", help="Optional S3 bucket for upload.") parser.add_argument("--s3-prefix", help="Optional S3 prefix. Defaults to salesforce/full.") parser.add_argument("--skip-s3-upload", action="store_true", help="Skip upload even when S3 bucket is provided.") parser.add_argument("--list-profiles", action="store_true", help="List configured profiles and exit.") return parser.parse_args() def list_profiles(export_config: dict[str, Any]) -> int: profiles = export_config.get("profiles", {}) for name, members in profiles.items(): count = len(members) if isinstance(members, list) else 0 print(f"{name}\t{count}") return 0 def build_object_catalog(export_config: dict[str, Any]) -> dict[str, dict[str, Any]]: catalog: dict[str, dict[str, Any]] = {} for raw in export_config["objects"]: if not isinstance(raw, dict) or not isinstance(raw.get("name"), str): raise ValueError("each object config entry must include a string name") name = raw["name"].strip() if not name: raise ValueError("object name cannot be blank") if name in catalog: raise ValueError(f"duplicate object config for {name}") catalog[name] = raw return catalog def select_objects(args: argparse.Namespace, export_config: dict[str, Any]) -> list[dict[str, Any]]: catalog = build_object_catalog(export_config) if args.objects: selected_names = [item.strip() for item in args.objects.split(",") if item.strip()] else: profile_name = args.profile or export_config.get("default_profile") if not profile_name: raise ValueError("either --profile or --objects must be provided, or export config must set default_profile") profiles = export_config.get("profiles", {}) if profile_name not in profiles or not isinstance(profiles[profile_name], list): raise ValueError(f"profile {profile_name!r} not found in export config") selected_names = [str(item).strip() for item in profiles[profile_name] if str(item).strip()] missing = [name for name in selected_names if name not in catalog] if missing: raise ValueError(f"selected objects missing from config: {', '.join(missing)}") return [catalog[name] for name in selected_names if catalog[name].get("enabled", True)] def build_export_target(args: argparse.Namespace, export_config: dict[str, Any]) -> ExportTarget: output_root = Path(args.output_root).expanduser().resolve() snapshot_date = args.snapshot_date if args.run_label: run_label = sanitize_path_component(args.run_label) else: run_label = sanitize_path_component(args.profile or export_config.get("default_profile") or DEFAULT_EXPORT_MODE) local_dir = output_root / snapshot_date / run_label ensure_dir(local_dir) s3_uri = None if args.s3_bucket: raw_prefix = args.s3_prefix or export_config.get("default_s3_prefix") or DEFAULT_S3_PREFIX prefix = raw_prefix.strip("/") s3_uri = f"s3://{args.s3_bucket}/{prefix}/{snapshot_date}/{run_label}/" return ExportTarget( snapshot_date=snapshot_date, run_label=run_label, local_dir=local_dir, manifest_path=local_dir / args.manifest_name, s3_uri=s3_uri, ) def pick_fields(describe: dict[str, Any], object_config: dict[str, Any], fallback_field_limit: int | None) -> list[str]: fields: list[str] = [] excluded = {str(item) for item in object_config.get("exclude_fields", [])} for field in describe.get("fields", []): name = field.get("name") if not name or name in excluded: continue if field.get("deprecatedAndHidden"): continue if field.get("type") in {"address", "location"}: continue fields.append(str(name)) required_fields = [str(item) for item in object_config.get("required_fields", []) if str(item).strip()] ordered: list[str] = [] seen: set[str] = set() for name in required_fields + ["Id", "Name", "OwnerId", "CreatedDate", "LastModifiedDate", "SystemModstamp"]: if name in fields and name not in seen: ordered.append(name) seen.add(name) for name in fields: if name not in seen: ordered.append(name) seen.add(name) field_limit = fallback_field_limit or object_config.get("field_limit") or 190 return ordered[: int(field_limit)] def build_soql(object_name: str, fields: list[str], object_config: dict[str, Any], default_limit: int | None) -> str: soql = f"SELECT {', '.join(fields)} FROM {object_name}" where_clause = str(object_config.get("where", "")).strip() if where_clause: soql += f" WHERE {where_clause}" order_by = str(object_config.get("order_by", "")).strip() if order_by: soql += f" ORDER BY {order_by}" limit = default_limit if default_limit is not None else object_config.get("limit") if limit is not None: soql += f" LIMIT {int(limit)}" return soql def is_query_too_complicated_error(exc: Exception) -> bool: return "QUERY_TOO_COMPLICATED" in str(exc) def chunk_fields(fields: list[str], chunk_size: int) -> list[list[str]]: if chunk_size <= 0: raise ValueError("chunk_size must be > 0") if not fields: return [] chunks: list[list[str]] = [] current: list[str] = [] for field in fields: if field == "Id": continue current.append(field) if len(current) >= chunk_size - 1: chunks.append(["Id", *current]) current = [] if current: chunks.append(["Id", *current]) if not chunks: chunks.append(["Id"]) return chunks def load_csv_rows(content: bytes) -> tuple[list[str], list[dict[str, str]]]: decoded = content.decode("utf-8-sig") reader = csv.DictReader(decoded.splitlines()) header = list(reader.fieldnames or []) rows: list[dict[str, str]] = [] for row in reader: normalized: dict[str, str] = {} for key, value in (row or {}).items(): if key: normalized[key] = value or "" rows.append(normalized) return header, rows def dump_csv_rows(field_order: list[str], rows: list[dict[str, str]]) -> bytes: buffer = io.StringIO() writer = csv.DictWriter(buffer, fieldnames=field_order, lineterminator="\n") writer.writeheader() for row in rows: writer.writerow({field: row.get(field, "") for field in field_order}) return buffer.getvalue().encode("utf-8") def run_bulk_query(client: SalesforceClient, soql: str, poll_seconds: int, timeout_seconds: int) -> tuple[str, dict[str, Any], bytes]: job_id = client.create_bulk_query_job(soql) job_info = client.wait_for_bulk_query_job(job_id, poll_seconds=poll_seconds, timeout_seconds=timeout_seconds) content = client.download_bulk_query_results(job_id) return job_id, job_info, content def export_with_field_chunks( client: SalesforceClient, object_name: str, fields: list[str], object_config: dict[str, Any], default_limit: int | None, poll_seconds: int, timeout_seconds: int, chunk_size: int, ) -> tuple[bytes, list[dict[str, Any]], int]: candidate_sizes: list[int] = [] seen_sizes: set[int] = set() current = max(1, chunk_size) while current >= 1 and current not in seen_sizes: candidate_sizes.append(current) seen_sizes.add(current) if current == 1: break current = max(1, current // 2) last_error: Exception | None = None for candidate in candidate_sizes: try: chunks = chunk_fields(fields, candidate) merged_rows: dict[str, dict[str, str]] = {} row_order: list[str] = [] chunk_runs: list[dict[str, Any]] = [] for index, chunk_fields_list in enumerate(chunks, start=1): soql = build_soql(object_name, chunk_fields_list, object_config, default_limit) job_id, job_info, content = run_bulk_query(client, soql, poll_seconds, timeout_seconds) header, rows = load_csv_rows(content) if "Id" not in header: raise ValueError(f"{object_name} chunk {index} missing Id column") for row in rows: row_id = row.get("Id", "").strip() if not row_id: continue if row_id not in merged_rows: merged_rows[row_id] = {"Id": row_id} row_order.append(row_id) merged_rows[row_id].update(row) chunk_runs.append( { "chunk_index": index, "field_count": len(chunk_fields_list), "fields": chunk_fields_list, "job_id": job_id, "job_state": job_info.get("state"), "query": soql, } ) merged_content = dump_csv_rows(fields, [merged_rows[row_id] for row_id in row_order]) return merged_content, chunk_runs, candidate except Exception as exc: last_error = exc if not is_query_too_complicated_error(exc): raise assert last_error is not None raise last_error def write_gzip_bytes(path: Path, content: bytes) -> None: with gzip.open(path, "wb") as fh: fh.write(content) def count_csv_records(content: bytes) -> int: decoded = content.decode("utf-8-sig") reader = csv.reader(decoded.splitlines()) record_count = 0 for idx, _ in enumerate(reader): if idx == 0: continue record_count += 1 return record_count def require_minimum_record_count( object_config: dict[str, Any], record_count: int, ) -> int | None: """Block a full snapshot when a configured source-visibility floor collapses.""" raw_minimum = object_config.get("minimum_record_count") if str(object_config.get("name") or "").strip() == "ContentDocument": if raw_minimum is not None: try: declared_minimum = int(raw_minimum) except (TypeError, ValueError) as exc: raise ValueError( "ContentDocument minimum_record_count must be an integer" ) from exc if declared_minimum != MINIMUM_FULL_CONTENT_DOCUMENT_COUNT: raise ValueError( "ContentDocument minimum_record_count must match the shared " "file-history contract" ) raw_minimum = MINIMUM_FULL_CONTENT_DOCUMENT_COUNT if raw_minimum is None: return None try: minimum = int(raw_minimum) except (TypeError, ValueError) as exc: raise ValueError( f"{object_config.get('name', 'object')} minimum_record_count must be an integer" ) from exc if minimum < 0: raise ValueError( f"{object_config.get('name', 'object')} minimum_record_count must be non-negative" ) if record_count < minimum: object_name = str(object_config.get("name") or "object") raise RuntimeError( f"{object_name} export returned {record_count} rows below the " f"configured floor {minimum}; source principal visibility or " "permissions may have changed" ) return minimum def upload_directory_to_s3(local_dir: Path, s3_uri: str) -> None: aws_path = shutil.which("aws") if not aws_path: raise RuntimeError("aws CLI is required for S3 upload but was not found in PATH") subprocess.run( [aws_path, "s3", "cp", str(local_dir), s3_uri, "--recursive"], check=True, ) def main() -> int: args = parse_args() export_config_path = Path(args.export_config).expanduser().resolve() readonly_config_path = Path(args.config).expanduser().resolve() export_config = load_export_config(export_config_path) if args.list_profiles: return list_profiles(export_config) api_version = args.api_version or export_config.get("api_version") or DEFAULT_API_VERSION export_mode = export_config.get("export_mode") or DEFAULT_EXPORT_MODE selected_objects = select_objects(args, export_config) target = build_export_target(args, export_config) readonly_config = load_kv_config(readonly_config_path) client = SalesforceClient(readonly_config, api_version=str(api_version)) auth = client.login() organization_id = client.organization_id() manifest: dict[str, Any] = { "manifest_version": 1, "snapshot_time": utc_now().isoformat(), "snapshot_date": target.snapshot_date, "run_label": target.run_label, "export_mode": export_mode, "profile": args.profile or export_config.get("default_profile"), "source_org": organization_id, "source_org_kind": "salesforce_organization_id", "source_instance_url": auth.instance_url, "api_version": api_version, "config_path": str(export_config_path), "config_sha256": sha256_file(export_config_path), "object_count": len(selected_objects), "local_output_dir": str(target.local_dir), "s3_uri": target.s3_uri, "objects": [], } for object_config in selected_objects: object_name = str(object_config["name"]) describe = client.describe(object_name) fields = pick_fields(describe, object_config, args.field_limit) reference_fields = selected_reference_fields(describe, fields) poll_seconds = int(object_config.get("poll_seconds", DEFAULT_POLL_SECONDS)) timeout_seconds = int(object_config.get("timeout_seconds", DEFAULT_TIMEOUT_SECONDS)) soql = build_soql(object_name, fields, object_config, args.limit) chunk_size = int(object_config.get("query_field_chunk_size", 0) or 0) chunk_runs: list[dict[str, Any]] = [] used_chunk_size: int | None = None try: job_id, job_info, content = run_bulk_query(client, soql, poll_seconds=poll_seconds, timeout_seconds=timeout_seconds) except Exception as exc: if chunk_size <= 0 or not is_query_too_complicated_error(exc): raise content, chunk_runs, used_chunk_size = export_with_field_chunks( client, object_name, fields, object_config, args.limit, poll_seconds, timeout_seconds, chunk_size, ) job_id = chunk_runs[0]["job_id"] job_info = {"state": "JobComplete"} soql = f"chunked[{used_chunk_size}] {soql}" file_name = f"{object_name}.csv.gz" file_path = target.local_dir / file_name record_count = count_csv_records(content) minimum_record_count = require_minimum_record_count( object_config, record_count, ) write_gzip_bytes(file_path, content) manifest["objects"].append( { "object": object_name, "target_table": object_config.get("target_table"), "incremental_watermark_field": object_config.get("incremental_watermark_field"), "priority": object_config.get("priority"), "category": object_config.get("category"), "file": file_name, "record_count": record_count, "minimum_record_count": minimum_record_count, "sha256": sha256_file(file_path), "job_id": job_id, "job_state": job_info.get("state"), "query": soql, "field_count": len(fields), "fields": fields, "reference_fields": reference_fields, "chunk_runs": chunk_runs, "query_field_chunk_size_used": used_chunk_size, "config": { "required_fields": object_config.get("required_fields", []), "exclude_fields": object_config.get("exclude_fields", []), "where": object_config.get("where"), "order_by": object_config.get("order_by"), "limit": args.limit if args.limit is not None else object_config.get("limit"), "query_field_chunk_size": object_config.get("query_field_chunk_size"), "minimum_record_count": minimum_record_count, }, } ) print(f"[ok] {object_name}: file={file_path}") manifest["manifest_sha256_input"] = sha256_json( { "snapshot_date": manifest["snapshot_date"], "run_label": manifest["run_label"], "objects": manifest["objects"], "config_sha256": manifest["config_sha256"], } ) target.manifest_path.write_text(json.dumps(manifest, indent=2), encoding="utf-8") print(f"[ok] manifest={target.manifest_path}") if target.s3_uri and not args.skip_s3_upload: upload_directory_to_s3(target.local_dir, target.s3_uri) print(f"[ok] uploaded={target.s3_uri}") elif target.s3_uri: print(f"[skip] s3 upload skipped for {target.s3_uri}") return 0 if __name__ == "__main__": raise SystemExit(main())