@dbx-tools/cli 0.9.58 → 0.9.60

package/src/auth/cli.ts

@@ -0,0 +1,165 @@
+/**
+ * `dbx auth` Commander program for Databricks OAuth.
+ *
+ * The command delegates profile resolution and token lifecycle to
+ * `@dbx-tools/auth`.
+ *
+ * @module
+ */
+
+import * as databricks from "@dbx-tools/auth";
+import { addArgs, parseArgs } from "@dbx-tools/cli-args";
+import type { DatabricksAuthClientInfo } from "@dbx-tools/shared-auth/client";
+import { Command, CommanderError } from "commander";
+
+import {
+  AuthCliOptionsSchema,
+  databricksAuthOptions,
+  TokenCommandOptionsSchema,
+  type AuthCliOptions,
+} from "./options.ts";
+
+interface AuthContext {
+  auth: databricks.AuthClient;
+  close(): Promise<void>;
+}
+
+interface AuthCliDependencies {
+  createAuthClient: typeof databricks.client.createAuthClient;
+  writeJson(value: unknown): void;
+  writeText(value: string): void;
+}
+
+const DEFAULT_DEPENDENCIES: AuthCliDependencies = {
+  createAuthClient: databricks.client.createAuthClient,
+  writeJson: (value) => {
+    process.stdout.write(`${JSON.stringify(value, null, 2)}\n`);
+  },
+  writeText: (value) => {
+    process.stdout.write(`${value}\n`);
+  },
+};
+
+/** Open the selected authentication client. */
+async function openAuth(
+  options: AuthCliOptions,
+  dependencies: AuthCliDependencies,
+): Promise<AuthContext> {
+  return {
+    auth: await dependencies.createAuthClient(databricksAuthOptions(options)),
+    close: async () => {},
+  };
+}
+
+/** Execute an auth action and close resources owned by the command. */
+async function withAuth(
+  options: AuthCliOptions,
+  dependencies: AuthCliDependencies,
+  action: (context: AuthContext) => Promise<void>,
+): Promise<void> {
+  const context = await openAuth(options, dependencies);
+  try {
+    await action(context);
+  } finally {
+    await context.close();
+  }
+}
+
+/** Shape a generated token record for stable CLI JSON output. */
+function tokenJson(token: databricks.AccessToken): Record<string, unknown> {
+  return {
+    access_token: token.accessToken,
+    token_type: token.tokenType,
+    ...(token.expiry ? { expiry: token.expiry } : {}),
+    ...(token.scopes.length ? { scopes: token.scopes } : {}),
+  };
+}
+
+/** Select the serializable identity fields from an authentication client. */
+function clientInfo(auth: databricks.AuthClient): DatabricksAuthClientInfo {
+  return {
+    ...(auth.profile ? { profile: auth.profile } : {}),
+    host: auth.host,
+    ...(auth.accountId ? { accountId: auth.accountId } : {}),
+    ...(auth.workspaceId ? { workspaceId: auth.workspaceId } : {}),
+    target: auth.target,
+    authType: auth.authType,
+    principal: auth.principal,
+  };
+}
+
+/** Build the `dbx auth` Commander program without parsing arguments. */
+export function buildProgram(
+  name = "dbx auth",
+  overrides: Partial<AuthCliDependencies> = {},
+): Command {
+  const dependencies = { ...DEFAULT_DEPENDENCIES, ...overrides };
+  const program = addArgs(
+    new Command()
+      .name(name)
+      .description("Authenticate to Databricks with user or machine OAuth")
+      .showHelpAfterError(),
+    AuthCliOptionsSchema,
+  );
+  const options = (): AuthCliOptions => parseArgs(program, AuthCliOptionsSchema);
+
+  program
+    .command("login")
+    .description("Force browser login and return an access token")
+    .action(async () => {
+      await withAuth(options(), dependencies, async ({ auth }) => {
+        dependencies.writeJson(tokenJson(await auth.token({ login: true })));
+      });
+    });
+
+  const tokenCommand = addArgs(
+    program.command("token").description("Return a valid access token, logging in when needed"),
+    TokenCommandOptionsSchema,
+  );
+  tokenCommand.action(async () => {
+    const tokenOptions = parseArgs(tokenCommand, TokenCommandOptionsSchema);
+    await withAuth(options(), dependencies, async ({ auth }) => {
+      const token = await auth.token({
+        ...(tokenOptions.login ? {} : { login: false }),
+        ...(tokenOptions.forceRefresh ? { refresh: true } : {}),
+      });
+      dependencies.writeJson(tokenJson(token));
+    });
+  });
+
+  program
+    .command("profile")
+    .description("Print the configured or automatically detected profile")
+    .action(async () => {
+      await withAuth(options(), dependencies, async ({ auth }) => {
+        dependencies.writeText(auth.profile ?? "ambient");
+      });
+    });
+
+  program
+    .command("logout")
+    .description("Delete the stored credential for the selected profile")
+    .action(async () => {
+      await withAuth(options(), dependencies, async ({ auth }) => {
+        await auth.logout();
+      });
+    });
+
+  program
+    .command("status")
+    .description("Show the resolved authentication client configuration")
+    .action(async () => {
+      await withAuth(options(), dependencies, async ({ auth }) => {
+        dependencies.writeJson(clientInfo(auth));
+      });
+    });
+
+  return program;
+}
+
+/** Parse `argv` and run the selected auth operation. */
+export async function runCli(argv: string[]): Promise<void> {
+  await buildProgram().parseAsync(argv);
+}
+
+export { CommanderError };

