client.ts

#!/usr/bin/env bun

const startTime = performance.now();

import { mkdir, readFile, rm, stat, writeFile } from "node:fs/promises";
import { homedir, tmpdir } from "node:os";
import { join } from "node:path";
import { parseArgs } from "node:util";
import z from "zod";

const GITHUB_RAW_URL =
  "https://raw.githubusercontent.com/futuritywork/pg-dump-cache/main/client.ts";
const CONFIG_DIR = join(homedir(), ".config", "pg-dump-cache");
const UPDATE_CHECK_FILE = join(CONFIG_DIR, "update-checked");
const UPDATE_CHECK_INTERVAL = 24 * 60 * 60 * 1000; // 24 hours in ms
const NO_UPDATE_CHECK = process.env.DUMP_CACHE_NO_UPDATE_CHECK === "1";

const CACHE_SERVER_URL = new URL(
  process.env.CACHE_SERVER_URL ?? "http://localhost:3000",
);
const LOCAL_DB_URL = process.env.DATABASE_URL;
const API_KEY = process.env.API_KEY;

if (!API_KEY) {
  console.error("API_KEY environment variable is required");
  process.exit(1);
}

const headers = { Authorization: `Bearer ${API_KEY}` };

const { values } = parseArgs({
  options: {
    wait: { type: "boolean", short: "w", default: false },
    fresh: { type: "boolean", short: "f", default: false },
    status: { type: "boolean", short: "s", default: false },
    help: { type: "boolean", short: "h", default: false },
    "no-update-check": { type: "boolean", default: false },
    "force-update": { type: "boolean", short: "u", default: false },
  },
});

if (values.help) {
  console.log(`
pg-dump-cache client - Fetch and restore PostgreSQL dumps

Usage: ./client.ts [options]

Options:
  -w, --wait           Wait for fresh dump if cache is stale
  -f, --fresh          Force refresh before fetching
  -s, --status         Show server status and exit
  -h, --help           Show this help
  -u, --force-update   Force update check (ignore 24h TTL)
  --no-update-check    Disable auto-update check for this invocation

Environment variables:
  CACHE_SERVER_URL           Server URL (default: http://localhost:3000)
  API_KEY                    Shared API key for authentication (required)
  LOCAL_DB_URL               Local PostgreSQL connection string (required for restore)
  DUMP_CACHE_NO_UPDATE_CHECK Set to 1 to disable auto-update checks
`);
  process.exit(0);
}

const Z_Status = z.object({
  hasCache: z.boolean(),
  cacheTimestamp: z.string().nullable(),
  cacheAgeSeconds: z.number().nullable(),
  refreshing: z.boolean(),
  ttl: z.number(),
});

async function getStatus(): Promise<{
  hasCache: boolean;
  cacheTimestamp: string | null;
  cacheAgeSeconds: number | null;
  refreshing: boolean;
  ttl: number;
}> {
  const res = await fetch(`${CACHE_SERVER_URL}status`, { headers });
  if (!res.ok) {
    throw new Error(`Status check failed: ${res.status} ${res.statusText}`);
  }
  return Z_Status.parse(await res.json());
}

function formatBytes(bytes: number): string {
  const mb = bytes / 1024 / 1024;
  return mb >= 1000 ? `${(mb / 1024).toFixed(2)} GB` : `${mb.toFixed(2)} MB`;
}

function formatSpeed(bytes: number, seconds: number): string {
  const mbps = bytes / 1024 / 1024 / seconds;
  return mbps >= 1000
    ? `${(mbps / 1024).toFixed(1)} GB/s`
    : `${mbps.toFixed(1)} MB/s`;
}

