Python: Data Warehouse – Incremental Loads from Postgres to BigQuery

python postgres_to_bigquery_incremental.py \
    --pg-dsn "postgresql://etl_reader@<replica-host>:5432/app" \
    --source-table public.orders \
    --project my-project \
    --dataset warehouse \
    --destination-table orders \
    --key-column id \
    --watermark-column updated_at \
    --lookback-minutes 10

2025-12-17 09:00:01 INFO Last watermark for public.orders: 2025-12-16 23:58:02+00:00 (extracting from 2025-12-16 23:48:02+00:00)
2025-12-17 09:00:03 INFO Extracted 7143 rows to /tmp/tmpk2v9x1.ndjson
2025-12-17 09:00:09 INFO Loaded 7143 rows into my-project.warehouse.orders__staging
2025-12-17 09:00:14 INFO Merge committed: 7143 staged rows, watermark now 2025-12-17 08:59:58+00:00

2025-12-17 10:00:01 INFO Last watermark for public.orders: 2025-12-17 08:59:58+00:00 (extracting from 2025-12-17 08:49:58+00:00)
2025-12-17 10:00:01 INFO No rows newer than the overlap window. Nothing to load.

# m h dom mon dow  command
*/15 * * * * flock -n /var/lock/pg-to-bq-orders.lock python /opt/etl/postgres_to_bigquery_incremental.py --pg-dsn "postgresql://etl_reader@<replica-host>:5432/app" --source-table public.orders --project my-project --dataset warehouse --destination-table orders >> /var/log/etl/orders.log 2>&1

#!/usr/bin/env python3
"""
postgres_to_bigquery_incremental.py

Incrementally copies new and changed rows from a Postgres table into a
BigQuery table, using a timestamp watermark column (e.g. updated_at).

Steps:
    1. Read the last watermark from <project>.<dataset>.etl_watermarks.
    2. Stream rows with watermark > (last watermark - lookback) out of
       Postgres through a named (server-side) cursor into a local NDJSON file.
    3. Load that file into <destination>__staging with one load job
       (WRITE_TRUNCATE, schema copied from the destination table).
    4. In one BigQuery transaction: MERGE staging into the destination on the
       key column, then advance the watermark. Any failure rolls back both.

Reads:  the Postgres source table; the destination table's schema; the
        watermark control table.
Writes: <destination>__staging (replaced every run), the destination table,
        and one row per source table in etl_watermarks.
"""

import argparse
import base64
import datetime as dt
import decimal
import json
import logging
import re
import sys
import tempfile
import uuid

import psycopg2
from google.cloud import bigquery
from psycopg2 import sql

logging.basicConfig(level=logging.INFO, format="%(asctime)s %(levelname)s %(message)s")
logger = logging.getLogger("pg_to_bq_incremental")

WATERMARK_TABLE = "etl_watermarks"
IDENTIFIER = re.compile(r"^[A-Za-z_][A-Za-z0-9_]*$")


def json_default(value):
    """Serialize the Postgres types psycopg2 returns that json.dumps cannot."""
    if isinstance(value, (dt.datetime, dt.date, dt.time)):
        return value.isoformat()
    if isinstance(value, decimal.Decimal):
        return str(value)  # keeps full precision for NUMERIC columns
    if isinstance(value, uuid.UUID):
        return str(value)
    if isinstance(value, (bytes, memoryview)):
        return base64.b64encode(bytes(value)).decode("ascii")  # BYTES must be base64 in JSON
    raise TypeError(f"Cannot serialize {type(value).__name__}")


def get_last_watermark(bq_client, project, dataset, source_table):
    """Return the stored watermark (aware datetime) for source_table, or None."""
    query = f"""
        SELECT watermark_value
        FROM `{project}.{dataset}.{WATERMARK_TABLE}`
        WHERE source_table = @source_table
    """
    job_config = bigquery.QueryJobConfig(
        query_parameters=[bigquery.ScalarQueryParameter("source_table", "STRING", source_table)]
    )
    rows = list(bq_client.query(query, job_config=job_config).result())
    return rows[0]["watermark_value"] if rows else None


def extract_to_file(pg_dsn, source_table, columns, watermark_column, since, batch_size, handle):
    """Stream matching rows into handle as NDJSON. Returns the row count."""
    schema_name, table_name = source_table.split(".", 1) if "." in source_table else ("public", source_table)
    query = sql.SQL("SELECT {cols} FROM {table}").format(
        cols=sql.SQL(", ").join(sql.Identifier(c) for c in columns),
        table=sql.Identifier(schema_name, table_name),
    )
    params = None
    if since is not None:
        query = query + sql.SQL(" WHERE {wm} > %s").format(wm=sql.Identifier(watermark_column))
        params = (since,)

    conn = psycopg2.connect(pg_dsn)
    count = 0
    try:
        # A named cursor lives inside a transaction; the with block commits it.
        with conn:
            with conn.cursor(name="pg_to_bq_extract") as cursor:
                cursor.itersize = batch_size
                cursor.execute(query, params)
                for record in cursor:
                    row = dict(zip(columns, record))
                    handle.write(json.dumps(row, default=json_default).encode("utf-8"))
                    handle.write(b"\n")
                    count += 1
    finally:
        conn.close()  # leaving "with conn" ends the transaction, not the connection
    return count


def load_staging(bq_client, handle, staging_id, schema):
    """Replace the staging table with the file contents using one load job."""
    job_config = bigquery.LoadJobConfig(
        source_format=bigquery.SourceFormat.NEWLINE_DELIMITED_JSON,
        write_disposition=bigquery.WriteDisposition.WRITE_TRUNCATE,
        schema=schema,
    )
    load_job = bq_client.load_table_from_file(handle, staging_id, rewind=True, job_config=job_config)
    load_job.result()  # raises google.api_core.exceptions.GoogleAPICallError on failure
    return load_job.output_rows