package/src/auth/options.ts

@@ -0,0 +1,79 @@
+/**
+ * Databricks auth command-line option schemas and runtime projection.
+ *
+ * @module
+ */
+
+import {
+  AuthOptionsSchema,
+  DatabricksAuthOptionsSchema,
+  type DatabricksAuthOptions,
+} from "@dbx-tools/auth/config";
+import { stringUtils } from "@dbx-tools/shared-core";
+import { z } from "zod";
+
+export const AuthCliOptionsSchema = z
+  .object({
+    profile: DatabricksAuthOptionsSchema.shape.profile,
+    host: DatabricksAuthOptionsSchema.shape.host,
+    accountId: DatabricksAuthOptionsSchema.shape.accountId,
+    workspaceId: DatabricksAuthOptionsSchema.shape.workspaceId,
+    configFile: DatabricksAuthOptionsSchema.shape.configFile,
+    clientId: DatabricksAuthOptionsSchema.shape.clientId,
+    groupId: DatabricksAuthOptionsSchema.shape.groupId,
+    authType: DatabricksAuthOptionsSchema.shape.authType,
+    scopes: z
+      .preprocess(
+        (value) =>
+          Array.isArray(value)
+            ? value.flatMap((entry) => stringUtils.parseList(String(entry)))
+            : value,
+        z.array(z.string().trim().min(1)),
+      )
+      .default([])
+      .describe("OAuth scopes."),
+    target: DatabricksAuthOptionsSchema.shape.target,
+    lockTimeoutMs: AuthOptionsSchema.shape.lockTimeoutMs,
+    loginTimeoutMs: AuthOptionsSchema.shape.loginTimeoutMs,
+    refreshBufferMs: AuthOptionsSchema.shape.refreshBufferMs,
+    preferUserToMachine: z
+      .boolean()
+      .default(true)
+      .describe("Prefer a matching user profile over selected machine credentials."),
+  })
+  .strict()
+  .describe("Options shared by every Databricks auth command.");
+
+export const TokenCommandOptionsSchema = z
+  .object({
+    forceRefresh: z.boolean().default(false).describe("Refresh the token before returning it."),
+    login: z.boolean().default(true).describe("Log in when credentials are missing or invalid."),
+  })
+  .strict()
+  .describe("Options for returning a Databricks access token.");
+
+export type AuthCliOptions = z.output<typeof AuthCliOptionsSchema>;
+
+export type TokenCommandOptions = z.output<typeof TokenCommandOptionsSchema>;
+
+/** Convert flat CLI fields into the Node auth package's nested options. */
+export function databricksAuthOptions(options: AuthCliOptions): DatabricksAuthOptions {
+  return {
+    profile: options.profile,
+    host: options.host,
+    accountId: options.accountId,
+    workspaceId: options.workspaceId,
+    configFile: options.configFile,
+    clientId: options.clientId,
+    groupId: options.groupId,
+    authType: options.authType,
+    scopes: options.scopes.length > 0 ? options.scopes : undefined,
+    target: options.target,
+    auth: {
+      lockTimeoutMs: options.lockTimeoutMs,
+      loginTimeoutMs: options.loginTimeoutMs,
+      refreshBufferMs: options.refreshBufferMs,
+    },
+    preferUserToMachine: options.preferUserToMachine,
+  };
+}

