Postgres is Crying: A Zero-Downtime, 15TB Migration to BigQuery with Python & Arrow

              +--------------------------------+
              |                                |
              |      PostgreSQL Primary DB     |
              |          (15 TB, on fire)      |
              |                                |
              +----------------+---------------+
                               |
      +------------------------+--------------------------+
      |                                                   |
      | (Phase 1: Bulk Export)                            | (Phase 2: CDC Stream)
      v                                                   v
+-----+------------------+                     +------------------------+
| Python ETL Worker      |                     |   CDC Listener         |
| (Streaming with Arrow) |                     |   (e.g., Debezium,     |
+-----+------------------+                     |    Logical Replication)|
      |                                        +----------+-------------+
      | Chunks of data                                    |
      | as Parquet files                                  | Live changes
      v                                                   v
+-----+------------------+                     +------------------------+
| Google Cloud Storage   |                     |   Pub/Sub or Kafka     |
| (Staging Area)         |                     +------------------------+
+-----+------------------+                                |
      |                                                   |
      | Load Job                                          v
      v                                        +------------------------+
+-----+--------------------------------------+ | Cloud Function/Stream  |
|                                            | | Processor              |
|             Google BigQuery                +-> (Applies changes)    |
|                                            | +------------------------+
|                                            |
+--------------------------------------------+

# requirements: psycopg[binary], pyarrow, google-cloud-storage, google-cloud-bigquery
import os
import uuid
from contextlib import contextmanager
from datetime import datetime, timezone

import pyarrow as pa
import pyarrow.parquet as pq
import psycopg
from google.cloud import bigquery, storage

# --- Configuration ---
# Use env variables in production!
PG_CONN_STRING = "postgresql://user:password@host:port/dbname"
GCS_BUCKET_NAME = "your-bq-staging-bucket"
GCS_PREFIX = f"pg_migration/{datetime.now(timezone.utc).strftime('%Y-%m-%d')}"
BQ_DATASET_ID = "your_dataset"
BQ_TABLE_ID = "your_migrated_table"
CHUNK_SIZE = 500_000  # Number of rows to process at a time

# --- Database & Cloud Clients ---
# In a real app, manage these clients more robustly (e.g., singletons)
storage_client = storage.Client()
bq_client = bigquery.Client()
gcs_bucket = storage_client.bucket(GCS_BUCKET_NAME)

@contextmanager
def get_server_side_cursor(table_name: str, columns: list[str]):
    """
    Creates a server-side cursor to stream data from Postgres without
    loading it all into the client's memory. This is CRITICAL.
    """
    conn = psycopg.connect(PG_CONN_STRING, autocommit=True)
    # A server-side cursor is named. The client only ever holds a small buffer.
    cursor_name = f"migration_cursor_{uuid.uuid4().hex}"
    cursor = conn.cursor(name=cursor_name)
    cursor.itersize = CHUNK_SIZE  # How many rows psycopg fetches at once internally

    print(f"INFO: Starting stream from table {table_name}...")
    try:
        column_str = ", ".join(f'"{c}"' for c in columns)
        cursor.execute(f"SELECT {column_str} FROM {table_name}")
        yield cursor
    finally:
        print("INFO: Closing connection.")
        cursor.close()
        conn.close()

def upload_to_gcs(file_path: str, gcs_path: str):
    """Uploads a local file to GCS."""
    blob = gcs_bucket.blob(gcs_path)
    blob.upload_from_filename(file_path)
    print(f"INFO: Uploaded {file_path} to gs://{GCS_BUCKET_NAME}/{gcs_path}")

def stream_pg_to_gcs_parquet(table_name: str, columns: list[str]):
    """
    The main orchestration function.
    - Streams from Postgres using a server-side cursor.
    - Converts chunks to Arrow RecordBatches.
    - Writes to local Parquet files.
    - Uploads to GCS.
    """
    gcs_paths = []
    with get_server_side_cursor(table_name, columns) as cursor:
        chunk_num = 0
        while True:
            # fetchmany is the key to streaming.
            rows = cursor.fetchmany(CHUNK_SIZE)
            if not rows:
                break

            chunk_num += 1
            print(f"INFO: Processing chunk {chunk_num} with {len(rows)} rows...")

            # Zero-copy conversion to an Arrow Table
            # PyArrow can often directly map Postgres memory to Arrow memory.
            try:
                record_batch = pa.RecordBatch.from_pylist(
                    [dict(zip(columns, row)) for row in rows]
                )
                arrow_table = pa.Table.from_batches([record_batch])
            except Exception as e:
                print(f"ERROR: Could not convert chunk to Arrow: {e}")
                # Add dead-letter queue logic here
                continue

            # Write to a local Parquet file
            local_filename = f"/tmp/{table_name}_chunk_{chunk_num}.parquet"
            pq.write_table(arrow_table, local_filename, compression='SNAPPY')

            # Upload to GCS
            gcs_path = f"{GCS_PREFIX}/{table_name}/part-{chunk_num:05d}.parquet"
            upload_to_gcs(local_filename, gcs_path)
            gcs_paths.append(f"gs://{GCS_BUCKET_NAME}/{gcs_path}")

            os.remove(local_filename)

    print("INFO: Stream finished.")
    return gcs_paths

