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

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},
     )