package/src/cli.ts

@@ -106,47 +106,63 @@ export function buildProgram(name: string = PROGRAM_NAMES[0]): Command {
     program,
     "appkit",
     "AppKit helpers (env: print the environment an AppKit app resolves)",
-    async () => (await import("@dbx-tools/cli-appkit-env/cli")).buildProgram,
+    async () => (await import("./appkit/cli.ts")).buildProgram,
   );
 
   addForwardedCommand(
     program,
     "auth",
     "Authenticate to Databricks and manage OAuth tokens",
-    async () => (await import("@dbx-tools/cli-auth/cli")).buildProgram,
+    async () => (await import("./auth/cli.ts")).buildProgram,
   );
 
   addForwardedCommand(
     program,
     "tunnel",
     "Run a public portr tunnel with an email-OTP gate",
-    async () => (await import("@dbx-tools/cli-tunnel/cli")).buildProgram,
+    async () => (await import("./tunnel/cli.ts")).buildProgram,
   );
 
   addForwardedCommand(
     program,
     "lakebase-proxy",
     "Run the Node Databricks Lakebase PostgreSQL proxy",
-    async () => (await import("@dbx-tools/cli-lakebase-proxy/cli")).buildProgram,
+    async () => (await import("./lakebase-proxy/cli.ts")).buildProgram,
   );
 
   addForwardedCommand(
     program,
     "model-gateway",
     "Run the foreground AppKit Databricks model gateway",
-    async () => (await import("@dbx-tools/cli-model-gateway/cli")).buildProgram,
+    async () => (await import("./model-gateway/cli.ts")).buildProgram,
   );
 
   addForwardedCommand(
     program,
     "graphiti",
     "Run Graphiti or manage its current-user desktop service",
-    async () => (await import("@dbx-tools/cli-graphiti/cli")).buildProgram,
+    async () => (await import("./graphiti/cli.ts")).buildProgram,
   );
 
   return program;
 }
 