def load_gcs_to_bigquery(gcs_paths: list[str]):
    """Kicks off a BigQuery load job from GCS."""
    job_config = bigquery.LoadJobConfig(
        source_format=bigquery.SourceFormat.PARQUET,
        write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE, # Overwrite table
    )

    table_ref = f"{bq_client.project}.{BQ_DATASET_ID}.{BQ_TABLE_ID}"
    load_job = bq_client.load_table_from_uri(gcs_paths, table_ref, job_config=job_config)
    print(f"INFO: Starting BigQuery load job {load_job.job_id}")

    load_job.result() # Waits for the job to complete.
    print("INFO: BigQuery load job finished.")

if __name__ == "__main__":
    # In a real scenario, you'd get these columns from information_schema
    # and handle data type mapping carefully!
    table_columns = ["id", "user_id", "event_type", "event_timestamp", "payload"]

    gcs_file_paths = stream_pg_to_gcs_parquet("events", table_columns)

    if gcs_file_paths:
        load_gcs_to_bigquery(gcs_file_paths)

# A simplified validation script
def validate_data_integrity(pg_table, bq_table, column):
    """Compares checksums between Postgres and BigQuery."""

    # --- Postgres Checksum ---
    with psycopg.connect(PG_CONN_STRING) as conn:
        with conn.cursor() as cur:
            # BIT_XOR is a great way to 'fingerprint' a numeric column
            cur.execute(f"SELECT BIT_XOR(CAST({column} AS BIGINT)) FROM {pg_table};")
            pg_checksum = cur.fetchone()[0]
            cur.execute(f"SELECT COUNT(*) FROM {pg_table};")
            pg_count = cur.fetchone()[0]

    # --- BigQuery Checksum ---
    bq_client = bigquery.Client()
    # BigQuery's equivalent for BIT_XOR
    query = f"""
        SELECT BIT_XOR(CAST({column} AS INT64)) as checksum, COUNT(*) as count
        FROM `{bq_table}`
    """
    query_job = bq_client.query(query)
    results = query_job.result()
    for row in results:
        bq_checksum = row.checksum
        bq_count = row.count

    print("--- DATA INTEGRITY CHECK ---")
    print(f"Postgres Count: {pg_count}, BigQuery Count: {bq_count}")
    print(f"Postgres XOR:   {pg_checksum}, BigQuery XOR:   {bq_checksum}")

    if pg_count == bq_count and pg_checksum == bq_checksum:
        print("✅ SUCCESS: Data is consistent!")
        return True
    else:
        print("❌ FAILURE: Data mismatch detected!")
        return False

# Run this right before the final cutover
# validate_data_integrity("public.events", "your_project.your_dataset.your_migrated_table", "id")

┌────────────────────────────────────────────────────────┐
│        🚨 INCIDENT: MIGRATION SLOW OR FAILING?         │
└──────────────────────────┬─────────────────────────────┘
                           │
             ┌─────────────┴─────────────┐
             ▼                           ▼
  [Check Source Postgres]     [Check Python Worker]
  • High CPU / Disk I/O?      • High RAM / OOMKilled?
    ├─ YES ──► Scale replica    ├─ YES ──► Stream with PyArrow
    └─ NO  ──► Source is fine   └─ NO  ──► Worker is fine
             │                           │
             ▼                           ▼
     [Check Network]             [Check BigQuery Load]
  • Low throughput?           • Schema validation errors?
    ├─ YES ──► Snappy Parquet   ├─ YES ──► Define explicit schema
    └─ NO  ──► Network is fine  └─ NO  ──► Pipeline is optimal!