#!/usr/bin/env python3 """Upload a CSV, create a Cust import, and report the finished job. Python 3.9+.""" import argparse import csv import io import json import os import re import sys import time import uuid from datetime import datetime, timezone from pathlib import Path from urllib.error import HTTPError, URLError from urllib.parse import urlsplit from urllib.request import HTTPRedirectHandler, Request, build_opener API_HOSTS = { "eu": "https://partner-api.cust.co", "us": "https://partner-api.us.cust.co", } class ImportErrorWithContext(Exception): """A request or import result needs attention.""" class NoRedirect(HTTPRedirectHandler): def redirect_request(self, req, fp, code, msg, headers, newurl): return None def send(request, *, api_response): # Explicitly reject redirects so credentials and signed fields stay on # the original destination. Storage receives no Cust Authorization header. try: with build_opener(NoRedirect).open(request, timeout=30) as response: body = response.read() except HTTPError as error: try: detail = error.read().decode("utf-8", errors="replace")[:2000] finally: error.close() if not api_response: detail = "Request fresh upload details if the signed upload expired." raise ImportErrorWithContext(f"HTTP {error.code}: {detail}") from error except (URLError, TimeoutError) as error: raise ImportErrorWithContext( "Request did not finish. Check Data import before resubmitting a " "create request; use --job-id to resume an existing job." ) from error if not api_response: return None try: return json.loads(body) except (ValueError, UnicodeDecodeError) as error: raise ImportErrorWithContext("The API returned an invalid JSON response.") from error def api_request(base_url, api_key, path, payload=None): data = None if payload is None else json.dumps(payload).encode("utf-8") request = Request( f"{base_url}{path}", data=data, headers={ "Authorization": f"Token {api_key}", "Accept": "application/json", "Content-Type": "application/json", }, method="GET" if payload is None else "POST", ) return send(request, api_response=True) def read_csv(path): content = path.read_bytes() if path.suffix.lower() != ".csv": raise ImportErrorWithContext("Choose a .csv file.") try: reader = csv.DictReader(io.StringIO(content.decode("utf-8-sig")), strict=True) if not reader.fieldnames or len(set(reader.fieldnames)) != len(reader.fieldnames): raise ValueError("CSV headers are missing or duplicated.") rows = list(reader) if not rows or any(None in row or None in row.values() for row in rows): raise ValueError("CSV needs data rows with the same number of columns as the header.") except (UnicodeError, csv.Error, ValueError) as error: raise ImportErrorWithContext(str(error)) from error return content def multipart_body(fields, file_name, content): boundary = f"cust-{uuid.uuid4().hex}" parts = [] for name, value in fields.items(): if any(char in name for char in '\r\n"'): raise ImportErrorWithContext("Invalid upload field name.") parts.append( f'--{boundary}\r\nContent-Disposition: form-data; name="{name}"\r\n\r\n' f'{value}\r\n'.encode("utf-8") ) if any(char in file_name for char in '\r\n"\\'): raise ImportErrorWithContext("Use a simple CSV filename without quotes or line breaks.") parts.append( f'--{boundary}\r\nContent-Disposition: form-data; name="file"; ' f'filename="{file_name}"\r\nContent-Type: text/csv\r\n\r\n'.encode("utf-8") ) parts.extend([content, f"\r\n--{boundary}--\r\n".encode("utf-8")]) return b"".join(parts), f"multipart/form-data; boundary={boundary}" def create_import(base_url, api_key, path, resource_type): content = read_csv(path) upload = api_request(base_url, api_key, "/v1/file-uploads/", { "file_name": path.name, "content_type": "text/csv", }) expires_at = datetime.fromisoformat(upload["expires_at"].replace("Z", "+00:00")) if expires_at <= datetime.now(timezone.utc): raise ImportErrorWithContext("Upload details expired. Start again to request fresh details.") # A production HTTPS API must return an HTTPS storage destination. if urlsplit(base_url).scheme == "https" and urlsplit(upload["upload_url"]).scheme != "https": raise ImportErrorWithContext("The storage upload URL must use HTTPS.") body, content_type = multipart_body(upload["upload_fields"], path.name, content) send(Request( upload["upload_url"], data=body, headers={"Content-Type": content_type}, method="POST", ), api_response=False) return api_request(base_url, api_key, "/v1/bulk-imports/", { "resource_type": resource_type, "source_file": upload["id"], }) def wait_for_import(base_url, api_key, job_id, timeout, interval=5): deadline = time.monotonic() + timeout while True: job = api_request(base_url, api_key, f"/v1/bulk-imports/{job_id}/") status = job.get("status") if status in {"completed", "failed"}: return job if status not in {"pending", "processing"}: raise ImportErrorWithContext(f"Unexpected job status: {status!r}") if time.monotonic() >= deadline: raise ImportErrorWithContext( f"Job {job_id} is still {status}. Resume with --job-id {job_id}." ) time.sleep(min(interval, max(0, deadline - time.monotonic()))) def result_is_clean(job): # The processing lifecycle can be completed even when rows failed validation. message = job.get("result_message", "") return ( job.get("status") == "completed" and message.startswith("Import completed successfully.") and not re.search(r"Records (?:skipped|with errors): [1-9]\d*", message) ) def main(): parser = argparse.ArgumentParser(description=__doc__) parser.add_argument("--region", choices=API_HOSTS, required=True) source = parser.add_mutually_exclusive_group(required=True) source.add_argument("--csv", type=Path, help="CSV file to upload and import") source.add_argument("--job-id", type=int, help="Resume polling an existing job") parser.add_argument("--resource", default="company", choices=[ "company", "contact", "metric_data", "metric_snapshot", ]) parser.add_argument("--timeout", type=int, default=300, help="Polling limit in seconds") args = parser.parse_args() api_key = os.environ.get("CUST_API_KEY", "").strip() if not api_key: parser.error("Set CUST_API_KEY in your environment.") if args.timeout <= 0 or (args.job_id is not None and args.job_id <= 0): parser.error("Timeout and job ID must be positive.") base_url = API_HOSTS[args.region] try: job_id = args.job_id if args.csv: job_id = create_import(base_url, api_key, args.csv, args.resource)["id"] print(f"Polling import job {job_id}", file=sys.stderr) job = wait_for_import(base_url, api_key, job_id, args.timeout) print(json.dumps(job, indent=2)) if not result_is_clean(job): print("Review result_message, correct failed or skipped rows, and verify imported records.", file=sys.stderr) return 2 print("Processing finished. Open a sample record in Cust and verify its fields.", file=sys.stderr) return 0 except (ImportErrorWithContext, OSError, KeyError, TypeError, ValueError) as error: print(f"Import stopped: {error}", file=sys.stderr) return 1 if __name__ == "__main__": sys.exit(main())