buildai-cli 0.3.53__tar.gz → 0.3.55__tar.gz

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/CLAUDE.md

@@ -7,13 +7,14 @@ Typer-based CLI with two modes: standalone (PyPI, API-backed) and workspace (rep
 | Mode | Install | Auth | DB Access |
 |------|---------|------|-----------|
 | Standalone (`buildai`) | `uv tool install buildai-cli` | API key / JWT | Through API |
-| Workspace (`uv run buildai`) | Repo-local editable install | IAM / password | Direct via `db` |
+| Workspace (`uv run buildai`) | Repo-local editable install | IAM / password | Direct via `db`, through the profile-keyed local broker |
 
 ## Key Commands
 
 ```bash
 buildai auth whoami                                  # API auth inspection
 buildai db query "SELECT count(*) FROM core.clips"   # DB-direct
+buildai db broker status                             # Active local DB broker state
 buildai db schema tables                             # Schema introspection
 buildai db schema describe core.clips                # Table details
 buildai db --write migrate all                       # Run migrations
@@ -23,6 +24,10 @@ buildai db status                                    # Migration status
 ## Guards
 
 - `db` subcommands require workspace install + gcloud IAM.
+- `db` inspection/migration commands use a profile-keyed `alloydb-auth-proxy`
+  broker owned by the local machine, with state in `~/.buildai/db-brokers` so
+  all worktrees reuse the same listener; use `buildai db broker status|ensure|stop`
+  for that local transport.
 - Writes require `--write` flag.
 - Production migrations prompt for confirmation.
 - Worktree app/runtime targeting comes from explicit env vars and the repo Makefile, not a saved CLI context layer.

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: buildai-cli
-Version: 0.3.53
+Version: 0.3.55
 Summary: Build AI CLI (Typer)
 Requires-Python: >=3.11
 Requires-Dist: httpx>=0.27.0

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/cli/auth_local.py

@@ -98,7 +98,7 @@ class ResolvedLocalAuthProfile:
             exports.extend(
                 [
                     f'export ALLOYDB_RUNTIME_IMPERSONATE_SA_{suffix}="{self.target_service_account}"',
-                    f'export ALLOYDB_USE_AUTH_PROXY_{suffix}="false"',
+                    f'export ALLOYDB_USE_AUTH_PROXY_{suffix}="true"',
                 ]
             )
         elif self.name == "db-admin-local" and self.target_db_user and self.target_service_account:

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/cli/commands/db/__init__.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import typer
 
+from . import broker as broker_mod
 from . import migrate as migrate_mod
 from . import query as query_mod
 from . import schema as schema_mod
@@ -18,4 +19,5 @@ app = typer.Typer(
 app.command("query")(query_mod.query)
 app.command("status")(status_mod.status)
 app.command("migrate")(migrate_mod.migrate)
+app.add_typer(broker_mod.app, name="broker")
 app.add_typer(schema_mod.app, name="schema")

buildai_cli-0.3.55/cli/commands/db/broker.py

@@ -0,0 +1,60 @@
+"""Local AlloyDB broker commands for the DB CLI surface."""
+
+from __future__ import annotations
+
+import typer
+
+from cli.context import resolve_broker_config
+from cli.db_broker import BrokerConfig, broker_status, ensure_broker, stop_broker
+from cli.output import Format, format_option, render
+
+app = typer.Typer(
+    name="broker",
+    help="Inspect and manage the profile-keyed local AlloyDB broker.",
+    no_args_is_help=True,
+)
+
+
+def _config_from_context(ctx: typer.Context) -> BrokerConfig:
+    """Resolve the active DB command context into a broker config."""
+
+    settings = ctx.obj["settings"]
+    profile = ctx.obj.get("cli_profile")
+    return resolve_broker_config(settings, profile=profile)
+
+
+@app.command("status")
+def status(
+    ctx: typer.Context,
+    format: Format = format_option(),
+) -> None:
+    """Show whether the active profile broker is alive and listening."""
+
+    config = _config_from_context(ctx)
+    render(broker_status(config), format=format)
+
+
+@app.command("ensure")
+def ensure(
+    ctx: typer.Context,
+    format: Format = format_option(),
+) -> None:
+    """Start or reuse the active profile broker, then print its status."""
+
+    config = _config_from_context(ctx)
+    ensure_broker(config)
+    render(broker_status(config), format=format)
+
+
+@app.command("stop")
+def stop(
+    ctx: typer.Context,
+    format: Format = format_option(),
+) -> None:
+    """Stop the active profile broker if this workspace owns it."""
+
+    config = _config_from_context(ctx)
+    stopped = stop_broker(config)
+    payload = broker_status(config)
+    payload["stopped"] = stopped
+    render(payload, format=format)

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/cli/commands/db/common.py

@@ -9,6 +9,12 @@ from typing import Any
 import asyncpg
 from infra.settings import Settings
 
+from cli.db_broker import (
+    BrokeredDatabase,
+    broker_config_for_identity,
+    open_brokered_database,
+)
+
 MIGRATIONS_SA_EMAIL = "buildai-migrations-sa@data-470400.iam.gserviceaccount.com"
 MIGRATIONS_SA_DB_USER = "buildai-migrations-sa@data-470400.iam"
 REPO_ROOT = Path(__file__).resolve().parents[5]
@@ -62,35 +68,22 @@ def migration_label(item: Any) -> str:
     return str(item)
 
 
-async def get_migrations_connection(settings: Settings):
+async def get_migrations_connection(settings: Settings) -> BrokeredDatabase:
     """Open the canonical migrations identity for reviewed DDL work."""
 
-    from infra.database import Database
-
     settings.validate_environment_access()
 
     if not settings.alloydb_instance_uri:
         raise RuntimeError("AlloyDB instance URI is not configured for this environment.")
 
-    ip_type = "PRIVATE" if settings.use_private_ip else "PUBLIC"
-    db = Database(
-        instance_uri=settings.alloydb_instance_uri,
-        database=settings.db_name,
-        user=MIGRATIONS_SA_DB_USER,
+    config = broker_config_for_identity(
+        settings,
+        profile="migrations-local",
+        db_user=MIGRATIONS_SA_DB_USER,
         use_iam_auth=True,
-        ip_type=ip_type,
-        auth_proxy_host=(
-            settings.alloydb_auth_proxy_host if settings.effective_use_alloydb_auth_proxy else None
-        ),
-        auth_proxy_port=(
-            settings.effective_alloydb_auth_proxy_port
-            if settings.effective_use_alloydb_auth_proxy
-            else None
-        ),
-        impersonate_sa=MIGRATIONS_SA_EMAIL,
+        target_service_account=MIGRATIONS_SA_EMAIL,
     )
-    await db.connect()
-    return db
+    return await open_brokered_database(config)
 
 
 async def get_table_row_count(

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/cli/commands/db/query.py

@@ -17,6 +17,7 @@ QUERY_HINTS = {
     "public._migrations": "Tip: Use 'buildai db migrate' to inspect migration state.",
     "core.clips": None,
 }
+_QUERY_WRITE_PROFILES = frozenset({"internal_admin", "db-admin-local"})
 
 
 def _show_query_hints(sql: str) -> None:
@@ -117,7 +118,10 @@ def query(
     ] = False,
     write: Annotated[
         bool,
-        typer.Option("--write", help="Allow write queries (requires internal_admin profile)"),
+        typer.Option(
+            "--write",
+            help="Allow write queries (requires internal_admin or db-admin-local profile)",
+        ),
     ] = False,
 ) -> None:
     """Execute SQL against the selected database lane in inspection mode."""
@@ -130,11 +134,13 @@ def query(
         cli_profile = (ctx.obj or {}).get("cli_profile", "internal_viewer")
 
         if is_write and not write:
-            error("Write query blocked. Re-run with --write and an internal_admin profile.")
+            error(
+                "Write query blocked. Re-run with --write and an admin-capable DB profile."
+            )
             raise typer.Exit(1)
 
-        if is_write and cli_profile != "internal_admin":
-            error("Write query requires --profile internal_admin.")
+        if is_write and cli_profile not in _QUERY_WRITE_PROFILES:
+            error("Write query requires --profile internal_admin or --profile db-admin-local.")
             raise typer.Exit(1)
 
         if settings.is_production and is_ddl:

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/cli/context.py

@@ -28,7 +28,15 @@ import asyncpg
 from dal import scopes as dal_scopes
 from infra.settings import Settings, get_settings
 
-from infra import Database, get_logger
+from cli.auth_local import resolve_sanctioned_profile
+from cli.db_broker import (
+    BrokerConfig,
+    BrokeredDatabase,
+    broker_config_for_identity,
+    connect_via_broker,
+    open_brokered_database,
+)
+from infra import get_logger
 
 if TYPE_CHECKING:
     from dal.context import Context
@@ -72,6 +80,10 @@ _UNRESTRICTED_CLI_PROFILES = frozenset(
         "agent",
     }
 )
+DB_IDENTITY_PROFILE_ALIASES: dict[str, str] = {
+    "internal_admin": "db-admin-local",
+}
+_ADMIN_DB_IDENTITY_PROFILES = frozenset({"internal_admin", "db-admin-local"})
 
 
 def scopes_for_cli_profile(profile: str) -> frozenset[str]:
@@ -82,6 +94,12 @@ def scopes_for_cli_profile(profile: str) -> frozenset[str]:
     return resolved
 
 
+def db_identity_profile_for_cli_profile(profile: str) -> str:
+    """Return the DB identity profile backing one CLI-facing profile."""
+
+    return DB_IDENTITY_PROFILE_ALIASES.get(profile, profile)
+
+
 @dataclass(frozen=True)
 class AdminConnectionConfig:
     """Describe the canonical admin connection the ops-plane CLI should prefer.
@@ -166,7 +184,72 @@ def resolve_admin_connection_config(
     return None
 
 
-async def open_admin_database(settings: Settings) -> Database:
+def _fixed_profile_connection_config(
+    settings: Settings,
+    *,
+    profile: str,
+) -> AdminConnectionConfig | None:
+    """Resolve a sanctioned local auth profile into a DB login contract."""
+
+    try:
+        resolved = resolve_sanctioned_profile(
+            profile,
+            environment=settings.app_env.value,
+        )
+    except ValueError:
+        return None
+
+    if not resolved.db_access or not resolved.target_db_user or not resolved.target_service_account:
+        return None
+
+    return AdminConnectionConfig(
+        user=resolved.target_db_user,
+        use_iam_auth=True,
+        impersonate_sa=resolved.target_service_account,
+    )
+
+
+def resolve_broker_config(
+    settings: Settings,
+    *,
+    profile: str | None = None,
+) -> BrokerConfig:
+    """Resolve the CLI profile and settings into one local DB broker identity."""
+
+    from cli.config import resolve_cli_profile
+
+    resolved_profile = resolve_cli_profile(profile)
+    db_identity_profile = db_identity_profile_for_cli_profile(resolved_profile)
+    config = _fixed_profile_connection_config(settings, profile=db_identity_profile)
+
+    if config is None and resolved_profile in _ADMIN_DB_IDENTITY_PROFILES:
+        config = resolve_admin_connection_config(settings)
+
+    if config is None:
+        if resolved_profile in _ADMIN_DB_IDENTITY_PROFILES:
+            raise RuntimeError(
+                f"Profile '{resolved_profile}' requires admin DB credentials. "
+                "Set ALLOYDB_ADMIN_IAM_USER_* / ALLOYDB_ADMIN_IMPERSONATE_SA_* "
+                "or use the sanctioned db-admin-local profile."
+            )
+        config = AdminConnectionConfig(
+            user=settings.effective_db_user,
+            use_iam_auth=settings.effective_use_iam_auth,
+            impersonate_sa=settings.effective_alloydb_runtime_impersonate_sa or None,
+            password=settings.effective_alloydb_password,
+        )
+
+    return broker_config_for_identity(
+        settings,
+        profile=resolved_profile,
+        db_user=config.user,
+        use_iam_auth=config.use_iam_auth,
+        target_service_account=config.impersonate_sa,
+        password=config.password,
+    )
+
+
+async def open_admin_database(settings: Settings) -> BrokeredDatabase:
     """Open the canonical admin DB connection for privileged ops-plane access.
 
     Commands that need deterministic production introspection should use this
@@ -180,33 +263,26 @@ async def open_admin_database(settings: Settings) -> Database:
     if not settings.alloydb_instance_uri:
         raise RuntimeError("AlloyDB instance URI is not configured for this environment.")
 
-    config = resolve_admin_connection_config(settings)
+    config = resolve_admin_connection_config(settings) or _fixed_profile_connection_config(
+        settings,
+        profile="db-admin-local",
+    )
     if config is None:
         raise RuntimeError(
             "Admin credentials not configured for this environment. "
-            "Set ALLOYDB_ADMIN_IAM_USER_* / ALLOYDB_ADMIN_IMPERSONATE_SA_*."
+            "Set ALLOYDB_ADMIN_IAM_USER_* / ALLOYDB_ADMIN_IMPERSONATE_SA_* "
+            "or use the sanctioned db-admin-local profile."
         )
 
-    ip_type = "PRIVATE" if settings.use_private_ip else "PUBLIC"
-    db = Database(
-        instance_uri=settings.alloydb_instance_uri,
-        database=settings.db_name,
-        user=config.user,
-        password=config.password,
+    broker_config = broker_config_for_identity(
+        settings,
+        profile="db-admin-local",
+        db_user=config.user,
         use_iam_auth=config.use_iam_auth,
-        ip_type=ip_type,
-        auth_proxy_host=(
-            settings.alloydb_auth_proxy_host if settings.effective_use_alloydb_auth_proxy else None
-        ),
-        auth_proxy_port=(
-            settings.effective_alloydb_auth_proxy_port
-            if settings.effective_use_alloydb_auth_proxy
-            else None
-        ),
-        impersonate_sa=config.impersonate_sa,
+        target_service_account=config.impersonate_sa,
+        password=config.password,
     )
-    await db.connect()
-    return db
+    return await open_brokered_database(broker_config)
 
 
 @asynccontextmanager
@@ -240,14 +316,13 @@ async def get_inspection_connection(
             await _stamp_internal_inspection_session_contract(conn)
             yield conn
     else:
-        # Development/Production: use AlloyDB via infra.Database
-        db = Database.from_settings(settings)
+        config = resolve_broker_config(settings)
+        conn = await connect_via_broker(config)
         try:
-            await db.connect()
-            await _stamp_internal_inspection_session_contract(db.conn)
-            yield db.conn
+            await _stamp_internal_inspection_session_contract(conn)
+            yield conn
         finally:
-            await db.close()
+            await conn.close()
 
 
 @asynccontextmanager
@@ -283,7 +358,7 @@ async def _local_connection(
     - DB_PORT (default: 5432)
     - DB_USER (falls back to settings.db_user)
     - DB_PASSWORD (falls back to "postgres")
-    - DB_NAME (falls back to settings.db_name) - FIX: respect env var for CI
+    - DB_NAME (falls back to settings.db_name)
     """
     conn = await asyncpg.connect(
         host=os.getenv("DB_HOST", "localhost"),
@@ -304,7 +379,7 @@ async def get_cli_context(
     settings: Settings | None = None,
     *,
     profile: str | None = None,
-) -> AsyncGenerator[tuple[Database | None, "Context"], None]:
+) -> AsyncGenerator[tuple[BrokeredDatabase | None, "Context"], None]:
     """
     Get a database connection and canonical DAL context for CLI commands.
 
@@ -317,7 +392,7 @@ async def get_cli_context(
     Yields:
         Tuple of (Database instance or None, Context for DAL operations)
         - For test/CI: Database is None, Context wraps the raw connection
-        - For dev/prod: Database is the infra.Database instance
+        - For dev/prod: Database is a brokered Database-like wrapper
 
     Usage:
         async with get_cli_context(settings) as (db, ctx):
@@ -338,10 +413,9 @@ async def get_cli_context(
             ctx = Context.for_cli(conn, scopes=scopes_for_cli_profile(resolved_profile))
             yield None, ctx
     else:
-        # Development/Production: use infra.Database
-        db = Database.from_settings(settings)
+        broker_config = resolve_broker_config(settings, profile=resolved_profile)
+        db = await open_brokered_database(broker_config)
         try:
-            await db.connect()
             await _stamp_scoped_cli_session_contract(db.conn, profile=resolved_profile)
             ctx = Context.for_cli(db, scopes=scopes_for_cli_profile(resolved_profile))
             yield db, ctx

buildai_cli-0.3.55/cli/db_broker.py

@@ -0,0 +1,632 @@
+"""Profile-keyed local AlloyDB broker for CLI database operations."""
+
+from __future__ import annotations
+
+import fcntl
+import hashlib
+import json
+import os
+import re
+import shlex
+import shutil
+import socket
+import subprocess
+import time
+from contextlib import contextmanager
+from dataclasses import asdict, dataclass
+from pathlib import Path
+from typing import Any, Iterator
+
+import asyncpg
+from infra.settings import Settings
+
+REPO_ROOT = Path(__file__).resolve().parents[3]
+LEGACY_STATE_DIR = REPO_ROOT / ".buildai" / "db-brokers"
+STATE_DIR = Path(os.getenv("BUILDAI_DB_BROKER_STATE_DIR", Path.home() / ".buildai" / "db-brokers"))
+_PROFILE_PORT_OFFSETS = {
+    "engineers-dev": 0,
+    "internal_viewer": 1,
+    "internal_admin": 21,
+    "migrations-local": 10,
+    "db-admin-local": 20,
+}
+
+
+@dataclass(frozen=True)
+class BrokerConfig:
+    """Complete identity and endpoint contract for one local DB broker."""
+
+    environment: str
+    profile: str
+    instance_uri: str
+    database: str
+    db_user: str
+    use_iam_auth: bool
+    target_service_account: str | None
+    host: str
+    port: int
+    use_private_ip: bool
+    password: str = ""
+
+    @property
+    def identity(self) -> str:
+        """Return the stable identity hash for state files and diagnostics."""
+
+        payload = {
+            "environment": self.environment,
+            "profile": self.profile,
+            "instance_uri": self.instance_uri,
+            "database": self.database,
+            "db_user": self.db_user,
+            "use_iam_auth": self.use_iam_auth,
+            "target_service_account": self.target_service_account,
+            "host": self.host,
+            "port": self.port,
+            "use_private_ip": self.use_private_ip,
+        }
+        raw = json.dumps(payload, sort_keys=True, separators=(",", ":"))
+        return hashlib.sha256(raw.encode("utf-8")).hexdigest()[:16]
+
+    @property
+    def state_path(self) -> Path:
+        """Return the broker state file path for this identity."""
+
+        return STATE_DIR / f"{self.identity}.json"
+
+    @property
+    def log_path(self) -> Path:
+        """Return the broker log path for this identity."""
+
+        return STATE_DIR / f"{self.identity}.log"
+
+
+@dataclass(frozen=True)
+class BrokerState:
+    """Persist the child process and identity that own one broker listener."""
+
+    pid: int
+    config: dict[str, Any]
+    command: list[str]
+    log_path: str
+
+
+@dataclass(frozen=True)
+class ListenerProcess:
+    """Describe one local process already bound to a broker TCP endpoint."""
+
+    pid: int
+    command: list[str]
+
+
+@dataclass
+class BrokeredDatabase:
+    """Expose the small Database-like surface migration commands require."""
+
+    config: BrokerConfig
+    conn: asyncpg.Connection
+
+    async def fetch(self, query: str, *args: Any, timeout: float | None = None) -> list[Any]:
+        """Run ``fetch`` through the brokered connection."""
+
+        result = await self.conn.fetch(query, *args, timeout=timeout)
+        return list(result)
+
+    async def fetchrow(self, query: str, *args: Any, timeout: float | None = None) -> Any | None:
+        """Run ``fetchrow`` through the brokered connection."""
+
+        return await self.conn.fetchrow(query, *args, timeout=timeout)
+
+    async def fetchval(self, query: str, *args: Any, timeout: float | None = None) -> Any:
+        """Run ``fetchval`` through the brokered connection."""
+
+        return await self.conn.fetchval(query, *args, timeout=timeout)
+
+    async def execute(self, query: str, *args: Any, timeout: float | None = None) -> str:
+        """Run ``execute`` through the brokered connection."""
+
+        return str(await self.conn.execute(query, *args, timeout=timeout))
+
+    async def executemany(
+        self,
+        query: str,
+        args: list[tuple[Any, ...]],
+        timeout: float | None = None,
+    ) -> None:
+        """Run ``executemany`` through the brokered connection."""
+
+        await self.conn.executemany(query, args, timeout=timeout)
+
+    def transaction(self) -> Any:
+        """Return the asyncpg transaction context for this connection."""
+
+        return self.conn.transaction()
+
+    async def close(self) -> None:
+        """Close the underlying local asyncpg connection."""
+
+        await self.conn.close()
+
+
+async def _init_broker_connection(conn: asyncpg.Connection) -> None:
+    """Register codecs expected by DAL code using brokered CLI connections."""
+
+    await conn.set_type_codec(
+        "jsonb",
+        encoder=json.dumps,
+        decoder=json.loads,
+        schema="pg_catalog",
+    )
+    await conn.set_type_codec(
+        "json",
+        encoder=json.dumps,
+        decoder=json.loads,
+        schema="pg_catalog",
+    )
+
+
+def default_proxy_binary() -> str | None:
+    """Return the preferred AlloyDB Auth Proxy binary path if available."""
+
+    on_path = shutil.which("alloydb-auth-proxy")
+    if on_path:
+        return on_path
+    local_bin = Path.home() / ".local" / "bin" / "alloydb-auth-proxy"
+    if local_bin.is_file():
+        return str(local_bin)
+    return None
+
+
+def proxy_is_listening(*, host: str, port: int) -> bool:
+    """Return True when a TCP listener is present at the local broker endpoint."""
+
+    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
+        sock.settimeout(0.25)
+        return sock.connect_ex((host, port)) == 0
+
+
+def process_is_running(pid: int) -> bool:
+    """Return whether the broker PID still exists."""
+
+    try:
+        os.kill(pid, 0)
+    except ProcessLookupError:
+        return False
+    except PermissionError:
+        return True
+    return True
+
+
+@contextmanager
+def _broker_lock(config: BrokerConfig) -> Iterator[None]:
+    """Serialize local broker startup for one profile identity on this machine."""
+
+    STATE_DIR.mkdir(parents=True, exist_ok=True)
+    lock_path = STATE_DIR / f"{config.identity}.lock"
+    with lock_path.open("w", encoding="utf-8") as handle:
+        fcntl.flock(handle.fileno(), fcntl.LOCK_EX)
+        try:
+            yield
+        finally:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+
+
+def _state_paths(config: BrokerConfig) -> tuple[Path, ...]:
+    """Return current and legacy broker state paths for one identity."""
+
+    paths = [config.state_path]
+    legacy_path = LEGACY_STATE_DIR / f"{config.identity}.json"
+    if legacy_path != config.state_path:
+        paths.append(legacy_path)
+    return tuple(paths)
+
+
+def _read_broker_state_file(path: Path) -> BrokerState | None:
+    """Read one broker state file without trusting malformed stale content."""
+
+    try:
+        data = json.loads(path.read_text(encoding="utf-8"))
+        return BrokerState(
+            pid=int(data["pid"]),
+            config=dict(data["config"]),
+            command=list(data["command"]),
+            log_path=str(data["log_path"]),
+        )
+    except Exception:
+        return None
+
+
+def _unlink_state(config: BrokerConfig) -> None:
+    """Remove current and legacy state files for a stale broker identity."""
+
+    for path in _state_paths(config):
+        try:
+            path.unlink()
+        except FileNotFoundError:
+            pass
+
+
+def _state_payload(config: BrokerConfig, *, pid: int, command: list[str]) -> BrokerState:
+    """Build a serializable broker state object."""
+
+    return BrokerState(
+        pid=pid,
+        config=asdict(config),
+        command=command,
+        log_path=str(config.log_path),
+    )
+
+
+def read_broker_state(config: BrokerConfig) -> BrokerState | None:
+    """Read the state file for a broker identity when present and valid."""
+
+    for path in _state_paths(config):
+        state = _read_broker_state_file(path)
+        if state is not None:
+            return state
+    return None
+
+
+def _write_broker_state(config: BrokerConfig, state: BrokerState) -> None:
+    """Atomically write the broker state file."""
+
+    STATE_DIR.mkdir(parents=True, exist_ok=True)
+    tmp_path = config.state_path.with_suffix(".tmp")
+    tmp_path.write_text(
+        json.dumps(asdict(state), indent=2, sort_keys=True) + "\n",
+        encoding="utf-8",
+    )
+    os.replace(tmp_path, config.state_path)
+
+
+def _process_command(pid: int) -> list[str]:
+    """Return argv-like command text for a local PID when the OS exposes it."""
+
+    try:
+        completed = subprocess.run(
+            ["ps", "-p", str(pid), "-o", "command="],
+            capture_output=True,
+            text=True,
+            timeout=2,
+            check=False,
+        )
+    except Exception:
+        return []
+    if completed.returncode != 0:
+        return []
+    raw = completed.stdout.strip()
+    if not raw:
+        return []
+    try:
+        return shlex.split(raw)
+    except ValueError:
+        return raw.split()
+
+
+def _listening_pids(*, host: str, port: int) -> list[int]:
+    """Return local PIDs listening on a TCP port when OS tools expose them."""
+
+    if shutil.which("lsof") is not None:
+        try:
+            completed = subprocess.run(
+                ["lsof", "-nP", f"-iTCP:{port}", "-sTCP:LISTEN", "-t"],
+                capture_output=True,
+                text=True,
+                timeout=2,
+                check=False,
+            )
+        except Exception:
+            completed = None
+        if completed is not None and completed.returncode == 0:
+            pids: list[int] = []
+            for line in completed.stdout.splitlines():
+                try:
+                    pids.append(int(line.strip()))
+                except ValueError:
+                    continue
+            if pids:
+                return pids
+
+    if shutil.which("ss") is not None:
+        try:
+            completed = subprocess.run(
+                ["ss", "-H", "-ltnp", f"sport = :{port}"],
+                capture_output=True,
+                text=True,
+                timeout=2,
+                check=False,
+            )
+        except Exception:
+            completed = None
+        if completed is not None and completed.returncode == 0:
+            candidates: list[int] = []
+            for line in completed.stdout.splitlines():
+                if f":{port}" not in line:
+                    continue
+                candidates.extend(int(pid) for pid in re.findall(r"pid=(\d+)", line))
+            if candidates:
+                return candidates
+
+    del host
+    return []
+
+
+def _listener_processes(*, host: str, port: int) -> list[ListenerProcess]:
+    """Return process metadata for listeners on the broker endpoint."""
+
+    return [
+        ListenerProcess(pid=pid, command=_process_command(pid))
+        for pid in _listening_pids(host=host, port=port)
+    ]
+
+
+def _flag_value(command: list[str], flag: str) -> str | None:
+    """Return the value passed to a command-line flag in split or equals form."""
+
+    for index, value in enumerate(command):
+        if value == flag and index + 1 < len(command):
+            return command[index + 1]
+        prefix = f"{flag}="
+        if value.startswith(prefix):
+            return value[len(prefix) :]
+    return None
+
+
+def _is_matching_proxy_process(process: ListenerProcess, config: BrokerConfig) -> bool:
+    """Return True when a listener command is the broker this profile needs."""
+
+    command = process.command
+    if not command:
+        return False
+    if Path(command[0]).name != "alloydb-auth-proxy":
+        return False
+    if config.instance_uri not in command:
+        return False
+    if _flag_value(command, "--address") != config.host:
+        return False
+    if _flag_value(command, "--port") != str(config.port):
+        return False
+    if ("--auto-iam-authn" in command) != config.use_iam_auth:
+        return False
+    impersonate = _flag_value(command, "--impersonate-service-account")
+    if impersonate != config.target_service_account:
+        return False
+    if not config.use_private_ip and "--public-ip" not in command:
+        return False
+    if config.use_private_ip and "--public-ip" in command:
+        return False
+    return True
+
+
+def _adopt_matching_listener(config: BrokerConfig) -> BrokerState | None:
+    """Persist state for an already-running broker from another worktree."""
+
+    for process in _listener_processes(host=config.host, port=config.port):
+        if _is_matching_proxy_process(process, config):
+            state = _state_payload(config, pid=process.pid, command=process.command)
+            _write_broker_state(config, state)
+            return state
+    return None
+
+
+def _broker_command(config: BrokerConfig, *, proxy_binary: str) -> list[str]:
+    """Build the AlloyDB Auth Proxy command for one broker config."""
+
+    command = [
+        proxy_binary,
+        config.instance_uri,
+        "--address",
+        config.host,
+        "--port",
+        str(config.port),
+        "--disable-built-in-telemetry",
+        "--quiet",
+    ]
+    if config.use_iam_auth:
+        command.append("--auto-iam-authn")
+    if config.target_service_account:
+        command.append(f"--impersonate-service-account={config.target_service_account}")
+    if not config.use_private_ip:
+        command.append("--public-ip")
+    return command
+
+
+def ensure_broker(config: BrokerConfig) -> BrokerState:
+    """Start or reuse the profile-keyed local AlloyDB broker."""
+
+    with _broker_lock(config):
+        state = read_broker_state(config)
+        if state:
+            alive = process_is_running(state.pid)
+            listening = proxy_is_listening(host=config.host, port=config.port)
+            if alive and listening:
+                _write_broker_state(config, state)
+                return state
+            if alive:
+                raise RuntimeError(
+                    f"Broker process {state.pid} for {config.profile} exists, but "
+                    f"{config.host}:{config.port} is not accepting connections. "
+                    f"Inspect {state.log_path} or run `buildai db broker stop`."
+                )
+            _unlink_state(config)
+
+        adopted = _adopt_matching_listener(config)
+        if adopted is not None:
+            return adopted
+
+        if proxy_is_listening(host=config.host, port=config.port):
+            raise RuntimeError(
+                f"{config.host}:{config.port} is already listening, but not for the expected "
+                f"{config.profile} broker identity. Stop that process or choose another broker port."
+            )
+
+        proxy_binary = default_proxy_binary()
+        if proxy_binary is None:
+            raise RuntimeError(
+                "alloydb-auth-proxy is not installed. Install it before using brokered DB access."
+            )
+        if not config.instance_uri:
+            raise RuntimeError("AlloyDB instance URI is not configured for this environment.")
+
+        command = _broker_command(config, proxy_binary=proxy_binary)
+        with config.log_path.open("ab") as log_file:
+            process = subprocess.Popen(
+                command,
+                stdout=log_file,
+                stderr=log_file,
+                start_new_session=True,
+                env=os.environ.copy(),
+            )
+
+        deadline = time.monotonic() + 15.0
+        while time.monotonic() < deadline:
+            if process.poll() is not None:
+                adopted = _adopt_matching_listener(config)
+                if adopted is not None:
+                    return adopted
+                break
+            if proxy_is_listening(host=config.host, port=config.port):
+                state = _state_payload(config, pid=process.pid, command=command)
+                _write_broker_state(config, state)
+                return state
+            time.sleep(0.25)
+
+        if process.poll() is not None:
+            raise RuntimeError(
+                f"AlloyDB broker process exited with code {process.returncode}. "
+                f"Broker log: {config.log_path}"
+            )
+
+        raise RuntimeError(
+            f"Timed out waiting for AlloyDB broker on {config.host}:{config.port}. "
+            f"Broker log: {config.log_path}"
+        )
+
+
+def stop_broker(config: BrokerConfig) -> bool:
+    """Stop the machine-local broker for one identity."""
+
+    with _broker_lock(config):
+        state = read_broker_state(config) or _adopt_matching_listener(config)
+        if state is None:
+            return False
+
+        def stopped() -> bool:
+            return not process_is_running(state.pid) and not proxy_is_listening(
+                host=config.host,
+                port=config.port,
+            )
+
+        if process_is_running(state.pid):
+            os.kill(state.pid, 15)
+            deadline = time.monotonic() + 5.0
+            while time.monotonic() < deadline:
+                if stopped():
+                    break
+                time.sleep(0.1)
+
+            if not stopped():
+                try:
+                    os.kill(state.pid, 9)
+                except ProcessLookupError:
+                    pass
+
+                deadline = time.monotonic() + 2.0
+                while time.monotonic() < deadline:
+                    if stopped():
+                        break
+                    time.sleep(0.1)
+
+        if not stopped():
+            raise RuntimeError(
+                f"Failed to stop AlloyDB broker process {state.pid} for {config.profile}. "
+                f"State file left in place: {config.state_path}"
+            )
+
+        _unlink_state(config)
+        return True
+
+
+def broker_status(config: BrokerConfig) -> dict[str, Any]:
+    """Return a human and JSON friendly status payload for one broker."""
+
+    state = read_broker_state(config)
+    if state is None and proxy_is_listening(host=config.host, port=config.port):
+        with _broker_lock(config):
+            state = read_broker_state(config) or _adopt_matching_listener(config)
+    pid = state.pid if state else None
+    alive = process_is_running(pid) if pid is not None else False
+    listening = proxy_is_listening(host=config.host, port=config.port)
+    return {
+        "environment": config.environment,
+        "profile": config.profile,
+        "database": config.database,
+        "db_user": config.db_user,
+        "instance_uri": config.instance_uri,
+        "host": config.host,
+        "port": config.port,
+        "state_path": str(config.state_path),
+        "log_path": state.log_path if state else str(config.log_path),
+        "pid": pid,
+        "process_alive": alive,
+        "listening": listening,
+        "healthy": bool(state and alive and listening),
+    }
+
+
+def _profile_port(settings: Settings, profile: str, *, db_user: str) -> int:
+    """Return a deterministic local port for an environment/profile/user tuple."""
+
+    base = settings.effective_alloydb_auth_proxy_port
+    if profile in _PROFILE_PORT_OFFSETS:
+        return base + _PROFILE_PORT_OFFSETS[profile]
+    digest = hashlib.sha256(f"{profile}:{db_user}".encode("utf-8")).digest()
+    return base + 100 + int.from_bytes(digest[:2], "big") % 200
+
+
+def broker_config_for_identity(
+    settings: Settings,
+    *,
+    profile: str,
+    db_user: str,
+    use_iam_auth: bool,
+    target_service_account: str | None,
+    password: str = "",
+) -> BrokerConfig:
+    """Create a broker config for an explicit DB identity."""
+
+    return BrokerConfig(
+        environment=settings.app_env.value,
+        profile=profile,
+        instance_uri=settings.alloydb_instance_uri,
+        database=settings.db_name,
+        db_user=db_user,
+        use_iam_auth=use_iam_auth,
+        target_service_account=target_service_account,
+        host=settings.alloydb_auth_proxy_host,
+        port=_profile_port(settings, profile, db_user=db_user),
+        use_private_ip=settings.use_private_ip,
+        password=password,
+    )
+
+
+async def connect_via_broker(config: BrokerConfig) -> asyncpg.Connection:
+    """Ensure the broker is running and return a local asyncpg connection."""
+
+    ensure_broker(config)
+    conn = await asyncpg.connect(
+        host=config.host,
+        port=config.port,
+        user=config.db_user,
+        password=config.password if config.password else None,
+        database=config.database,
+        ssl=False,
+        statement_cache_size=0,
+    )
+    await _init_broker_connection(conn)
+    return conn
+
+
+async def open_brokered_database(config: BrokerConfig) -> BrokeredDatabase:
+    """Open one brokered connection behind a Database-like wrapper."""
+
+    return BrokeredDatabase(config=config, conn=await connect_via_broker(config))

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/cli/ops_init.py

@@ -61,6 +61,12 @@ def _apply_default_db_target(*, env_prefix: str) -> None:
     os.environ.setdefault(f"ALLOYDB_IAM_AUTH_{env_prefix}", "true")
 
 
+def _short_service_account(email: str) -> str:
+    """Return a compact service-account label for CLI banners."""
+
+    return email.split("@", 1)[0]
+
+
 def init_ops_context(ctx: typer.Context):
     """Heavy DB/auth/observability init — called lazily by ops-plane commands.
 
@@ -80,6 +86,8 @@ def init_ops_context(ctx: typer.Context):
     from infra.auth import AuthError, get_effective_user_for_iam, validate_auth_config
     from infra.settings import Environment, ExecutionContext, Settings, get_settings
 
+    from cli.auth_local import resolve_sanctioned_profile
+    from cli.context import db_identity_profile_for_cli_profile
     from infra import init_observability
 
     verbose = ctx.obj.get("_verbose", False)
@@ -114,6 +122,26 @@ def init_ops_context(ctx: typer.Context):
     auth_info: list[str] = []
     env_prefix = _auth_env_prefix(settings.app_env)
     _apply_default_db_target(env_prefix=env_prefix)
+    settings = Settings(app_env=app_env, execution_context=ExecutionContext.HUMAN)
+
+    try:
+        db_identity_profile = db_identity_profile_for_cli_profile(profile)
+        profile_contract = resolve_sanctioned_profile(
+            db_identity_profile,
+            environment=settings.app_env.value,
+        )
+    except ValueError:
+        profile_contract = None
+    profile_db_user = (
+        profile_contract.target_db_user
+        if profile_contract is not None and profile_contract.db_access
+        else None
+    )
+    profile_impersonate_sa = (
+        profile_contract.target_service_account
+        if profile_contract is not None and profile_contract.db_access
+        else None
+    )
 
     default_user = settings.effective_db_user
     default_use_iam = settings.effective_use_iam_auth
@@ -121,15 +149,15 @@ def init_ops_context(ctx: typer.Context):
 
     if auth_flag is not None:
         use_iam_auth = auth_flag == AuthMethod.IAM
-    elif profile == "engineers-dev" and not ctx.obj.get("allow_write"):
+    elif profile_db_user and profile_impersonate_sa:
         use_iam_auth = True
     else:
         use_iam_auth = default_use_iam
 
     if user_flag is not None:
         effective_user = user_flag
-    elif profile == "engineers-dev" and not ctx.obj.get("allow_write"):
-        effective_user = "engineers-dev-sa@data-470400.iam"
+    elif profile_db_user:
+        effective_user = profile_db_user
     elif use_iam_auth:
         try:
             effective_user = get_effective_user_for_iam(None, default_user)
@@ -158,12 +186,11 @@ def init_ops_context(ctx: typer.Context):
 
     os.environ[f"ALLOYDB_USER_{env_prefix}"] = effective_user
     auth_info.append(f"user={effective_user}")
-    if profile == "engineers-dev" and not ctx.obj.get("allow_write"):
-        os.environ[f"ALLOYDB_RUNTIME_IMPERSONATE_SA_{env_prefix}"] = (
-            "engineers-dev-sa@data-470400.iam.gserviceaccount.com"
-        )
-        os.environ[f"ALLOYDB_USE_AUTH_PROXY_{env_prefix}"] = "false"
-        auth_info.append("impersonate=engineers-dev-sa")
+    if profile_impersonate_sa:
+        os.environ[f"ALLOYDB_RUNTIME_IMPERSONATE_SA_{env_prefix}"] = profile_impersonate_sa
+        os.environ[f"ALLOYDB_USE_AUTH_PROXY_{env_prefix}"] = "true"
+        auth_info.append(f"impersonate={_short_service_account(profile_impersonate_sa)}")
+        auth_info.append("broker=alloydb-auth-proxy")
 
     # Reload settings with overrides
     get_settings.cache_clear()

{buildai_cli-0.3.53 → buildai_cli-0.3.55}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "hatchling.build"
 
 [project]
 name = "buildai-cli"
-version = "0.3.53"
+version = "0.3.55"
 description = "Build AI CLI (Typer)"
 requires-python = ">=3.11"
 dependencies = [

buildai_cli-0.3.53/cli/auth_proxy.py

@@ -1,99 +0,0 @@
-"""Helpers for the sanctioned local AlloyDB Auth Proxy lane."""
-
-from __future__ import annotations
-
-import os
-import shutil
-import socket
-import subprocess
-import time
-from pathlib import Path
-
-from infra.settings import Settings
-
-
-def _default_proxy_binary() -> str | None:
-    """Return the preferred AlloyDB Auth Proxy binary path if available."""
-
-    on_path = shutil.which("alloydb-auth-proxy")
-    if on_path:
-        return on_path
-    local_bin = Path.home() / ".local" / "bin" / "alloydb-auth-proxy"
-    if local_bin.is_file():
-        return str(local_bin)
-    return None
-
-
-def proxy_is_listening(*, host: str, port: int) -> bool:
-    """Return True when a TCP listener is already present on the proxy endpoint."""
-
-    with socket.socket(socket.AF_INET, socket.SOCK_STREAM) as sock:
-        sock.settimeout(0.25)
-        return sock.connect_ex((host, port)) == 0
-
-
-def ensure_alloydb_auth_proxy(
-    *,
-    settings: Settings,
-    target_service_account: str,
-) -> None:
-    """Start the sanctioned AlloyDB Auth Proxy if it is not already running.
-
-    The engineer DB lane should use one explicit transport path across CLI,
-    `psql`, and desktop tools. If the proxy is already listening, leave it
-    alone. Otherwise start it in the background and wait briefly for the port
-    to come up.
-    """
-
-    if not settings.effective_use_alloydb_auth_proxy:
-        return
-
-    host = settings.alloydb_auth_proxy_host
-    port = settings.effective_alloydb_auth_proxy_port
-    if proxy_is_listening(host=host, port=port):
-        return
-
-    proxy_binary = _default_proxy_binary()
-    if proxy_binary is None:
-        raise RuntimeError(
-            "alloydb-auth-proxy is not installed. Re-run ./scripts/setup/setup.py to provision "
-            "the canonical local DB transport."
-        )
-
-    if not settings.alloydb_instance_uri:
-        raise RuntimeError(
-            "AlloyDB instance URI is not configured, so the local Auth Proxy cannot start."
-        )
-
-    command = [
-        proxy_binary,
-        settings.alloydb_instance_uri,
-        "--address",
-        host,
-        "--port",
-        str(port),
-        "--auto-iam-authn",
-        f"--impersonate-service-account={target_service_account}",
-        "--disable-built-in-telemetry",
-    ]
-    if not settings.use_private_ip:
-        command.append("--public-ip")
-
-    subprocess.Popen(
-        command,
-        stdout=subprocess.DEVNULL,
-        stderr=subprocess.DEVNULL,
-        start_new_session=True,
-        env=os.environ.copy(),
-    )
-
-    deadline = time.monotonic() + 10.0
-    while time.monotonic() < deadline:
-        if proxy_is_listening(host=host, port=port):
-            return
-        time.sleep(0.25)
-
-    raise RuntimeError(
-        f"Timed out waiting for alloydb-auth-proxy on {host}:{port}. "
-        "Run `uv run buildai dev db info` to inspect the expected proxy recipe."
-    )