+/** Build the complete command tree used for generated documentation. */
+export async function buildDocumentationProgram(name: string = PROGRAM_NAMES[0]): Promise<Command> {
+  const program = buildProgram(name);
+  (program.commands as Command[]).splice(
+    1,
+    program.commands.length - 1,
+    (await import("./appkit/cli.ts")).buildProgram("appkit"),
+    (await import("./auth/cli.ts")).buildProgram("auth"),
+    (await import("./tunnel/cli.ts")).buildProgram("tunnel"),
+    (await import("./lakebase-proxy/cli.ts")).buildProgram("lakebase-proxy"),
+    (await import("./model-gateway/cli.ts")).buildProgram("model-gateway"),
+    (await import("./graphiti/cli.ts")).buildProgram("graphiti"),
+  );
+  return program;
+}
+
 /**
  * Program name for help output: whichever bin the user actually typed, so
  * `dbx-tools --help` doesn't document itself as `dbx`. Falls back to `dbx`.

package/src/graphiti/cli.ts

@@ -0,0 +1,75 @@
+/**
+ * Commander-owned Graphiti CLI and shared current-user service definition.
+ *
+ * This module owns every user-facing Graphiti command. The shared Graphiti Zod
+ * schema owns its options, environment names, defaults, and validation. The
+ * Node Graphiti package owns runtime and Python interaction.
+ *
+ * @module
+ */
+import { addArgs, parseArgs, serializeArgs } from "@dbx-tools/cli-args";
+import { buildServiceCommand, type CliServiceCliDependencies } from "@dbx-tools/cli-service/cli";
+import { defineService, type CliServiceDefinition } from "@dbx-tools/cli-service/definition";
+import { runGraphiti, type GraphitiRuntimeOptions } from "@dbx-tools/graphiti/runtime";
+import { Command } from "commander";
+
+import { PACKAGE_VERSION } from "../../index.ts";
+import {
+  GraphitiCliOptionsSchema,
+  GraphitiOptionsSchema,
+  type GraphitiOptions,
+} from "./options.ts";
+
+/** Injectable runtime and service lifecycle boundaries for CLI callers. */
+export interface GraphitiCliDependencies {
+  /** Run the Graphiti stack through the Node runtime owner. */
+  run(options: GraphitiRuntimeOptions): Promise<void>;
+  /** Shared service command dependencies for tests and embedding. */
+  readonly service?: CliServiceCliDependencies;
+}
+
+/** Runtime options persisted when installing the Graphiti desktop service. */
+export type GraphitiServiceOptions = GraphitiOptions;
+
+/** Define a Graphiti service without duplicating installation or lifecycle policy. */
+export function graphitiServiceDefinition(
+  options: GraphitiServiceOptions = {},
+): CliServiceDefinition {
+  return defineService(import.meta.url, {
+    id: "dbx-tools.cli-graphiti",
+    name: "dbx graphiti",
+    pythonPackage: { name: "dbx-tools-graphiti" },
+    command: {
+      binName: "dbx-graphiti",
+      arguments: serializeArgs(GraphitiOptionsSchema.parse(options)),
+    },
+  });
+}
+
+/** Build foreground execution and shared desktop-service commands. */
+export function buildProgram(
+  name = "dbx graphiti",
+  dependencies: GraphitiCliDependencies = { run: runGraphiti },
+): Command {
+  const program = addArgs(
+    new Command()
+      .name(name)
+      .description("Run Graphiti or manage its current-user desktop service")
+      .version(PACKAGE_VERSION, "-v, --version"),
+    GraphitiCliOptionsSchema,
+  ).action(async () => {
+    await dependencies.run(graphitiOptions(program));
+  });
+
+  program.addCommand(
+    buildServiceCommand(
+      () => graphitiServiceDefinition(graphitiOptions(program)),
+      dependencies.service,
+    ),
+  );
+  return program;
+}
+
+function graphitiOptions(command: Command): GraphitiOptions {
+  return GraphitiOptionsSchema.parse(parseArgs(command, GraphitiCliOptionsSchema));
+}

package/src/graphiti/options.ts

@@ -0,0 +1,20 @@
+/**
+ * Graphiti CLI option surface.
+ *
+ * The Node Graphiti runtime owns the complete shared option contract. This
+ * module re-exports that surface without creating parallel CLI policy.
+ *
+ * @module
+ */
+
+export {
+  GRAPHITI_DEFAULTS,
+  GraphitiCliOptionsSchema,
+  GraphitiOptionsSchema,
+  ResolvedGraphitiOptionsSchema,
+  graphitiOptionOverrides,
+  graphitiOptionsEnvironment,
+  graphitiOptionsFromEnvironment,
+  resolveGraphitiOptions,
+} from "@dbx-tools/graphiti/options";
+export type { GraphitiOptions, ResolvedGraphitiOptions } from "@dbx-tools/graphiti/options";

package/src/lakebase-proxy/cancellation.ts