async function downloadDump(options: {
  wait: boolean;
  fresh: boolean;
}): Promise<{ dumpPath: string; tempDir: string }> {
  const tempDir = join(tmpdir(), `pg-dump-cache-${Date.now()}`);
  await mkdir(tempDir, { recursive: true });

  const url = new URL(`${CACHE_SERVER_URL}dump`);
  if (options.fresh) url.searchParams.set("fresh", "true");
  else if (options.wait) url.searchParams.set("wait", "true");

  console.log(`Fetching dump from ${url}...`);
  const fetchStart = performance.now();
  const res = await fetch(url, { headers });

  if (!res.ok) {
    const body = (await res.json().catch(() => ({}))) as { error?: unknown };
    throw new Error(
      `Download failed: ${res.status} ${body.error ?? res.statusText}`,
    );
  }

  const dumpPath = join(tempDir, "dump.dump");
  const reader = res.body!.getReader();
  const decoder = new TextDecoder();
  let buffer = new Uint8Array(0);
  let foundDelimiter = false;
  const binaryChunks: Uint8Array[] = [];
  let downloadStart = fetchStart;

  // Read stream — status lines arrive before \0\n delimiter, then binary data
  for (;;) {
    const { done, value } = await reader.read();
    if (done) break;

    if (foundDelimiter) {
      binaryChunks.push(value);
      continue;
    }

    // Append to buffer
    const merged = new Uint8Array(buffer.length + value.length);
    merged.set(buffer);
    merged.set(value, buffer.length);
    buffer = merged;

    // Scan for \0\n delimiter
    let delimIdx = -1;
    for (let i = 0; i < buffer.length - 1; i++) {
      if (buffer[i] === 0x00 && buffer[i + 1] === 0x0a) {
        delimIdx = i;
        break;
      }
    }

    if (delimIdx >= 0) {
      // Print any remaining status lines before delimiter
      const statusText = decoder.decode(buffer.slice(0, delimIdx));
      for (const line of statusText.split("\n").filter(Boolean)) {
        try {
          const msg = JSON.parse(line) as { status: string };
          console.log(`  ${msg.status}`);
        } catch {}
      }

      // Everything after delimiter is binary — download timer starts here
      downloadStart = performance.now();
      const remaining = buffer.slice(delimIdx + 2);
      if (remaining.length > 0) binaryChunks.push(remaining);
      foundDelimiter = true;
      buffer = new Uint8Array(0);
    } else {
      // Print complete lines as they arrive for real-time feedback
      let lastNewline = -1;
      for (let i = buffer.length - 1; i >= 0; i--) {
        if (buffer[i] === 0x0a) {
          lastNewline = i;
          break;
        }
      }
      if (lastNewline >= 0) {
        const complete = decoder.decode(buffer.slice(0, lastNewline + 1));
        for (const line of complete.split("\n").filter(Boolean)) {
          try {
            const msg = JSON.parse(line) as { status: string };
            console.log(`  ${msg.status}`);
          } catch {}
        }
        buffer = buffer.slice(lastNewline + 1);
      }
    }
  }

  if (!foundDelimiter) {
    throw new Error("Server closed connection before sending dump data");
  }

  // Combine binary chunks and write
  const totalBytes = binaryChunks.reduce((s, c) => s + c.length, 0);
  const data = new Uint8Array(totalBytes);
  let off = 0;
  for (const chunk of binaryChunks) {
    data.set(chunk, off);
    off += chunk.length;
  }
  await Bun.write(dumpPath, data);

  const downloadSeconds = (performance.now() - downloadStart) / 1000;
  console.log(
    `Downloaded ${formatBytes(totalBytes)} in ${downloadSeconds.toFixed(1)}s (${formatSpeed(totalBytes, downloadSeconds)})`,
  );

  return { dumpPath, tempDir };
}

async function restoreDump(dumpPath: string): Promise<void> {
  if (!LOCAL_DB_URL) {
    console.log("LOCAL_DB_URL not set, skipping restore");
    return;
  }

  console.log(`Restoring to ${LOCAL_DB_URL.replace(/:[^:@]+@/, ":***@")}...`);

  const restoreStart = performance.now();
  const proc = Bun.spawn(
    [
      "pg_restore",
      "--no-owner",
      "--no-acl",
      "--clean",
      "--if-exists",
      `--dbname=${LOCAL_DB_URL}`,
      dumpPath,
    ],
    {
      stdout: "inherit",
      stderr: "inherit",
    },
  );

  const exitCode = await proc.exited;
  const restoreSeconds = ((performance.now() - restoreStart) / 1000).toFixed(1);

  if (exitCode !== 0) {
    throw new Error(`pg_restore failed with exit code ${exitCode}`);
  }

  console.log(`Restore completed in ${restoreSeconds}s`);
}

