airbyte-internal-ops 0.1.9__py3-none-any.whl → 0.1.10__py3-none-any.whl

{airbyte_internal_ops-0.1.9.dist-info → airbyte_internal_ops-0.1.10.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: airbyte-internal-ops
-Version: 0.1.9
+Version: 0.1.10
 Summary: MCP and API interfaces that let the agents do the admin work
 Author-email: Aaron Steers <aj@airbyte.io>
 Keywords: admin,airbyte,api,mcp

{airbyte_internal_ops-0.1.9.dist-info → airbyte_internal_ops-0.1.10.dist-info}/RECORD

@@ -1,6 +1,6 @@
 airbyte_ops_mcp/__init__.py,sha256=HhzURuYr29_UIdMrnWYaZB8ENr_kFkBdm4uqeiIW3Vw,760
 airbyte_ops_mcp/_annotations.py,sha256=MO-SBDnbykxxHDESG7d8rviZZ4WlZgJKv0a8eBqcEzQ,1757
-airbyte_ops_mcp/constants.py,sha256=419-AlRfwbbxeEEV9lhmXhpTUjsSdzJpfcuL_MZZtXM,1982
+airbyte_ops_mcp/constants.py,sha256=col6-5BUWuIYhbtKmlvSRR8URBoSNExoz94cn4_kujI,2333
 airbyte_ops_mcp/gcp_auth.py,sha256=5k-k145ZoYhHLjyDES8nrA8f8BBihRI0ykrdD1IcfOs,3599
 airbyte_ops_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 airbyte_ops_mcp/_legacy/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -351,7 +351,7 @@ airbyte_ops_mcp/cli/__init__.py,sha256=XpL7FyVfgabfBF2JR7u7NwJ2krlYqjd_OwLcWf-Xc
 airbyte_ops_mcp/cli/_base.py,sha256=I8tWnyQf0ks4r3J8N8h-5GZxyn37T-55KsbuHnxYlcg,415
 airbyte_ops_mcp/cli/_shared.py,sha256=jg-xMyGzTCGPqKd8VTfE_3kGPIyO_3Kx5sQbG4rPc0Y,1311
 airbyte_ops_mcp/cli/app.py,sha256=SEdBpqFUG2O8zGV5ifwptxrLGFph_dLr66-MX9d69gQ,789
-airbyte_ops_mcp/cli/cloud.py,sha256=OTc8o_MIw5DGGPnQ2kiUT98GCXS2I95vvsSN9_d7LVo,32077
+airbyte_ops_mcp/cli/cloud.py,sha256=semAuPqruWfe7JiVmtbELZgcpeqtoIn4LGD0kEvIS30,38981
 airbyte_ops_mcp/cli/gh.py,sha256=91b1AxFXvHQCFyXhrrym-756ZjnMCqvxFdmwCtma1zI,2046
 airbyte_ops_mcp/cli/registry.py,sha256=-yiLJWSslV_qGi6ImXZYfXOJSE4oJBO7yICkyA_RiUo,5792
 airbyte_ops_mcp/cli/repo.py,sha256=G1hoQpH0XYhUH3FFOsia9xabGB0LP9o3XcwBuqvFVo0,16331
@@ -400,14 +400,14 @@ airbyte_ops_mcp/mcp/registry.py,sha256=PW-VYUj42qx2pQ_apUkVaoUFq7VgB9zEU7-aGrkSC
 airbyte_ops_mcp/mcp/server.py,sha256=7zi91xioVTx1q-bEleekZH2c2JnbzDQt_6zxdEwnLbg,2958
 airbyte_ops_mcp/mcp/server_info.py,sha256=Yi4B1auW64QZGBDas5mro_vwTjvrP785TFNSBP7GhRg,2361
 airbyte_ops_mcp/prod_db_access/__init__.py,sha256=5pxouMPY1beyWlB0UwPnbaLTKTHqU6X82rbbgKY2vYU,1069
-airbyte_ops_mcp/prod_db_access/db_engine.py,sha256=ia1KcuQOXi3Qhy_MnxYmccCBJ4rAt_d4nVDjcyzje6o,4289
+airbyte_ops_mcp/prod_db_access/db_engine.py,sha256=11xNZTk4I8SKYhsnmE7-LVrkJXN4dCRbBeD1_hj3f-s,9027
 airbyte_ops_mcp/prod_db_access/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 airbyte_ops_mcp/prod_db_access/queries.py,sha256=q7PcI15EGh6jFS9MVB_gZt1a56YvrZV5hnwa5lgU2q0,10844
 airbyte_ops_mcp/prod_db_access/sql.py,sha256=tWQAwMk8DzG8HpLIYglljlReI2oeYulQPsV31ocUJSw,16251
 airbyte_ops_mcp/registry/__init__.py,sha256=iEaPlt9GrnlaLbc__98TguNeZG8wuQu7S-_2QkhHcbA,858
 airbyte_ops_mcp/registry/models.py,sha256=B4L4TKr52wo0xs0CqvCBrpowqjShzVnZ5eTr2-EyhNs,2346
 airbyte_ops_mcp/registry/publish.py,sha256=VoPxsM2_0zJ829orzCRN-kjgcJtuBNyXgW4I9J680ro,12717
-airbyte_internal_ops-0.1.9.dist-info/METADATA,sha256=1tBlf96RtcJNsKARaZt4711QiVISWcAkex4OLBlUmjk,5282
-airbyte_internal_ops-0.1.9.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
-airbyte_internal_ops-0.1.9.dist-info/entry_points.txt,sha256=eUgJ9xIy9PlR-CgRbqRMsh1NVp6jz08v9bul9vCYlU4,111
-airbyte_internal_ops-0.1.9.dist-info/RECORD,,
+airbyte_internal_ops-0.1.10.dist-info/METADATA,sha256=I3WAz4JM1tgyGMcFaV5ElRwvwYjlRLR9kXKzifetd1Q,5283
+airbyte_internal_ops-0.1.10.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
+airbyte_internal_ops-0.1.10.dist-info/entry_points.txt,sha256=eUgJ9xIy9PlR-CgRbqRMsh1NVp6jz08v9bul9vCYlU4,111
+airbyte_internal_ops-0.1.10.dist-info/RECORD,,

airbyte_ops_mcp/cli/cloud.py

@@ -13,6 +13,12 @@ Commands:
 from __future__ import annotations
 
 import json
+import os
+import shutil
+import signal
+import socket
+import subprocess
+import time
 from pathlib import Path
 from typing import Annotated, Literal
 
@@ -30,6 +36,12 @@ from airbyte_ops_mcp.cli._shared import (
     print_success,
 )
 from airbyte_ops_mcp.cloud_admin.connection_config import fetch_connection_config
+from airbyte_ops_mcp.constants import (
+    CLOUD_SQL_INSTANCE,
+    CLOUD_SQL_PROXY_PID_FILE,
+    DEFAULT_CLOUD_SQL_PROXY_PORT,
+    ENV_GCP_PROD_DB_ACCESS_CREDENTIALS,
+)
 from airbyte_ops_mcp.live_tests.ci_output import (
     generate_regression_report,
     get_report_summary,
@@ -75,6 +87,184 @@ connector_app = App(
 )
 cloud_app.command(connector_app)
 
+# Create the db sub-app under cloud
+db_app = App(name="db", help="Database operations for Airbyte Cloud Prod DB Replica.")
+cloud_app.command(db_app)
+
+
+@db_app.command(name="start-proxy")
+def start_proxy(
+    port: Annotated[
+        int,
+        Parameter(help="Port for the Cloud SQL Proxy to listen on."),
+    ] = DEFAULT_CLOUD_SQL_PROXY_PORT,
+    daemon: Annotated[
+        bool,
+        Parameter(
+            help="Run as daemon in background (default). Use --no-daemon for foreground."
+        ),
+    ] = True,
+) -> None:
+    """Start the Cloud SQL Proxy for database access.
+
+    This command starts the Cloud SQL Auth Proxy to enable connections to the
+    Airbyte Cloud Prod DB Replica. The proxy is required for database query tools.
+
+    By default, runs as a daemon (background process). Use --no-daemon to run in
+    foreground mode where you can see logs and stop with Ctrl+C.
+
+    Credentials are read from the GCP_PROD_DB_ACCESS_CREDENTIALS environment variable,
+    which should contain the service account JSON credentials.
+
+    After starting the proxy, set these environment variables to use database tools:
+        export USE_CLOUD_SQL_PROXY=1
+        export DB_PORT={port}
+
+    Example:
+        airbyte-ops cloud db start-proxy
+        airbyte-ops cloud db start-proxy --port 15432
+        airbyte-ops cloud db start-proxy --no-daemon
+    """
+    # Check if proxy is already running on the requested port (idempotency)
+    try:
+        with socket.create_connection(("127.0.0.1", port), timeout=0.5):
+            # Something is already listening on this port
+            pid_file = Path(CLOUD_SQL_PROXY_PID_FILE)
+            pid_info = ""
+            if pid_file.exists():
+                pid_info = f" (PID: {pid_file.read_text().strip()})"
+            print_success(
+                f"Cloud SQL Proxy is already running on port {port}{pid_info}"
+            )
+            print_success("")
+            print_success("To use database tools, set these environment variables:")
+            print_success("  export USE_CLOUD_SQL_PROXY=1")
+            print_success(f"  export DB_PORT={port}")
+            return
+    except (OSError, TimeoutError, ConnectionRefusedError):
+        pass  # Port not in use, proceed with starting proxy
+
+    # Check if cloud-sql-proxy is installed
+    proxy_path = shutil.which("cloud-sql-proxy")
+    if not proxy_path:
+        exit_with_error(
+            "cloud-sql-proxy not found in PATH. "
+            "Install it from: https://cloud.google.com/sql/docs/mysql/sql-proxy"
+        )
+
+    # Get credentials from environment
+    creds_json = os.getenv(ENV_GCP_PROD_DB_ACCESS_CREDENTIALS)
+    if not creds_json:
+        exit_with_error(
+            f"{ENV_GCP_PROD_DB_ACCESS_CREDENTIALS} environment variable is not set. "
+            "This should contain the GCP service account JSON credentials."
+        )
+
+    # Build the command using --json-credentials to avoid writing to disk
+    cmd = [
+        proxy_path,
+        CLOUD_SQL_INSTANCE,
+        f"--port={port}",
+        f"--json-credentials={creds_json}",
+    ]
+
+    print_success(f"Starting Cloud SQL Proxy on port {port}...")
+    print_success(f"Instance: {CLOUD_SQL_INSTANCE}")
+    print_success("")
+    print_success("To use database tools, set these environment variables:")
+    print_success("  export USE_CLOUD_SQL_PROXY=1")
+    print_success(f"  export DB_PORT={port}")
+    print_success("")
+
+    if daemon:
+        # Run in background (daemon mode) with log file for diagnostics
+        log_file_path = Path("/tmp/airbyte-cloud-sql-proxy.log")
+        log_file = log_file_path.open("ab")
+        process = subprocess.Popen(
+            cmd,
+            stdout=subprocess.DEVNULL,
+            stderr=log_file,
+            start_new_session=True,
+        )
+
+        # Brief wait to verify the process started successfully
+        time.sleep(0.5)
+        if process.poll() is not None:
+            # Process exited immediately - read any error output
+            log_file.close()
+            error_output = ""
+            if log_file_path.exists():
+                error_output = log_file_path.read_text()[-1000:]  # Last 1000 chars
+            exit_with_error(
+                f"Cloud SQL Proxy failed to start (exit code: {process.returncode}).\n"
+                f"Check logs at {log_file_path}\n"
+                f"Recent output: {error_output}"
+            )
+
+        # Write PID to file for stop-proxy command
+        pid_file = Path(CLOUD_SQL_PROXY_PID_FILE)
+        pid_file.write_text(str(process.pid))
+        print_success(f"Cloud SQL Proxy started as daemon (PID: {process.pid})")
+        print_success(f"Logs: {log_file_path}")
+        print_success("To stop: airbyte-ops cloud db stop-proxy")
+    else:
+        # Run in foreground - replace current process
+        # Signals (Ctrl+C) will be handled directly by the cloud-sql-proxy process
+        print_success("Running in foreground. Press Ctrl+C to stop the proxy.")
+        print_success("")
+        os.execv(proxy_path, cmd)
+
+
+@db_app.command(name="stop-proxy")
+def stop_proxy() -> None:
+    """Stop the Cloud SQL Proxy daemon.
+
+    This command stops a Cloud SQL Proxy that was started with 'start-proxy'.
+    It reads the PID from the PID file and sends a SIGTERM signal to stop the process.
+
+    Example:
+        airbyte-ops cloud db stop-proxy
+    """
+    pid_file = Path(CLOUD_SQL_PROXY_PID_FILE)
+
+    if not pid_file.exists():
+        exit_with_error(
+            f"PID file not found at {CLOUD_SQL_PROXY_PID_FILE}. "
+            "No Cloud SQL Proxy daemon appears to be running."
+        )
+
+    pid_str = pid_file.read_text().strip()
+    if not pid_str.isdigit():
+        pid_file.unlink()
+        exit_with_error(f"Invalid PID in {CLOUD_SQL_PROXY_PID_FILE}: {pid_str}")
+
+    pid = int(pid_str)
+
+    # Check if process is still running
+    try:
+        os.kill(pid, 0)  # Signal 0 just checks if process exists
+    except ProcessLookupError:
+        pid_file.unlink()
+        print_success(
+            f"Cloud SQL Proxy (PID: {pid}) is not running. Cleaned up PID file."
+        )
+        return
+    except PermissionError:
+        exit_with_error(f"Permission denied to check process {pid}.")
+
+    # Send SIGTERM to stop the process
+    try:
+        os.kill(pid, signal.SIGTERM)
+        print_success(f"Sent SIGTERM to Cloud SQL Proxy (PID: {pid}).")
+    except ProcessLookupError:
+        print_success(f"Cloud SQL Proxy (PID: {pid}) already stopped.")
+    except PermissionError:
+        exit_with_error(f"Permission denied to stop process {pid}.")
+
+    # Clean up PID file
+    pid_file.unlink(missing_ok=True)
+    print_success("Cloud SQL Proxy stopped.")
+
 
 @connector_app.command(name="get-version-info")
 def get_version_info(

airbyte_ops_mcp/constants.py

@@ -27,6 +27,15 @@ EXPECTED_ADMIN_EMAIL_DOMAIN = "@airbyte.io"
 GCP_PROJECT_NAME = "prod-ab-cloud-proj"
 """The GCP project name for Airbyte Cloud production."""
 
+CLOUD_SQL_INSTANCE = "prod-ab-cloud-proj:us-west3:prod-pgsql-replica"
+"""The Cloud SQL instance connection name for the Prod DB Replica."""
+
+DEFAULT_CLOUD_SQL_PROXY_PORT = 15432
+"""Default port for Cloud SQL Proxy connections."""
+
+CLOUD_SQL_PROXY_PID_FILE = "/tmp/airbyte-cloud-sql-proxy.pid"
+"""PID file for tracking the Cloud SQL Proxy process."""
+
 CLOUD_REGISTRY_URL = (
     "https://connectors.airbyte.com/files/registries/v0/cloud_registry.json"
 )

airbyte_ops_mcp/prod_db_access/db_engine.py

@@ -11,7 +11,9 @@ from __future__ import annotations
 
 import json
 import os
-import traceback
+import shutil
+import socket
+import subprocess
 from typing import Any, Callable
 
 import sqlalchemy
@@ -21,9 +23,127 @@ from google.cloud.sql.connector.enums import IPTypes
 
 from airbyte_ops_mcp.constants import (
     CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID,
+    DEFAULT_CLOUD_SQL_PROXY_PORT,
 )
 
 PG_DRIVER = "pg8000"
+PROXY_CHECK_TIMEOUT = 0.5  # seconds
+DIRECT_CONNECTION_TIMEOUT = 5  # seconds - timeout for direct VPC/Tailscale connections
+
+
+class CloudSqlProxyNotRunningError(Exception):
+    """Raised when proxy mode is enabled but the Cloud SQL Proxy is not running."""
+
+    pass
+
+
+class VpnNotConnectedError(Exception):
+    """Raised when direct connection mode requires VPN but it's not connected."""
+
+    pass
+
+
+def _is_tailscale_connected() -> bool:
+    """Check if Tailscale VPN is likely connected.
+
+    This is a best-effort check that works on Linux and macOS.
+    Returns True if Tailscale appears to be connected, False otherwise.
+
+    Detection methods:
+    1. Check for tailscale0 network interface (Linux)
+    2. Run 'tailscale status --json' and check backend state (cross-platform)
+    """
+    # Method 1: Check for tailscale0 interface (Linux)
+    try:
+        interfaces = [name for _, name in socket.if_nameindex()]
+        if "tailscale0" in interfaces:
+            return True
+    except (OSError, AttributeError):
+        pass  # if_nameindex not available on this platform
+
+    # Method 2: Check tailscale CLI status
+    tailscale_path = shutil.which("tailscale")
+    if tailscale_path:
+        try:
+            result = subprocess.run(
+                [tailscale_path, "status", "--json"],
+                capture_output=True,
+                text=True,
+                timeout=2,
+            )
+            if result.returncode == 0:
+                import json as json_module
+
+                status = json_module.loads(result.stdout)
+                # BackendState "Running" indicates connected
+                return status.get("BackendState") == "Running"
+        except (subprocess.TimeoutExpired, subprocess.SubprocessError, ValueError):
+            pass
+
+    return False
+
+
+def _check_vpn_or_proxy_available() -> None:
+    """Check if either VPN or proxy is available for database access.
+
+    This function checks if the environment is properly configured for
+    database access. It fails fast with a helpful error message if neither
+    Tailscale VPN nor the Cloud SQL Proxy appears to be available.
+
+    Raises:
+        VpnNotConnectedError: If no VPN or proxy is detected
+    """
+    # If proxy mode is explicitly enabled, don't check VPN
+    if os.getenv("CI") or os.getenv("USE_CLOUD_SQL_PROXY"):
+        return
+
+    # Check if Tailscale is connected
+    if _is_tailscale_connected():
+        return
+
+    # Neither proxy mode nor Tailscale detected
+    raise VpnNotConnectedError(
+        "No VPN or proxy detected for database access.\n\n"
+        "To connect to the Airbyte Cloud Prod DB Replica, you need either:\n\n"
+        "1. Tailscale VPN connected (for direct VPC access)\n"
+        "   - Install Tailscale: https://tailscale.com/download\n"
+        "   - Connect to the Airbyte network\n\n"
+        "2. Cloud SQL Proxy running locally\n"
+        "   - Start the proxy:\n"
+        "       airbyte-ops cloud db start-proxy\n"
+        "       uvx --from=airbyte-internal-ops airbyte-ops cloud db start-proxy\n"
+        "   - Set env vars: export USE_CLOUD_SQL_PROXY=1 DB_PORT=15432\n"
+    )
+
+
+def _check_proxy_is_running(host: str, port: int) -> None:
+    """Check if the Cloud SQL Proxy is running and accepting connections.
+
+    This performs a quick socket connection check to fail fast if the proxy
+    is not running, rather than waiting for a long connection timeout.
+
+    Args:
+        host: The host to connect to (typically 127.0.0.1)
+        port: The port to connect to
+
+    Raises:
+        CloudSqlProxyNotRunningError: If the proxy is not accepting connections
+    """
+    try:
+        with socket.create_connection((host, port), timeout=PROXY_CHECK_TIMEOUT):
+            pass  # Connection successful, proxy is running
+    except (OSError, TimeoutError, ConnectionRefusedError) as e:
+        raise CloudSqlProxyNotRunningError(
+            f"Cloud SQL Proxy is not running on {host}:{port}. "
+            f"Proxy mode is enabled (CI or USE_CLOUD_SQL_PROXY env var is set), "
+            f"but nothing is listening on the expected port.\n\n"
+            f"To start the proxy, run:\n"
+            f"  airbyte-ops cloud db start-proxy --port {port}\n"
+            f"  uvx --from=airbyte-internal-ops airbyte-ops cloud db start-proxy --port {port}\n\n"
+            f"Or unset USE_CLOUD_SQL_PROXY to use direct VPC connection.\n\n"
+            f"Original error: {e}"
+        ) from e
+
 
 # Lazy-initialized to avoid import-time GCP auth
 _connector: Connector | None = None
@@ -81,16 +201,20 @@ def get_pool(
     """Get a SQLAlchemy connection pool for the Airbyte Cloud database.
 
     This function supports two connection modes:
-    1. Direct connection via Cloud SQL Python Connector (default, requires VPC access)
+    1. Direct connection via Cloud SQL Python Connector (default, requires VPC/Tailscale)
     2. Connection via Cloud SQL Auth Proxy (when CI or USE_CLOUD_SQL_PROXY env var is set)
 
     For proxy mode, start the proxy with:
-        cloud-sql-proxy prod-ab-cloud-proj:us-west3:prod-pgsql-replica --port=<port>
+        airbyte-ops cloud db start-proxy
 
     Environment variables:
         CI: If set, uses proxy connection mode
         USE_CLOUD_SQL_PROXY: If set, uses proxy connection mode
-        DB_PORT: Port for proxy connection (default: 5432)
+        DB_PORT: Port for proxy connection (default: 15432)
+
+    Raises:
+        VpnNotConnectedError: If direct mode is used but no VPN/proxy is detected
+        CloudSqlProxyNotRunningError: If proxy mode is enabled but the proxy is not running
 
     Args:
         gsm_client: GCP Secret Manager client for retrieving credentials
@@ -98,6 +222,9 @@ def get_pool(
     Returns:
         SQLAlchemy Engine connected to the Prod DB Replica
     """
+    # Fail fast if no VPN or proxy is available
+    _check_vpn_or_proxy_available()
+
     pg_connection_details = json.loads(
         _get_secret_value(
             gsm_client, CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID
@@ -106,21 +233,21 @@ def get_pool(
 
     if os.getenv("CI") or os.getenv("USE_CLOUD_SQL_PROXY"):
         # Connect via Cloud SQL Auth Proxy, running on localhost
-        # Port can be configured via DB_PORT env var (default: 5432)
+        # Port can be configured via DB_PORT env var (default: DEFAULT_CLOUD_SQL_PROXY_PORT)
         host = "127.0.0.1"
-        port = os.getenv("DB_PORT", "5432")
-        try:
-            return sqlalchemy.create_engine(
-                f"postgresql+{PG_DRIVER}://{pg_connection_details['pg_user']}:{pg_connection_details['pg_password']}@{host}:{port}/{pg_connection_details['database_name']}",
-            )
-        except Exception as e:
-            raise AssertionError(
-                f"sqlalchemy.create_engine exception; could not connect to the proxy at {host}:{port}. "
-                f"Error: {traceback.format_exception(e)}"
-            ) from e
+        port = int(os.getenv("DB_PORT", str(DEFAULT_CLOUD_SQL_PROXY_PORT)))
+
+        # Fail fast if proxy is not running
+        _check_proxy_is_running(host, port)
+
+        return sqlalchemy.create_engine(
+            f"postgresql+{PG_DRIVER}://{pg_connection_details['pg_user']}:{pg_connection_details['pg_password']}@{host}:{port}/{pg_connection_details['database_name']}",
+        )
 
-    # Default: Connect via Cloud SQL Python Connector (requires VPC access)
+    # Default: Connect via Cloud SQL Python Connector (requires VPC/Tailscale access)
+    # Use a timeout to fail faster if the connection can't be established
     return sqlalchemy.create_engine(
         f"postgresql+{PG_DRIVER}://",
         creator=get_database_creator(pg_connection_details),
+        connect_args={"timeout": DIRECT_CONNECTION_TIMEOUT},
     )