diff --git a/AGENT.md b/AGENT.md index c4aa8ce..dd0840b 100644 --- a/AGENT.md +++ b/AGENT.md @@ -24,7 +24,11 @@ Prefer fixup commits over amending and force-pushing. ## Maintenance Tooling -`tools/justfile` contains DuckLake maintenance recipes (expire snapshots, cleanup files, delete orphans, compaction). It is copied to `/justfile` in the Docker image and uses the pod's existing env vars. The `DUCKDB` env var can override the duckdb binary path. +`tools/maintenance.py` is a self-contained Python script for DuckLake maintenance operations (snapshot expiry, file cleanup, orphan deletion, checkpoint). It connects to DuckLake using the same env vars as the main app (`DUCKLAKE_RDS_*`, `DUCKDB_S3_*`, `DUCKLAKE_DATA_PATH`) but does not import from the `millpond` package. Designed to run as a K8s CronJob reusing the same Docker image. + +If `PUSHGATEWAY_URL` is set, the script pushes two metrics: `maintenance_start_time{operation}` (pushed immediately on start) and `maintenance_duration_seconds{operation, status}` (pushed on completion). This enables Grafana annotation queries for maintenance windows. + +`tools/justfile` wraps the script for interactive use and is copied to `/justfile` in the Docker image. The `shell` and `drop` recipes still use the DuckDB CLI directly. The `DUCKDB` env var can override the duckdb binary path. ## What This Is @@ -375,7 +379,9 @@ millpond/ │ ├── metrics.py # Prometheus metric definitions │ └── server.py # HTTP server for /metrics and /healthz ├── tools/ -│ └── sizing-calculator.html # Interactive flush/object sizing calculator +│ ├── maintenance.py # Self-contained DuckLake maintenance script (K8s CronJob) +│ ├── justfile # Interactive wrapper for maintenance.py + DuckDB CLI shell +│ └── sizing-calculator.html # Interactive flush/object sizing calculator ├── tests/ │ ├── unit/ # Fast, no external deps │ ├── integration/ # Local DuckDB write path + schema evolution diff --git a/Dockerfile b/Dockerfile index 927f79e..4acca97 100644 --- a/Dockerfile +++ b/Dockerfile @@ -23,6 +23,7 @@ COPY --from=builder /app/.venv /app/.venv COPY --from=builder /app/millpond /app/millpond COPY --from=builder /root/.local/bin/duckdb /usr/local/bin/duckdb COPY tools/justfile /justfile +COPY tools/maintenance.py /app/tools/maintenance.py ENV PATH="/app/.venv/bin:$PATH" diff --git a/README.md b/README.md index 7225ebd..72ff3e6 100644 --- a/README.md +++ b/README.md @@ -83,7 +83,20 @@ Requires Docker (uses `keytool` from the Kafka container image for cert generati ### DuckLake Maintenance -`tools/justfile` contains DuckLake maintenance recipes and is baked into the Docker image at `/justfile`. From inside a running pod: +`tools/maintenance.py` is a self-contained Python script for DuckLake maintenance operations (snapshot expiry, file cleanup, orphan deletion, checkpoint). It is baked into the Docker image at `/app/tools/maintenance.py` and designed to run as a K8s CronJob reusing the same image and credentials as the main application. + +```bash +python /app/tools/maintenance.py maintain --days 7 # expire snapshots + cleanup files +python /app/tools/maintenance.py maintain --days 7 --dry-run # preview only +python /app/tools/maintenance.py expire --days 3 # expire snapshots only +python /app/tools/maintenance.py cleanup --days 1 # cleanup scheduled files only +python /app/tools/maintenance.py checkpoint # integrated merge + expire + cleanup +python /app/tools/maintenance.py orphans # delete orphaned S3 files +``` + +If `PUSHGATEWAY_URL` is set, the script pushes `maintenance_start_time` (on start) and `maintenance_duration_seconds` (on completion) to a Prometheus Pushgateway, enabling Grafana annotation queries for maintenance windows. + +`tools/justfile` wraps the script and is also baked into the image at `/justfile` for interactive use: ```bash just --list # see available recipes @@ -94,7 +107,7 @@ just drop events # drop a table (data files remain until cleanup) just orphans-dry-run # preview orphaned S3 files ``` -All recipes use the pod's existing env vars (`DUCKLAKE_RDS_*`, `DUCKDB_S3_*`, `DUCKLAKE_DATA_PATH`). +All commands use the pod's existing env vars (`DUCKLAKE_RDS_*`, `DUCKDB_S3_*`, `DUCKLAKE_DATA_PATH`). ## Configuration diff --git a/tools/justfile b/tools/justfile index 996ad0a..bd4cffc 100644 --- a/tools/justfile +++ b/tools/justfile @@ -1,4 +1,5 @@ # DuckLake maintenance tasks +# Delegates to maintenance.py for actual operations. # Requires: DUCKLAKE_RDS_PASSWORD, DUCKDB_S3_ACCESS_KEY_ID, DUCKDB_S3_SECRET_ACCESS_KEY duckdb := env("DUCKDB", "duckdb") @@ -13,7 +14,9 @@ s3_key := env("DUCKDB_S3_ACCESS_KEY_ID", "") s3_secret := env("DUCKDB_S3_SECRET_ACCESS_KEY", "") data_path := env("DUCKLAKE_DATA_PATH", "") -_setup := "LOAD httpfs; LOAD ducklake; LOAD postgres; SET s3_region = '" + s3_region + "'; SET s3_endpoint = 's3." + s3_region + ".amazonaws.com'; SET s3_access_key_id = '" + s3_key + "'; SET s3_secret_access_key = '" + s3_secret + "'; SET s3_use_ssl = true; SET s3_url_style = 'vhost'; ATTACH 'postgres:host=" + rds_host + " port=" + rds_port + " dbname=" + rds_db + " user=" + rds_user + " password=" + rds_password + "' AS lake (TYPE ducklake, DATA_PATH '" + data_path + "'); USE lake;" +_setup := "INSTALL httpfs; INSTALL ducklake; INSTALL postgres; LOAD httpfs; LOAD ducklake; LOAD postgres; SET s3_region = '" + s3_region + "'; SET s3_endpoint = 's3." + s3_region + ".amazonaws.com'; SET s3_access_key_id = '" + s3_key + "'; SET s3_secret_access_key = '" + s3_secret + "'; SET s3_use_ssl = true; SET s3_url_style = 'vhost'; ATTACH 'postgres:host=" + rds_host + " port=" + rds_port + " dbname=" + rds_db + " user=" + rds_user + " password=" + rds_password + "' AS lake (TYPE ducklake, DATA_PATH '" + data_path + "'); USE lake;" + +maintenance := env("MAINTENANCE_SCRIPT", "/app/tools/maintenance.py") default: @just --list @@ -28,42 +31,40 @@ drop table: # Expire snapshots (dry run) expire-dry-run days="7": - {{ duckdb }} -c "{{ _setup }} CALL ducklake_expire_snapshots('lake', older_than => now() - INTERVAL '{{ days }} days', dry_run => true);" + python {{ maintenance }} expire --days {{ days }} --dry-run # Expire snapshots expire days="7": - {{ duckdb }} -c "{{ _setup }} CALL ducklake_expire_snapshots('lake', older_than => now() - INTERVAL '{{ days }} days');" + python {{ maintenance }} expire --days {{ days }} # Preview files scheduled for deletion cleanup-dry-run days="1": - {{ duckdb }} -c "{{ _setup }} CALL ducklake_cleanup_old_files('lake', older_than => now() - INTERVAL '{{ days }} days', dry_run => true);" + python {{ maintenance }} cleanup --days {{ days }} --dry-run # Delete files scheduled for deletion cleanup days="1": - {{ duckdb }} -c "{{ _setup }} CALL ducklake_cleanup_old_files('lake', older_than => now() - INTERVAL '{{ days }} days');" + python {{ maintenance }} cleanup --days {{ days }} # Delete all files scheduled for deletion regardless of age cleanup-all: - {{ duckdb }} -c "{{ _setup }} CALL ducklake_cleanup_old_files('lake', cleanup_all => true);" + python {{ maintenance }} cleanup-all # Preview orphaned files orphans-dry-run: - {{ duckdb }} -c "{{ _setup }} CALL ducklake_delete_orphaned_files('lake', dry_run => true);" + python {{ maintenance }} orphans --dry-run # Delete orphaned files orphans: - {{ duckdb }} -c "{{ _setup }} CALL ducklake_delete_orphaned_files('lake');" + python {{ maintenance }} orphans # Full maintenance: expire snapshots + cleanup files maintain days="7": - just expire {{ days }} - just cleanup {{ days }} + python {{ maintenance }} maintain --days {{ days }} # Full maintenance dry run maintain-dry-run days="7": - just expire-dry-run {{ days }} - just cleanup-dry-run {{ days }} + python {{ maintenance }} maintain --days {{ days }} --dry-run # Run CHECKPOINT (integrated merge + expire + cleanup) checkpoint: - {{ duckdb }} -c "{{ _setup }} CHECKPOINT lake;" + python {{ maintenance }} checkpoint diff --git a/tools/maintenance.py b/tools/maintenance.py new file mode 100644 index 0000000..fbc747d --- /dev/null +++ b/tools/maintenance.py @@ -0,0 +1,292 @@ +#!/usr/bin/env python3 +"""DuckLake maintenance operations. + +Standalone script for running DuckLake maintenance tasks (snapshot expiry, +file cleanup, orphan deletion, checkpoint) from a K8s CronJob or manually. + +Requires the same env vars as the tools/justfile: + DUCKLAKE_RDS_HOST, DUCKLAKE_RDS_PORT, DUCKLAKE_RDS_DATABASE, + DUCKLAKE_RDS_USERNAME, DUCKLAKE_RDS_PASSWORD, DUCKLAKE_DATA_PATH, + DUCKDB_S3_REGION, DUCKDB_S3_ACCESS_KEY_ID, DUCKDB_S3_SECRET_ACCESS_KEY + (plus optional DUCKDB_S3_ENDPOINT, DUCKDB_S3_USE_SSL, DUCKDB_S3_URL_STYLE) + +Optional: + PUSHGATEWAY_URL — Prometheus Pushgateway address (e.g. http://pushgateway:9091). + If unset, metrics are not pushed and the script runs without instrumentation. +""" + +from __future__ import annotations + +import argparse +import logging +import os +import re +import sys +import time + +import duckdb +from prometheus_client import CollectorRegistry, Gauge, push_to_gateway + +log = logging.getLogger("maintenance") + +_SETTING_VALUE_RE = re.compile(r"^[a-zA-Z0-9_.:/\-@+=]+$") + + +def _setup_logging(verbose: bool = False) -> None: + level = "DEBUG" if verbose else os.environ.get("LOG_LEVEL", "INFO").upper() + logging.basicConfig( + level=level, + format="%(asctime)s %(levelname)-5s [%(name)s] %(message)s", + datefmt="%Y-%m-%dT%H:%M:%S", + stream=sys.stderr, + ) + + +def _require(name: str) -> str: + val = os.environ.get(name) + if not val: + raise RuntimeError(f"Required environment variable {name} is not set") + return val + + +def _escape_libpq(value: str | None) -> str: + if value is None: + return "''" + escaped = value.replace("\\", "\\\\").replace("'", "\\'") + return f"'{escaped}'" + + +def _sanitize_setting_value(val: str) -> str: + if not _SETTING_VALUE_RE.match(val): + raise ValueError(f"Illegal character in DuckDB setting value: {val!r}") + return val + + +def connect() -> duckdb.DuckDBPyConnection: + """Connect to DuckLake using environment variables.""" + conn = duckdb.connect() + + # S3 config from env vars + s3_region = os.environ.get("DUCKDB_S3_REGION", "us-east-1") + s3_defaults = { + "s3_region": s3_region, + "s3_endpoint": f"s3.{s3_region}.amazonaws.com", + "s3_use_ssl": "true", + "s3_url_style": "vhost", + } + for key in ( + "DUCKDB_S3_ENDPOINT", + "DUCKDB_S3_ACCESS_KEY_ID", + "DUCKDB_S3_SECRET_ACCESS_KEY", + "DUCKDB_S3_USE_SSL", + "DUCKDB_S3_URL_STYLE", + "DUCKDB_S3_REGION", + ): + val = os.environ.get(key) + setting = key.lower().replace("duckdb_", "") + if val is not None: + conn.execute(f"SET {setting} = '{_sanitize_setting_value(val)}'") + elif setting in s3_defaults: + conn.execute(f"SET {setting} = '{_sanitize_setting_value(s3_defaults[setting])}'") + + conn.execute("LOAD httpfs") + conn.execute("LOAD ducklake") + conn.execute("LOAD postgres") + + rds_host = _require("DUCKLAKE_RDS_HOST") + rds_port = os.environ.get("DUCKLAKE_RDS_PORT", "5432") + rds_database = os.environ.get("DUCKLAKE_RDS_DATABASE", "ducklake") + rds_username = os.environ.get("DUCKLAKE_RDS_USERNAME", "ducklake") + rds_password = _require("DUCKLAKE_RDS_PASSWORD") + data_path = _require("DUCKLAKE_DATA_PATH") + + pg_connstr = ( + f"host={rds_host} port={rds_port} " + f"dbname={_escape_libpq(rds_database)} user={_escape_libpq(rds_username)} " + f"password={_escape_libpq(rds_password)}" + ) + pg_connstr_sql = pg_connstr.replace("'", "''") + conn.execute(f""" + ATTACH 'ducklake:postgres:{pg_connstr_sql}' AS lake ( + DATA_PATH '{data_path.replace("'", "''")}' + ) + """) + + log.info( + "Connected: metadata=%s:%s/%s data=%s", + rds_host, + rds_port, + rds_database, + data_path, + ) + return conn + + +def expire(conn: duckdb.DuckDBPyConnection, days: int, dry_run: bool) -> None: + """Expire snapshots older than N days.""" + log.info("Expiring snapshots older than %d days (dry_run=%s)", days, dry_run) + result = conn.execute( + f"CALL ducklake_expire_snapshots('lake', " + f"older_than => now() - INTERVAL '{days} days', " + f"dry_run => {str(dry_run).lower()})" + ).fetchall() + for row in result: + log.info("expire: %s", row) + + +def cleanup(conn: duckdb.DuckDBPyConnection, days: int, dry_run: bool) -> None: + """Delete files scheduled for deletion older than N days.""" + log.info("Cleaning up files older than %d days (dry_run=%s)", days, dry_run) + result = conn.execute( + f"CALL ducklake_cleanup_old_files('lake', " + f"older_than => now() - INTERVAL '{days} days', " + f"dry_run => {str(dry_run).lower()})" + ).fetchall() + for row in result: + log.info("cleanup: %s", row) + + +def cleanup_all(conn: duckdb.DuckDBPyConnection, dry_run: bool) -> None: + """Delete all files scheduled for deletion regardless of age.""" + if dry_run: + log.info("cleanup-all has no dry-run mode; skipping") + return + log.info("Cleaning up all files scheduled for deletion") + result = conn.execute( + "CALL ducklake_cleanup_old_files('lake', cleanup_all => true)" + ).fetchall() + for row in result: + log.info("cleanup-all: %s", row) + + +def orphans(conn: duckdb.DuckDBPyConnection, dry_run: bool) -> None: + """Find and delete orphaned S3 files.""" + log.info("Deleting orphaned files (dry_run=%s)", dry_run) + result = conn.execute( + f"CALL ducklake_delete_orphaned_files('lake', " + f"dry_run => {str(dry_run).lower()})" + ).fetchall() + for row in result: + log.info("orphans: %s", row) + + +def checkpoint(conn: duckdb.DuckDBPyConnection) -> None: + """Run CHECKPOINT (integrated merge + expire + cleanup).""" + log.info("Running CHECKPOINT") + conn.execute("CHECKPOINT lake") + log.info("CHECKPOINT complete") + + +def maintain(conn: duckdb.DuckDBPyConnection, days: int, dry_run: bool) -> None: + """Full maintenance: expire snapshots then cleanup files.""" + expire(conn, days, dry_run) + cleanup(conn, days, dry_run) + + +def build_parser() -> argparse.ArgumentParser: + parser = argparse.ArgumentParser( + description="DuckLake maintenance operations", + formatter_class=argparse.RawDescriptionHelpFormatter, + ) + parser.add_argument("--verbose", action="store_true", help="Debug logging") + + sub = parser.add_subparsers(dest="command", required=True) + + # expire + p = sub.add_parser("expire", help="Expire snapshots older than N days") + p.add_argument("--days", type=int, default=7) + p.add_argument("--dry-run", action="store_true") + + # cleanup + p = sub.add_parser("cleanup", help="Delete files scheduled for deletion") + p.add_argument("--days", type=int, default=1) + p.add_argument("--dry-run", action="store_true") + + # cleanup-all + p = sub.add_parser("cleanup-all", help="Delete all scheduled files") + p.add_argument("--dry-run", action="store_true") + + # orphans + p = sub.add_parser("orphans", help="Delete orphaned S3 files") + p.add_argument("--dry-run", action="store_true") + + # maintain + p = sub.add_parser("maintain", help="Full maintenance (expire + cleanup)") + p.add_argument("--days", type=int, default=7) + p.add_argument("--dry-run", action="store_true") + + # checkpoint + sub.add_parser("checkpoint", help="CHECKPOINT (merge + expire + cleanup)") + + return parser + + +def _push(registry: CollectorRegistry, pushgateway: str) -> None: + """Push metrics to the pushgateway, logging but not raising on failure.""" + try: + push_to_gateway(pushgateway, job="ducklake-maintenance", registry=registry) + except Exception: + log.exception("Failed to push metrics to %s", pushgateway) + + +def main(argv: list[str] | None = None) -> None: + parser = build_parser() + args = parser.parse_args(argv) + _setup_logging(args.verbose) + + pushgateway = os.environ.get("PUSHGATEWAY_URL") + registry = CollectorRegistry() + start_time = Gauge( + "maintenance_start_time", + "Unix timestamp when the maintenance operation started", + ["operation"], + registry=registry, + ) + duration = Gauge( + "maintenance_duration_seconds", + "Duration of the maintenance operation", + ["operation", "status"], + registry=registry, + ) + operation = args.command + if hasattr(args, "days") and args.days < 1: + parser.error("--days must be >= 1") + + start_time.labels(operation=operation).set(time.time()) + if pushgateway: + _push(registry, pushgateway) + + t0 = time.monotonic() + status = "success" + conn = None + conn = connect() + try: + match args.command: + case "expire": + expire(conn, args.days, args.dry_run) + case "cleanup": + cleanup(conn, args.days, args.dry_run) + case "cleanup-all": + cleanup_all(conn, args.dry_run) + case "orphans": + orphans(conn, args.dry_run) + case "maintain": + maintain(conn, args.days, args.dry_run) + case "checkpoint": + checkpoint(conn) + except Exception: + status = "error" + log.exception("Maintenance operation %s failed", operation) + raise + finally: + elapsed = time.monotonic() - t0 + duration.labels(operation=operation, status=status).set(elapsed) + if conn is not None: + conn.close() + if pushgateway: + _push(registry, pushgateway) + log.info("Operation %s finished: status=%s duration=%.1fs", operation, status, elapsed) + + +if __name__ == "__main__": + main()