@@ -0,0 +1,116 @@
+/** PostgreSQL cancellation key translation and TLS forwarding. */
+
+import { randomBytes } from "node:crypto";
+import { createConnection, isIP, type Socket } from "node:net";
+import { connect as connectTls, type TLSSocket } from "node:tls";
+
+import { cancelRequest, sslRequest, type CancelMessage } from "./protocol.ts";
+
+/** Upstream PostgreSQL backend key used to forward a cancellation request. */
+export interface CancellationTarget {
+  host: string;
+  port: number;
+  processId: number;
+  secretKey: number;
+}
+
+/** Sends one PostgreSQL cancellation request to its upstream backend. */
+export type CancellationForwarder = (target: CancellationTarget) => Promise<void>;
+
+let nextProcessId = 10_000;
+
+/** Maps synthetic local backend keys onto real upstream cancellation targets. */
+export class CancellationRegistry {
+  private readonly targets = new Map<string, CancellationTarget>();
+
+  constructor(private readonly forwarder: CancellationForwarder = forwardCancellation) {}
+
+  register(target: CancellationTarget): CancelMessage {
+    const processId = nextProcessId++;
+    const secretKey = randomBytes(4).readInt32BE(0);
+    this.targets.set(key(processId, secretKey), target);
+    return { kind: "cancel", processId, secretKey };
+  }
+
+  remove(message: CancelMessage): void {
+    this.targets.delete(key(message.processId, message.secretKey));
+  }
+
+  async forward(message: CancelMessage): Promise<boolean> {
+    const target = this.targets.get(key(message.processId, message.secretKey));
+    if (!target) return false;
+    await this.forwarder(target);
+    return true;
+  }
+}
+
+/** Forward a PostgreSQL cancellation request over a verified TLS connection. */
+export async function forwardCancellation(target: CancellationTarget): Promise<void> {
+  const socket = createConnection({ host: target.host, port: target.port });
+  await event(socket, "connect");
+  socket.write(sslRequest());
+  const response = await readByte(socket);
+  if (response !== "S".charCodeAt(0)) {
+    socket.destroy();
+    throw new Error("Lakebase upstream refused TLS for cancellation");
+  }
+  const tls = connectTls({
+    socket,
+    rejectUnauthorized: true,
+    ...(isIP(target.host) === 0 ? { servername: target.host } : {}),
+  });
+  await event(tls, "secureConnect");
+  await writeAndClose(tls, cancelRequest(target.processId, target.secretKey));
+}
+
+function key(processId: number, secretKey: number): string {
+  return `${processId}:${secretKey}`;
+}
+
+function event(socket: Socket | TLSSocket, name: "connect" | "secureConnect"): Promise<void> {
+  return new Promise((resolve, reject) => {
+    const onEvent = () => finish();
+    const onError = (error: Error) => finish(error);
+    const finish = (error?: Error) => {
+      socket.off(name, onEvent);
+      socket.off("error", onError);
+      if (error) reject(error);
+      else resolve();
+    };
+    socket.once(name, onEvent);
+    socket.once("error", onError);
+  });
+}
+
+function readByte(socket: Socket): Promise<number> {
+  return new Promise((resolve, reject) => {
+    const onData = (chunk: Buffer) => {
+      finish();
+      if (chunk.length > 1) socket.unshift(chunk.subarray(1));
+      resolve(chunk[0]!);
+    };
+    const onError = (error: Error) => {
+      finish();
+      reject(error);
+    };
+    const onEnd = () => {
+      finish();
+      reject(new Error("Lakebase upstream closed during TLS negotiation"));
+    };
+    const finish = () => {
+      socket.off("data", onData);
+      socket.off("error", onError);
+      socket.off("end", onEnd);
+    };
+    socket.once("data", onData);
+    socket.once("error", onError);
+    socket.once("end", onEnd);
+  });
+}
+
+function writeAndClose(socket: TLSSocket, packet: Buffer): Promise<void> {
+  return new Promise((resolve, reject) => {
+    socket.end(packet, () => resolve());
+    socket.once("error", reject);
+  });
+}

package/src/lakebase-proxy/cli.ts