async function cleanup(tempDir: string): Promise<void> {
  await rm(tempDir, { recursive: true });
}

async function isInSourceRepo(): Promise<boolean> {
  try {
    const scriptDir = import.meta.dir;
    const gitDir = join(scriptDir, ".git");
    await stat(gitDir);
    // .git exists, check if it's the source repo
    const proc = Bun.spawn(["git", "remote", "get-url", "origin"], {
      cwd: scriptDir,
      stdout: "pipe",
      stderr: "pipe",
    });
    const output = await new Response(proc.stdout).text();
    return output.includes("futuritywork/pg-dump-cache");
  } catch {
    return false;
  }
}

async function checkForUpdates(force = false): Promise<void> {
  if (NO_UPDATE_CHECK || values["no-update-check"]) {
    return;
  }

  try {
    // Skip update if running from the source repository
    if (await isInSourceRepo()) {
      return;
    }

    // Check if we've checked recently (skip if forced)
    if (!force) {
      try {
        const stats = await stat(UPDATE_CHECK_FILE);
        const lastCheck = stats.mtimeMs;
        if (Date.now() - lastCheck < UPDATE_CHECK_INTERVAL) {
          return;
        }
      } catch {
        // File doesn't exist, proceed with check
      }
    }

    // Ensure config directory exists
    await mkdir(CONFIG_DIR, { recursive: true });

    // Fetch remote version
    const res = await fetch(GITHUB_RAW_URL);
    if (!res.ok) {
      return; // Silently fail on network errors
    }
    const remoteContent = await res.text();

    // Read local version
    const localPath = import.meta.path;
    const localContent = await readFile(localPath, "utf-8");

    // Update timestamp file
    await writeFile(UPDATE_CHECK_FILE, new Date().toISOString());

    // Compare and update if different
    if (remoteContent !== localContent) {
      console.log(`Updating client in place: ${localPath}`);
      await writeFile(localPath, remoteContent);
      console.log("Update complete. Re-running with new version...");

      // Re-execute with --no-update-check to prevent infinite loop
      const args = [...process.argv.slice(1), "--no-update-check"];
      const proc = Bun.spawn(["bun", ...args], {
        stdout: "inherit",
        stderr: "inherit",
        stdin: "inherit",
      });
      const exitCode = await proc.exited;
      process.exit(exitCode);
    }
  } catch (error) {
    // Silently continue on any errors
    if (process.env.DEBUG) {
      console.warn("Update check failed:", error);
    }
  }
}

async function main(): Promise<void> {
  await checkForUpdates(values["force-update"]);

  if (values.status) {
    const status = await getStatus();
    console.log("Server status:");
    console.log(`  Has cache: ${status.hasCache}`);
    console.log(`  Cache timestamp: ${status.cacheTimestamp ?? "N/A"}`);
    console.log(
      `  Cache age: ${status.cacheAgeSeconds !== null ? `${Math.round(status.cacheAgeSeconds / 60)} minutes` : "N/A"}`,
    );
    console.log(`  Refreshing: ${status.refreshing}`);
    console.log(`  TTL: ${status.ttl}s`);
    return;
  }

  const { dumpPath, tempDir } = await downloadDump({
    wait: values.wait ?? false,
    fresh: values.fresh ?? false,
  });
  try {
    await restoreDump(dumpPath);
  } finally {
    await cleanup(tempDir);
  }

  const elapsedSec = (performance.now() - startTime) / 1000;
  console.log(`Done in ${elapsedSec.toFixed(1)}s`);
}

main().catch((error) => {
  console.error("Error:", error.message);
  process.exit(1);
});