airbyte-internal-ops 0.3.1__py3-none-any.whl → 0.4.1__py3-none-any.whl

airbyte_ops_mcp/mcp/regression_tests.py

@@ -31,12 +31,10 @@ from airbyte.exceptions import (
 from fastmcp import FastMCP
 from pydantic import BaseModel, Field
 
-from airbyte_ops_mcp.github_actions import (
-    GITHUB_API_BASE,
-    resolve_github_token,
-    trigger_workflow_dispatch,
-)
+from airbyte_ops_mcp.github_actions import trigger_workflow_dispatch
+from airbyte_ops_mcp.github_api import GITHUB_API_BASE, resolve_github_token
 from airbyte_ops_mcp.mcp._mcp_utils import mcp_tool, register_mcp_tools
+from airbyte_ops_mcp.mcp.prerelease import ConnectorRepo
 
 logger = logging.getLogger(__name__)
 
@@ -284,7 +282,7 @@ class RunRegressionTestsResponse(BaseModel):
     )
     github_run_id: int | None = Field(
         default=None,
-        description="GitHub Actions workflow run ID (use with check_workflow_status)",
+        description="GitHub Actions workflow run ID (use with check_ci_workflow_status)",
     )
     github_run_url: str | None = Field(
         default=None,
@@ -309,7 +307,13 @@ def run_regression_tests(
     ],
     pr: Annotated[
         int,
-        "PR number from the airbyte monorepo to checkout and build from (e.g., 70847). Required.",
+        "PR number to checkout and build from (e.g., 70847). Required. "
+        "The PR must be from the repository specified by the 'repo' parameter.",
+    ],
+    repo: Annotated[
+        ConnectorRepo,
+        "Repository where the connector PR is located. "
+        "Use 'airbyte' for OSS connectors (default) or 'airbyte-enterprise' for enterprise connectors.",
     ],
     connection_id: Annotated[
         str | None,
@@ -347,6 +351,10 @@ def run_regression_tests(
     This tool triggers the regression test workflow which builds the connector
     from the specified PR and runs tests against it.
 
+    Supports both OSS connectors (from airbytehq/airbyte) and enterprise connectors
+    (from airbytehq/airbyte-enterprise). Use the 'repo' parameter to specify which
+    repository contains the connector PR.
+
     - skip_compare=False (default): Comparison mode - compares the PR version
       against the baseline (control) version.
     - skip_compare=True: Single-version mode - runs tests without comparison.
@@ -390,10 +398,11 @@ def run_regression_tests(
                 workflow_url=None,
             )
 
-    # Build workflow inputs - connector_name and pr are required
+    # Build workflow inputs - connector_name, pr, and repo are required
     workflow_inputs: dict[str, str] = {
         "connector_name": connector_name,
         "pr": str(pr),
+        "repo": repo,
     }
 
     # Add optional inputs
@@ -409,6 +418,7 @@ def run_regression_tests(
         workflow_inputs["override_control_image"] = override_control_image
 
     mode_description = "single-version" if skip_compare else "comparison"
+
     try:
         dispatch_result = trigger_workflow_dispatch(
             owner=REGRESSION_TEST_REPO_OWNER,
@@ -431,12 +441,13 @@ def run_regression_tests(
 
     view_url = dispatch_result.run_url or dispatch_result.workflow_url
     connection_info = f" for connection {connection_id}" if connection_id else ""
+    repo_info = f" from {repo}" if repo != ConnectorRepo.AIRBYTE else ""
     return RunRegressionTestsResponse(
         run_id=run_id,
         status=TestRunStatus.QUEUED,
         message=(
             f"{mode_description.capitalize()} regression test workflow triggered "
-            f"for {connector_name} (PR #{pr}){connection_info}. View progress at: {view_url}"
+            f"for {connector_name} (PR #{pr}{repo_info}){connection_info}. View progress at: {view_url}"
         ),
         workflow_url=dispatch_result.workflow_url,
         github_run_id=dispatch_result.run_id,

airbyte_ops_mcp/mcp/server.py

@@ -25,7 +25,7 @@ from airbyte_ops_mcp.mcp.cloud_connector_versions import (
     register_cloud_connector_version_tools,
 )
 from airbyte_ops_mcp.mcp.gcp_logs import register_gcp_logs_tools
-from airbyte_ops_mcp.mcp.github import register_github_tools
+from airbyte_ops_mcp.mcp.github_actions import register_github_actions_tools
 from airbyte_ops_mcp.mcp.github_repo_ops import register_github_repo_ops_tools
 from airbyte_ops_mcp.mcp.prerelease import register_prerelease_tools
 from airbyte_ops_mcp.mcp.prod_db_queries import register_prod_db_query_tools
@@ -59,7 +59,7 @@ def register_server_assets(app: FastMCP) -> None:
     """
     register_server_info_resources(app)
     register_github_repo_ops_tools(app)
-    register_github_tools(app)
+    register_github_actions_tools(app)
     register_prerelease_tools(app)
     register_cloud_connector_version_tools(app)
     register_prod_db_query_tools(app)

airbyte_ops_mcp/prod_db_access/queries.py

@@ -26,6 +26,7 @@ from airbyte_ops_mcp.prod_db_access.sql import (
     SELECT_CONNECTIONS_BY_DESTINATION_CONNECTOR_AND_ORG,
     SELECT_CONNECTOR_VERSIONS,
     SELECT_DATAPLANES_LIST,
+    SELECT_DESTINATION_CONNECTION_STATS,
     SELECT_FAILED_SYNC_ATTEMPTS_FOR_CONNECTOR,
     SELECT_NEW_CONNECTOR_RELEASES,
     SELECT_ORG_WORKSPACES,
@@ -35,6 +36,7 @@ from airbyte_ops_mcp.prod_db_access.sql import (
     SELECT_RECENT_SUCCESSFUL_SYNCS_FOR_SOURCE_CONNECTOR,
     SELECT_RECENT_SYNCS_FOR_DESTINATION_CONNECTOR,
     SELECT_RECENT_SYNCS_FOR_SOURCE_CONNECTOR,
+    SELECT_SOURCE_CONNECTION_STATS,
     SELECT_SUCCESSFUL_SYNCS_FOR_VERSION,
     SELECT_SYNC_RESULTS_FOR_VERSION,
     SELECT_WORKSPACE_INFO,
@@ -497,3 +499,69 @@ def query_workspaces_by_email_domain(
         query_name="SELECT_WORKSPACES_BY_EMAIL_DOMAIN",
         gsm_client=gsm_client,
     )
+
+
+def query_source_connection_stats(
+    connector_definition_id: str,
+    days: int = 7,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query aggregate connection stats for a SOURCE connector.
+
+    Returns counts of connections grouped by pinned version, including:
+    - Total, enabled, and active connection counts
+    - Pinned vs unpinned breakdown
+    - Latest attempt status breakdown (succeeded, failed, cancelled, running, unknown)
+
+    Args:
+        connector_definition_id: Source connector definition UUID
+        days: Number of days to look back for "active" connections (default: 7)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of dicts with aggregate counts grouped by pinned_version_id
+    """
+    cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
+    return _run_sql_query(
+        SELECT_SOURCE_CONNECTION_STATS,
+        parameters={
+            "connector_definition_id": connector_definition_id,
+            "cutoff_date": cutoff_date,
+        },
+        query_name="SELECT_SOURCE_CONNECTION_STATS",
+        gsm_client=gsm_client,
+    )
+
+
+def query_destination_connection_stats(
+    connector_definition_id: str,
+    days: int = 7,
+    *,
+    gsm_client: secretmanager.SecretManagerServiceClient | None = None,
+) -> list[dict[str, Any]]:
+    """Query aggregate connection stats for a DESTINATION connector.
+
+    Returns counts of connections grouped by pinned version, including:
+    - Total, enabled, and active connection counts
+    - Pinned vs unpinned breakdown
+    - Latest attempt status breakdown (succeeded, failed, cancelled, running, unknown)
+
+    Args:
+        connector_definition_id: Destination connector definition UUID
+        days: Number of days to look back for "active" connections (default: 7)
+        gsm_client: GCP Secret Manager client. If None, a new client will be instantiated.
+
+    Returns:
+        List of dicts with aggregate counts grouped by pinned_version_id
+    """
+    cutoff_date = datetime.now(timezone.utc) - timedelta(days=days)
+    return _run_sql_query(
+        SELECT_DESTINATION_CONNECTION_STATS,
+        parameters={
+            "connector_definition_id": connector_definition_id,
+            "cutoff_date": cutoff_date,
+        },
+        query_name="SELECT_DESTINATION_CONNECTION_STATS",
+        gsm_client=gsm_client,
+    )

airbyte_ops_mcp/prod_db_access/sql.py

@@ -826,3 +826,111 @@ SELECT_WORKSPACES_BY_EMAIL_DOMAIN = sqlalchemy.text(
     LIMIT :limit
     """
 )
+
+# =============================================================================
+# Connector Connection Stats Queries (Aggregate Counts)
+# =============================================================================
+
+# Count connections by SOURCE connector with latest attempt status breakdown
+# Groups by pinned version and provides counts of succeeded/failed/other attempts
+# Uses a CTE to get the latest attempt per connection, then aggregates
+SELECT_SOURCE_CONNECTION_STATS = sqlalchemy.text(
+    """
+    WITH latest_attempts AS (
+        SELECT DISTINCT ON (connection.id)
+            connection.id AS connection_id,
+            connection.status AS connection_status,
+            scoped_configuration.value AS pinned_version_id,
+            attempts.status::text AS latest_attempt_status
+        FROM connection
+        JOIN actor
+          ON connection.source_id = actor.id
+         AND actor.actor_definition_id = :connector_definition_id
+         AND actor.tombstone = false
+        JOIN workspace
+          ON actor.workspace_id = workspace.id
+         AND workspace.tombstone = false
+        LEFT JOIN scoped_configuration
+          ON scoped_configuration.scope_id = actor.id
+         AND scoped_configuration.key = 'connector_version'
+         AND scoped_configuration.scope_type = 'actor'
+        LEFT JOIN jobs
+          ON jobs.scope = connection.id::text
+         AND jobs.config_type = 'sync'
+         AND jobs.updated_at >= :cutoff_date
+        LEFT JOIN attempts
+          ON attempts.job_id = jobs.id
+        WHERE
+             connection.status != 'deprecated'
+        ORDER BY
+             connection.id,
+             attempts.ended_at DESC NULLS LAST
+    )
+    SELECT
+        pinned_version_id,
+        COUNT(*) AS total_connections,
+        COUNT(*) FILTER (WHERE connection_status = 'active') AS enabled_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status IS NOT NULL) AS active_connections,
+        COUNT(*) FILTER (WHERE pinned_version_id IS NOT NULL) AS pinned_connections,
+        COUNT(*) FILTER (WHERE pinned_version_id IS NULL) AS unpinned_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'succeeded') AS succeeded_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'failed') AS failed_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'cancelled') AS cancelled_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'running') AS running_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status IS NULL) AS unknown_connections
+    FROM latest_attempts
+    GROUP BY pinned_version_id
+    ORDER BY total_connections DESC
+    """
+)
+
+# Count connections by DESTINATION connector with latest attempt status breakdown
+SELECT_DESTINATION_CONNECTION_STATS = sqlalchemy.text(
+    """
+    WITH latest_attempts AS (
+        SELECT DISTINCT ON (connection.id)
+            connection.id AS connection_id,
+            connection.status AS connection_status,
+            scoped_configuration.value AS pinned_version_id,
+            attempts.status::text AS latest_attempt_status
+        FROM connection
+        JOIN actor
+          ON connection.destination_id = actor.id
+         AND actor.actor_definition_id = :connector_definition_id
+         AND actor.tombstone = false
+        JOIN workspace
+          ON actor.workspace_id = workspace.id
+         AND workspace.tombstone = false
+        LEFT JOIN scoped_configuration
+          ON scoped_configuration.scope_id = actor.id
+         AND scoped_configuration.key = 'connector_version'
+         AND scoped_configuration.scope_type = 'actor'
+        LEFT JOIN jobs
+          ON jobs.scope = connection.id::text
+         AND jobs.config_type = 'sync'
+         AND jobs.updated_at >= :cutoff_date
+        LEFT JOIN attempts
+          ON attempts.job_id = jobs.id
+        WHERE
+             connection.status != 'deprecated'
+        ORDER BY
+             connection.id,
+             attempts.ended_at DESC NULLS LAST
+    )
+    SELECT
+        pinned_version_id,
+        COUNT(*) AS total_connections,
+        COUNT(*) FILTER (WHERE connection_status = 'active') AS enabled_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status IS NOT NULL) AS active_connections,
+        COUNT(*) FILTER (WHERE pinned_version_id IS NOT NULL) AS pinned_connections,
+        COUNT(*) FILTER (WHERE pinned_version_id IS NULL) AS unpinned_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'succeeded') AS succeeded_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'failed') AS failed_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'cancelled') AS cancelled_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status = 'running') AS running_connections,
+        COUNT(*) FILTER (WHERE latest_attempt_status IS NULL) AS unknown_connections
+    FROM latest_attempts
+    GROUP BY pinned_version_id
+    ORDER BY total_connections DESC
+    """
+)

airbyte_ops_mcp/regression_tests/connector_runner.py

@@ -105,16 +105,19 @@ class ConnectorRunner:
         if self.config is not None:
             config_path = temp_dir / self.CONFIG_FILE
             config_path.write_text(json.dumps(self.config))
+            config_path.chmod(0o666)
             self.logger.debug(f"Wrote config to {config_path}")
 
         if self.configured_catalog is not None:
             catalog_path = temp_dir / self.CATALOG_FILE
             catalog_path.write_text(self.configured_catalog.json())
+            catalog_path.chmod(0o666)
             self.logger.debug(f"Wrote catalog to {catalog_path}")
 
         if self.state is not None:
             state_path = temp_dir / self.STATE_FILE
             state_path.write_text(json.dumps(self.state))
+            state_path.chmod(0o666)
             self.logger.debug(f"Wrote state to {state_path}")
 
     def _build_docker_command(self, temp_dir: Path) -> list[str]:
@@ -135,7 +138,7 @@ class ConnectorRunner:
             "--name",
             container_name,
             "-v",
-            f"{temp_dir}:{self.DATA_DIR}:ro",
+            f"{temp_dir}:{self.DATA_DIR}",
         ]
 
         if self.proxy_url:
@@ -168,9 +171,10 @@ class ConnectorRunner:
 
         with tempfile.TemporaryDirectory() as temp_dir:
             temp_path = Path(temp_dir)
-            # Make temp directory world-readable so non-root container users can access it
-            # Many connector images run as non-root users (e.g., 'airbyte' user)
-            temp_path.chmod(0o755)
+            # Make temp directory world-writable so non-root container users can read/write
+            # Many connector images run as non-root users (e.g., 'airbyte' user) with
+            # different UIDs than the host user, so they need write access for config migration
+            temp_path.chmod(0o777)
             self._prepare_data_directory(temp_path)
 
             docker_cmd = self._build_docker_command(temp_path)