def merge_and_advance(bq_client, dest_id, staging_id, watermark_id, columns, key, watermark_column, source_table):
    """MERGE staging into the destination and advance the watermark atomically."""
    col = lambda name: f"`{name}`"
    update_set = ", ".join(f"{col(c)} = S.{col(c)}" for c in columns if c != key)
    insert_cols = ", ".join(col(c) for c in columns)
    insert_vals = ", ".join(f"S.{col(c)}" for c in columns)

    script = f"""
    BEGIN TRANSACTION;

    MERGE `{dest_id}` AS T
    USING (
      SELECT * EXCEPT (_rn) FROM (
        SELECT *, ROW_NUMBER() OVER (PARTITION BY {col(key)} ORDER BY {col(watermark_column)} DESC) AS _rn
        FROM `{staging_id}`
      )
      WHERE _rn = 1
    ) AS S
    ON T.{col(key)} = S.{col(key)}
    WHEN MATCHED AND S.{col(watermark_column)} >= T.{col(watermark_column)} THEN
      UPDATE SET {update_set}
    WHEN NOT MATCHED THEN
      INSERT ({insert_cols}) VALUES ({insert_vals});

    MERGE `{watermark_id}` AS W
    USING (
      SELECT @source_table AS source_table,
             MAX({col(watermark_column)}) AS watermark_value,
             COUNT(*) AS row_count
      FROM `{staging_id}`
    ) AS S
    ON W.source_table = S.source_table
    WHEN MATCHED THEN
      UPDATE SET watermark_value = GREATEST(W.watermark_value, S.watermark_value),
                 loaded_at = CURRENT_TIMESTAMP(),
                 row_count = S.row_count
    WHEN NOT MATCHED THEN
      INSERT (source_table, watermark_value, loaded_at, row_count)
      VALUES (S.source_table, S.watermark_value, CURRENT_TIMESTAMP(), S.row_count);

    COMMIT TRANSACTION;
    """
    job_config = bigquery.QueryJobConfig(
        query_parameters=[bigquery.ScalarQueryParameter("source_table", "STRING", source_table)]
    )
    # No exception handler: if any statement fails, BigQuery rolls the transaction back
    # and the job (and therefore .result()) fails.
    bq_client.query(script, job_config=job_config).result()


def run(args):
    for name in (args.destination_table, args.key_column, args.watermark_column):
        if not IDENTIFIER.match(name):
            raise ValueError(f"Not a plain identifier: {name!r}")

    bq_client = bigquery.Client(project=args.project)
    dest_id = f"{args.project}.{args.dataset}.{args.destination_table}"
    staging_id = f"{dest_id}__staging"
    watermark_id = f"{args.project}.{args.dataset}.{WATERMARK_TABLE}"

    destination = bq_client.get_table(dest_id)
    columns = [field.name for field in destination.schema]
    for required in (args.key_column, args.watermark_column):
        if required not in columns:
            raise ValueError(f"Column {required!r} is not in {dest_id}")

    last_watermark = get_last_watermark(bq_client, args.project, args.dataset, args.source_table)
    since = None
    if last_watermark is not None:
        since = last_watermark - dt.timedelta(minutes=args.lookback_minutes)
    logger.info("Last watermark for %s: %s (extracting from %s)", args.source_table, last_watermark, since)

    with tempfile.NamedTemporaryFile(mode="w+b", suffix=".ndjson") as handle:
        count = extract_to_file(
            args.pg_dsn, args.source_table, columns, args.watermark_column, since, args.batch_size, handle
        )
        if count == 0:
            logger.info("No rows newer than the overlap window. Nothing to load.")
            return
        logger.info("Extracted %d rows to %s", count, handle.name)

        loaded = load_staging(bq_client, handle, staging_id, destination.schema)
        logger.info("Loaded %d rows into %s", loaded, staging_id)

    merge_and_advance(
        bq_client, dest_id, staging_id, watermark_id, columns,
        args.key_column, args.watermark_column, args.source_table,
    )
    new_watermark = get_last_watermark(bq_client, args.project, args.dataset, args.source_table)
    logger.info("Merge committed: %d staged rows, watermark now %s", loaded, new_watermark)


def parse_args():
    parser = argparse.ArgumentParser(description="Incrementally load a Postgres table into BigQuery.")
    parser.add_argument("--pg-dsn", required=True, help="libpq connection string or URI for the source")
    parser.add_argument("--source-table", required=True, help="Postgres table, optionally schema-qualified (public.orders)")
    parser.add_argument("--project", required=True, help="GCP project ID")
    parser.add_argument("--dataset", required=True, help="BigQuery dataset holding the destination and etl_watermarks")
    parser.add_argument("--destination-table", required=True, help="Existing BigQuery destination table")
    parser.add_argument("--key-column", default="id", help="Primary key used by the MERGE (default: id)")
    parser.add_argument("--watermark-column", default="updated_at", help="timestamptz column that changes on every write")
    parser.add_argument("--lookback-minutes", type=int, default=10, help="Overlap re-read each run (default: 10)")
    parser.add_argument("--batch-size", type=int, default=5000, help="Rows fetched per round trip (default: 5000)")
    args = parser.parse_args()
    if args.lookback_minutes < 0:
        parser.error("--lookback-minutes must be 0 or more")
    if args.batch_size < 1:
        parser.error("--batch-size must be at least 1")
    return args


def main():
    args = parse_args()
    try:
        run(args)
    except Exception:
        logger.exception("Incremental load failed")
        sys.exit(1)


if __name__ == "__main__":
    main()