@@ -0,0 +1,80 @@
+/** Commander entry point for the pure Node Lakebase proxy. */
+
+import { addArgs, parseArgs, serializeArgs } from "@dbx-tools/cli-args";
+import { buildServiceCommand, type CliServiceCliDependencies } from "@dbx-tools/cli-service/cli";
+import { defineService, type CliServiceDefinition } from "@dbx-tools/cli-service/definition";
+import { connectionUrl } from "@dbx-tools/lakebase";
+import { Command } from "commander";
+
+import { PACKAGE_VERSION } from "../../index.ts";
+import {
+  LakebaseProxyOptionsSchema,
+  LakebaseProxyUrlOptionsSchema,
+  resolveLakebaseProxyOptions,
+  type LakebaseProxyOptions,
+} from "./options.ts";
+import { LakebaseProxy } from "./proxy.ts";
+
+/** Injectable service boundary for Lakebase proxy CLI composition and tests. */
+export interface LakebaseProxyCliDependencies {
+  readonly service?: CliServiceCliDependencies;
+}
+
+/** Build the tray-only Lakebase proxy service definition. */
+export function lakebaseProxyServiceDefinition(
+  options: LakebaseProxyOptions = {},
+): CliServiceDefinition {
+  const resolved = resolveLakebaseProxyOptions(options);
+  return defineService(import.meta.url, {
+    id: "dbx-tools.cli-lakebase-proxy",
+    name: "dbx lakebase proxy",
+    command: {
+      binName: "dbx-lakebase-proxy",
+      arguments: serializeArgs(resolved),
+    },
+  });
+}
+
+/** Build the Lakebase proxy command-line program. */
+export function buildProgram(
+  name = "dbx lakebase-proxy",
+  dependencies: LakebaseProxyCliDependencies = {},
+): Command {
+  const program = addArgs(
+    new Command(name)
+      .description("Run a loopback PostgreSQL proxy for Databricks Lakebase")
+      .version(PACKAGE_VERSION)
+      .enablePositionalOptions(),
+    LakebaseProxyOptionsSchema,
+  ).action(async () => {
+    const proxy = new LakebaseProxy(parseArgs(program, LakebaseProxyOptionsSchema));
+    await proxy.listen();
+    await waitForShutdown();
+    await proxy.close();
+  });
+
+  const urlCommand = addArgs(
+    program.command("url").description("Format a local PostgreSQL URL for a Lakebase target"),
+    LakebaseProxyUrlOptionsSchema,
+  ).action(() => {
+    const options = parseArgs(urlCommand, LakebaseProxyUrlOptionsSchema);
+    const target = options.target ?? options.lakebaseEndpoint!;
+    process.stdout.write(`${connectionUrl(target, options.listen.host, options.listen.port)}\n`);
+  });
+  let installCommand!: Command;
+  const serviceCommand = buildServiceCommand(() => {
+    return lakebaseProxyServiceDefinition(parseArgs(installCommand, LakebaseProxyOptionsSchema));
+  }, dependencies.service);
+  installCommand = serviceCommand.commands.find((command) => command.name() === "install")!;
+  addArgs(installCommand, LakebaseProxyOptionsSchema);
+  program.addCommand(serviceCommand);
+  return program;
+}
+
+function waitForShutdown(): Promise<void> {
+  return new Promise((resolve) => {
+    const shutdown = () => resolve();
+    process.once("SIGINT", shutdown);
+    process.once("SIGTERM", shutdown);
+  });
+}

package/src/lakebase-proxy/options.ts

@@ -0,0 +1,56 @@
+/**
+ * Lakebase proxy listener, service, and URL command options.
+ *
+ * @module
+ */
+
+import { options } from "@dbx-tools/shared-core";
+import { z } from "zod";
+
+const text = (description: string) => z.string().trim().min(1).optional().describe(description);
+
+export const LakebaseProxyOptionsSchema = z
+  .object({
+    ...options.PostgresOptionsSchema.shape,
+    listen: options
+      .listenAddressSchema({ port: 5432, loopback: true })
+      .describe("Loopback listener address."),
+    startupTimeoutSeconds: z.coerce
+      .number<number>()
+      .int()
+      .nonnegative()
+      .default(30)
+      .describe("Startup timeout in seconds."),
+    profile: options.DatabricksOptionsSchema.shape.profile.describe("Exact Databricks profile."),
+  })
+  .strict()
+  .describe("Lakebase proxy listener and service options.");
+
+export const LakebaseProxyUrlOptionsSchema = z
+  .object({
+    target: text("Lakebase project, resource path, host, or URL."),
+    lakebaseEndpoint: options.LakebaseOptionsSchema.shape.lakebaseEndpoint,
+    listen: options.listenAddressSchema({ port: 5432 }).describe("Local proxy address."),
+  })
+  .strict()
+  .refine(
+    (configured) => configured.target !== undefined || configured.lakebaseEndpoint !== undefined,
+    {
+      message: "URL requires target or endpoint.",
+      path: ["target"],
+    },
+  )
+  .describe("Options for formatting a local Lakebase PostgreSQL URL.");
+
+export type LakebaseProxyOptions = z.input<typeof LakebaseProxyOptionsSchema>;
+
+export type ResolvedLakebaseProxyOptions = z.output<typeof LakebaseProxyOptionsSchema>;
+
+export type LakebaseProxyUrlOptions = z.output<typeof LakebaseProxyUrlOptionsSchema>;
+
+/** Validate and default Lakebase proxy listener options. */
+export function resolveLakebaseProxyOptions(
+  options: LakebaseProxyOptions = {},
+): ResolvedLakebaseProxyOptions {
+  return LakebaseProxyOptionsSchema.parse(options);
+}