airbyte-internal-ops 0.3.0__py3-none-any.whl → 0.4.0__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: airbyte-internal-ops
3
- Version: 0.3.0
3
+ Version: 0.4.0
4
4
  Summary: MCP and API interfaces that let the agents do the admin work
5
5
  Author-email: Aaron Steers <aj@airbyte.io>
6
6
  Keywords: admin,airbyte,api,mcp
@@ -1,7 +1,7 @@
1
1
  airbyte_ops_mcp/__init__.py,sha256=tuzdlMkfnWBnsri5KGHM2M_xuNnzFk2u_aR79mmN7Yg,772
2
2
  airbyte_ops_mcp/_annotations.py,sha256=MO-SBDnbykxxHDESG7d8rviZZ4WlZgJKv0a8eBqcEzQ,1757
3
3
  airbyte_ops_mcp/constants.py,sha256=khcv9W3WkApIyPygEGgE2noBIqLomjoOMLxFBU1ArjA,5308
4
- airbyte_ops_mcp/gcp_auth.py,sha256=5k-k145ZoYhHLjyDES8nrA8f8BBihRI0ykrdD1IcfOs,3599
4
+ airbyte_ops_mcp/gcp_auth.py,sha256=i0cm1_xX4fj_31iKlfARpNvTaSr85iGTSw9KMf4f4MU,7206
5
5
  airbyte_ops_mcp/github_actions.py,sha256=wKnuIVmF4u1gMYNdSoryD_PUmvMz5SaHgOvbU0dsolA,9957
6
6
  airbyte_ops_mcp/github_api.py,sha256=uupbYKAkm7yLHK_1cDXYKl1bOYhUygZhG5IHspS7duE,8104
7
7
  airbyte_ops_mcp/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
@@ -363,11 +363,11 @@ airbyte_ops_mcp/cloud_admin/auth.py,sha256=qE2Aqe0qbZB755KscL65s54Jz78-F-X5a8fXK
363
363
  airbyte_ops_mcp/cloud_admin/connection_config.py,sha256=9opGQer-cGMJANmm-LFLMwvMCNu3nzxa2n2XHkZj9Fw,4899
364
364
  airbyte_ops_mcp/cloud_admin/models.py,sha256=YZ3FbEW-tZa50khKTTl4Bzvy_LsGyyQd6qcpXo62jls,2670
365
365
  airbyte_ops_mcp/connection_config_retriever/__init__.py,sha256=Xoi-YvARrNPhECdpwEDDkdwEpnvj8zuUlwULpf4iRrU,800
366
- airbyte_ops_mcp/connection_config_retriever/audit_logging.py,sha256=GjT4dVa0TtvGDmiBz9qwzcYCnSf9hTo7UM6l7ubUNE8,2846
366
+ airbyte_ops_mcp/connection_config_retriever/audit_logging.py,sha256=QdOG9984NXeMaKeJnFUZ4oCOmqi37PBRG2NRBBjrZQQ,2753
367
367
  airbyte_ops_mcp/connection_config_retriever/retrieval.py,sha256=s6yeCyrboWkUd6KdaheEo87x-rLtQNTL8XeR8O9z2HI,12160
368
368
  airbyte_ops_mcp/connection_config_retriever/secrets_resolution.py,sha256=12g0lZzhCzAPl4Iv4eMW6d76mvXjIBGspOnNhywzks4,3644
369
369
  airbyte_ops_mcp/gcp_logs/__init__.py,sha256=IqkxclXJnD1U4L2at7aC9GYqPXnuLdYLgmkm3ZiIu6s,409
370
- airbyte_ops_mcp/gcp_logs/error_lookup.py,sha256=wtC2pXwUuJQcVyonIcduDyGxk8kjJ8Dj-Vyq9AdnYh4,12763
370
+ airbyte_ops_mcp/gcp_logs/error_lookup.py,sha256=Ufl1FtNQJKP_yWndVT1Xku1mT-gxW_0atmNMCYMXvOo,12757
371
371
  airbyte_ops_mcp/mcp/__init__.py,sha256=QqkNkxzdXlg-W03urBAQ3zmtOKFPf35rXgO9ceUjpng,334
372
372
  airbyte_ops_mcp/mcp/_guidance.py,sha256=48tQSnDnxqXtyGJxxgjz0ZiI814o_7Fj7f6R8jpQ7so,2375
373
373
  airbyte_ops_mcp/mcp/_http_headers.py,sha256=9TAH2RYhFR3z2JugW4Q3WrrqJIdaCzAbyA1GhtQ_EMM,7278
@@ -379,18 +379,18 @@ airbyte_ops_mcp/mcp/gcp_logs.py,sha256=IPtq4098_LN1Cgeba4jATO1iYFFFpL2-aRO0pGcOd
379
379
  airbyte_ops_mcp/mcp/github.py,sha256=h3M3VJrq09y_F9ueQVCq3bUbVBNFuTNKprHtGU_ttio,8045
380
380
  airbyte_ops_mcp/mcp/github_repo_ops.py,sha256=PiERpt8abo20Gz4CfXhrDNlVM4o4FOt5sweZJND2a0s,5314
381
381
  airbyte_ops_mcp/mcp/metadata.py,sha256=fwGW97WknR5lfKcQnFtK6dU87aA6TmLj1NkKyqDAV9g,270
382
- airbyte_ops_mcp/mcp/prerelease.py,sha256=nc6VU03ADVHWM3OjGKxbS5XqY4VoyRyrZNU_fyAtaOI,10465
383
- airbyte_ops_mcp/mcp/prod_db_queries.py,sha256=FfGoq3aEj6ZUT4ysBIs1w7LzzwBeRXTaRvPGEx62RzI,25474
382
+ airbyte_ops_mcp/mcp/prerelease.py,sha256=GI4p1rGDCLZ6QbEG57oD_M3_buIHwq9B0In6fbj7Ptk,11883
383
+ airbyte_ops_mcp/mcp/prod_db_queries.py,sha256=VsiBBnVbOjc8lBb2Xr1lmcH3wu7QHQfjd4lORarEE1s,42700
384
384
  airbyte_ops_mcp/mcp/prompts.py,sha256=mJld9mdPECXYZffWXGSvNs4Xevx3rxqUGNlzGKVC2_s,1599
385
385
  airbyte_ops_mcp/mcp/registry.py,sha256=PW-VYUj42qx2pQ_apUkVaoUFq7VgB9zEU7-aGrkSCCw,290
386
- airbyte_ops_mcp/mcp/regression_tests.py,sha256=S1h-5S5gcZA4WEtIZyAQ836hd04tjSRRqMiYMx0S93g,16079
386
+ airbyte_ops_mcp/mcp/regression_tests.py,sha256=VpXS36Ox2qPxtxnDhVoNfr83UfppWx8rMgCoDiKWzWg,16727
387
387
  airbyte_ops_mcp/mcp/server.py,sha256=lKAXxt4u4bz7dsKvAYFFHziMbun2pOnxYmrMtRxsZvM,5317
388
388
  airbyte_ops_mcp/mcp/server_info.py,sha256=Yi4B1auW64QZGBDas5mro_vwTjvrP785TFNSBP7GhRg,2361
389
389
  airbyte_ops_mcp/prod_db_access/__init__.py,sha256=5pxouMPY1beyWlB0UwPnbaLTKTHqU6X82rbbgKY2vYU,1069
390
390
  airbyte_ops_mcp/prod_db_access/db_engine.py,sha256=VUqEWZtharJUR-Cri_pMwtGh1C4Neu4s195mbEXlm-w,9190
391
391
  airbyte_ops_mcp/prod_db_access/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
392
- airbyte_ops_mcp/prod_db_access/queries.py,sha256=BBPAQEfcing4G0Q9PEmI8C_9kN26sZc65ZGXd9WuFSw,14257
393
- airbyte_ops_mcp/prod_db_access/sql.py,sha256=hTbPY4r_rrtJ34B5eVpwyuBMLotyuP--UTv0vl3ZwBw,19432
392
+ airbyte_ops_mcp/prod_db_access/queries.py,sha256=pyW5GxDZ5ibwXawxyI_IR7VFcmoX7pZyZ2jdADqhJRY,20276
393
+ airbyte_ops_mcp/prod_db_access/sql.py,sha256=lzFOYfkb-rFTaZ6vrAK9G8Ym4KTUdhPMBzK44NSRzcg,35362
394
394
  airbyte_ops_mcp/registry/__init__.py,sha256=iEaPlt9GrnlaLbc__98TguNeZG8wuQu7S-_2QkhHcbA,858
395
395
  airbyte_ops_mcp/registry/models.py,sha256=B4L4TKr52wo0xs0CqvCBrpowqjShzVnZ5eTr2-EyhNs,2346
396
396
  airbyte_ops_mcp/registry/publish.py,sha256=VoPxsM2_0zJ829orzCRN-kjgcJtuBNyXgW4I9J680ro,12717
@@ -399,8 +399,8 @@ airbyte_ops_mcp/regression_tests/cdk_secrets.py,sha256=iRjqqBS96KZoswfgT7ju-pE_p
399
399
  airbyte_ops_mcp/regression_tests/ci_output.py,sha256=rrvCVKKShc1iVPMuQJDBqSbsiAHIDpX8SA9j0Uwl_Cg,12718
400
400
  airbyte_ops_mcp/regression_tests/config.py,sha256=dwWeY0tatdbwl9BqbhZ7EljoZDCtKmGO5fvOAIxeXmA,5873
401
401
  airbyte_ops_mcp/regression_tests/connection_fetcher.py,sha256=5wIiA0VvCFNEc-fr6Po18gZMX3E5fyPOGf2SuVOqv5U,12799
402
- airbyte_ops_mcp/regression_tests/connection_secret_retriever.py,sha256=_qd_nBLx6Xc6yVQHht716sFELX8SgIE5q3R3R708tfw,4879
403
- airbyte_ops_mcp/regression_tests/connector_runner.py,sha256=bappfBSq8dn3IyVAMS_XuzYEwWus23hkDCHLa2RFysI,9920
402
+ airbyte_ops_mcp/regression_tests/connection_secret_retriever.py,sha256=FhWNVWq7sON4nwUmVJv8BgXBOqg1YV4b5WuWyCzZ0LU,4695
403
+ airbyte_ops_mcp/regression_tests/connector_runner.py,sha256=OZzUa2aLh0sHaEARsDePOA-e3qEX4cvh3Jhnvi8S1rY,10130
404
404
  airbyte_ops_mcp/regression_tests/evaluation_modes.py,sha256=lAL6pEDmy_XCC7_m4_NXjt_f6Z8CXeAhMkc0FU8bm_M,1364
405
405
  airbyte_ops_mcp/regression_tests/http_metrics.py,sha256=oTD7f2MnQOvx4plOxHop2bInQ0-whvuToSsrC7TIM-M,12469
406
406
  airbyte_ops_mcp/regression_tests/models.py,sha256=brtAT9oO1TwjFcP91dFcu0XcUNqQb-jf7di1zkoVEuo,8782
@@ -414,7 +414,7 @@ airbyte_ops_mcp/regression_tests/regression/comparators.py,sha256=MJkLZEKHivgrG0
414
414
  airbyte_ops_mcp/regression_tests/validation/__init__.py,sha256=MBEwGOoNuqT4_oCahtoK62OKWIjUCfWa7vZTxNj_0Ek,1532
415
415
  airbyte_ops_mcp/regression_tests/validation/catalog_validators.py,sha256=jqqVAMOk0mtdPgwu4d0hA0ZEjtsNh5gapvGydRv3_qk,12553
416
416
  airbyte_ops_mcp/regression_tests/validation/record_validators.py,sha256=RjauAhKWNwxMBTu0eNS2hMFNQVs5CLbQU51kp6FOVDk,7432
417
- airbyte_internal_ops-0.3.0.dist-info/METADATA,sha256=Gx40HXaZtFle9mxFDJQNYMGccjrZ3d0xirHsaWcg04s,5679
418
- airbyte_internal_ops-0.3.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
419
- airbyte_internal_ops-0.3.0.dist-info/entry_points.txt,sha256=WxP0l7bRFss4Cr5uQqVj9mTEKwnRKouNuphXQF0lotA,171
420
- airbyte_internal_ops-0.3.0.dist-info/RECORD,,
417
+ airbyte_internal_ops-0.4.0.dist-info/METADATA,sha256=K9rJIUSobD2QWdHccHpKZooawgipH8ZozDqTl0FrG-8,5679
418
+ airbyte_internal_ops-0.4.0.dist-info/WHEEL,sha256=WLgqFyCfm_KASv4WHyYy0P3pM_m7J5L9k2skdKLirC8,87
419
+ airbyte_internal_ops-0.4.0.dist-info/entry_points.txt,sha256=WxP0l7bRFss4Cr5uQqVj9mTEKwnRKouNuphXQF0lotA,171
420
+ airbyte_internal_ops-0.4.0.dist-info/RECORD,,
@@ -11,9 +11,8 @@ import logging
11
11
  import subprocess
12
12
  from typing import TYPE_CHECKING, Any, Callable
13
13
 
14
- from google.cloud import logging as gcloud_logging
15
-
16
14
  from airbyte_ops_mcp.constants import GCP_PROJECT_NAME
15
+ from airbyte_ops_mcp.gcp_auth import get_logging_client
17
16
 
18
17
  if TYPE_CHECKING:
19
18
  from airbyte_ops_mcp.connection_config_retriever.retrieval import (
@@ -23,21 +22,18 @@ if TYPE_CHECKING:
23
22
  LOGGER = logging.getLogger(__name__)
24
23
 
25
24
  # Lazy-initialized to avoid import-time GCP calls
26
- _logging_client: gcloud_logging.Client | None = None
27
25
  _airbyte_gcloud_logger: Any = None
28
26
 
29
27
 
30
28
  def _get_logger() -> Any:
31
29
  """Get the GCP Cloud Logger, initializing lazily on first use."""
32
- global _logging_client, _airbyte_gcloud_logger
30
+ global _airbyte_gcloud_logger
33
31
 
34
32
  if _airbyte_gcloud_logger is not None:
35
33
  return _airbyte_gcloud_logger
36
34
 
37
- _logging_client = gcloud_logging.Client(project=GCP_PROJECT_NAME)
38
- _airbyte_gcloud_logger = _logging_client.logger(
39
- "airbyte-cloud-connection-retriever"
40
- )
35
+ logging_client = get_logging_client(GCP_PROJECT_NAME)
36
+ _airbyte_gcloud_logger = logging_client.logger("airbyte-cloud-connection-retriever")
41
37
  return _airbyte_gcloud_logger
42
38
 
43
39
 
@@ -6,94 +6,187 @@ the airbyte-ops-mcp codebase. It supports both standard Application Default
6
6
  Credentials (ADC) and the GCP_PROD_DB_ACCESS_CREDENTIALS environment variable
7
7
  used internally at Airbyte.
8
8
 
9
+ The preferred approach is to pass credentials directly to GCP client constructors
10
+ rather than relying on file-based ADC discovery. This module provides helpers
11
+ that construct credentials from JSON content in environment variables.
12
+
9
13
  Usage:
10
- from airbyte_ops_mcp.gcp_auth import get_secret_manager_client
14
+ from airbyte_ops_mcp.gcp_auth import get_gcp_credentials, get_secret_manager_client
15
+
16
+ # Get credentials object to pass to any GCP client
17
+ credentials = get_gcp_credentials()
18
+ client = logging.Client(project="my-project", credentials=credentials)
11
19
 
12
- # Get a properly authenticated Secret Manager client
20
+ # Or use the convenience helper for Secret Manager
13
21
  client = get_secret_manager_client()
14
22
  """
15
23
 
16
24
  from __future__ import annotations
17
25
 
26
+ import json
18
27
  import logging
19
28
  import os
20
- import tempfile
21
- from pathlib import Path
29
+ import sys
30
+ import threading
22
31
 
32
+ import google.auth
33
+ from google.cloud import logging as gcp_logging
23
34
  from google.cloud import secretmanager
35
+ from google.oauth2 import service_account
24
36
 
25
37
  from airbyte_ops_mcp.constants import ENV_GCP_PROD_DB_ACCESS_CREDENTIALS
26
38
 
27
39
  logger = logging.getLogger(__name__)
28
40
 
29
- # Environment variable name (internal to GCP libraries)
30
- ENV_GOOGLE_APPLICATION_CREDENTIALS = "GOOGLE_APPLICATION_CREDENTIALS"
31
41
 
32
- # Module-level cache for the credentials file path
33
- _credentials_file_path: str | None = None
42
+ def _get_identity_from_service_account_info(info: dict) -> str | None:
43
+ """Extract service account identity from parsed JSON info.
44
+
45
+ Only accesses the 'client_email' key to avoid any risk of leaking
46
+ other credential material.
34
47
 
48
+ Args:
49
+ info: Parsed service account JSON as a dict.
35
50
 
36
- def ensure_adc_credentials() -> str | None:
37
- """Ensure GCP Application Default Credentials are available.
51
+ Returns:
52
+ The client_email if present and a string, otherwise None.
53
+ """
54
+ client_email = info.get("client_email")
55
+ if isinstance(client_email, str):
56
+ return client_email
57
+ return None
38
58
 
39
- If GOOGLE_APPLICATION_CREDENTIALS is not set but GCP_PROD_DB_ACCESS_CREDENTIALS is,
40
- write the JSON credentials to a temp file and set GOOGLE_APPLICATION_CREDENTIALS
41
- to point to that file. This provides a fallback for internal employees who use
42
- GCP_PROD_DB_ACCESS_CREDENTIALS as their standard credential source.
43
59
 
44
- Note: GOOGLE_APPLICATION_CREDENTIALS must be a file path, not JSON content.
45
- The GCP_PROD_DB_ACCESS_CREDENTIALS env var contains the JSON content directly,
46
- so we write it to a temp file first.
60
+ def _get_identity_from_credentials(
61
+ credentials: google.auth.credentials.Credentials,
62
+ ) -> str | None:
63
+ """Extract identity from a credentials object using safe attribute access.
47
64
 
48
- This function is idempotent and safe to call multiple times.
65
+ Only accesses known-safe attributes that don't trigger network calls
66
+ or token refresh.
67
+
68
+ Args:
69
+ credentials: A GCP credentials object.
49
70
 
50
71
  Returns:
51
- The path to the credentials file if one was created, or None if
52
- GOOGLE_APPLICATION_CREDENTIALS was already set.
72
+ The service account email if available, otherwise None.
53
73
  """
54
- global _credentials_file_path
74
+ # Try service_account_email first (most common for service accounts)
75
+ identity = getattr(credentials, "service_account_email", None)
76
+ if isinstance(identity, str):
77
+ return identity
78
+
79
+ # Try signer_email as fallback (sometimes present on impersonated creds)
80
+ identity = getattr(credentials, "signer_email", None)
81
+ if isinstance(identity, str):
82
+ return identity
83
+
84
+ return None
85
+
55
86
 
56
- # If GOOGLE_APPLICATION_CREDENTIALS is already set, nothing to do
57
- if ENV_GOOGLE_APPLICATION_CREDENTIALS in os.environ:
58
- return None
87
+ # Default scopes for GCP services used by this module
88
+ DEFAULT_GCP_SCOPES = ["https://www.googleapis.com/auth/cloud-platform"]
59
89
 
60
- # Check if we have the fallback credentials
61
- gsm_creds = os.getenv(ENV_GCP_PROD_DB_ACCESS_CREDENTIALS)
62
- if not gsm_creds:
63
- return None
90
+ # Module-level cache for credentials (thread-safe)
91
+ _cached_credentials: google.auth.credentials.Credentials | None = None
92
+ _credentials_lock = threading.Lock()
64
93
 
65
- # Reuse the same file path if we've already written credentials and file still exists
66
- if _credentials_file_path is not None and Path(_credentials_file_path).exists():
67
- os.environ[ENV_GOOGLE_APPLICATION_CREDENTIALS] = _credentials_file_path
68
- return _credentials_file_path
69
94
 
70
- # Write credentials to a temp file
71
- # Use a unique filename based on PID to avoid collisions between processes
72
- creds_file = Path(tempfile.gettempdir()) / f"gcp_prod_db_creds_{os.getpid()}.json"
73
- creds_file.write_text(gsm_creds)
95
+ def get_gcp_credentials() -> google.auth.credentials.Credentials:
96
+ """Get GCP credentials, preferring direct JSON parsing over file-based ADC.
74
97
 
75
- # Set restrictive permissions (owner read/write only)
76
- creds_file.chmod(0o600)
98
+ This function resolves credentials in the following order:
99
+ 1. GCP_PROD_DB_ACCESS_CREDENTIALS env var (JSON content) - parsed directly
100
+ 2. Standard ADC discovery (workload identity, gcloud auth, GOOGLE_APPLICATION_CREDENTIALS)
77
101
 
78
- _credentials_file_path = str(creds_file)
79
- os.environ[ENV_GOOGLE_APPLICATION_CREDENTIALS] = _credentials_file_path
102
+ The credentials are cached after first resolution for efficiency.
103
+ Uses the cloud-platform scope which provides access to all GCP services.
80
104
 
81
- logger.debug(
82
- f"Wrote {ENV_GCP_PROD_DB_ACCESS_CREDENTIALS} to {creds_file} and set "
83
- f"{ENV_GOOGLE_APPLICATION_CREDENTIALS}"
84
- )
105
+ Returns:
106
+ A Credentials object that can be passed to any GCP client constructor.
85
107
 
86
- return _credentials_file_path
108
+ Raises:
109
+ google.auth.exceptions.DefaultCredentialsError: If no credentials can be found.
110
+ """
111
+ global _cached_credentials
112
+
113
+ # Return cached credentials if available (fast path without lock)
114
+ if _cached_credentials is not None:
115
+ return _cached_credentials
116
+
117
+ # Acquire lock for thread-safe credential initialization
118
+ with _credentials_lock:
119
+ # Double-check after acquiring lock (another thread may have initialized)
120
+ if _cached_credentials is not None:
121
+ return _cached_credentials
122
+
123
+ # Try GCP_PROD_DB_ACCESS_CREDENTIALS first (JSON content in env var)
124
+ creds_json = os.getenv(ENV_GCP_PROD_DB_ACCESS_CREDENTIALS)
125
+ if creds_json:
126
+ try:
127
+ creds_dict = json.loads(creds_json)
128
+ credentials = service_account.Credentials.from_service_account_info(
129
+ creds_dict,
130
+ scopes=DEFAULT_GCP_SCOPES,
131
+ )
132
+ # Extract identity safely (only after successful credential creation)
133
+ identity = _get_identity_from_service_account_info(creds_dict)
134
+ identity_str = f" (identity: {identity})" if identity else ""
135
+ print(
136
+ f"GCP credentials loaded from {ENV_GCP_PROD_DB_ACCESS_CREDENTIALS}{identity_str}",
137
+ file=sys.stderr,
138
+ )
139
+ logger.debug(
140
+ f"Loaded GCP credentials from {ENV_GCP_PROD_DB_ACCESS_CREDENTIALS} env var"
141
+ )
142
+ _cached_credentials = credentials
143
+ return credentials
144
+ except (json.JSONDecodeError, ValueError) as e:
145
+ # Log only exception type to avoid any risk of leaking credential content
146
+ logger.warning(
147
+ f"Failed to parse {ENV_GCP_PROD_DB_ACCESS_CREDENTIALS}: "
148
+ f"{type(e).__name__}. Falling back to ADC discovery."
149
+ )
150
+
151
+ # Fall back to standard ADC discovery
152
+ credentials, project = google.auth.default(scopes=DEFAULT_GCP_SCOPES)
153
+ # Extract identity safely from ADC credentials
154
+ identity = _get_identity_from_credentials(credentials)
155
+ identity_str = f" (identity: {identity})" if identity else ""
156
+ project_str = f" (project: {project})" if project else ""
157
+ print(
158
+ f"GCP credentials loaded via ADC{project_str}{identity_str}",
159
+ file=sys.stderr,
160
+ )
161
+ logger.debug(f"Loaded GCP credentials via ADC discovery (project: {project})")
162
+ _cached_credentials = credentials
163
+ return credentials
87
164
 
88
165
 
89
166
  def get_secret_manager_client() -> secretmanager.SecretManagerServiceClient:
90
167
  """Get a Secret Manager client with proper credential handling.
91
168
 
92
- This function ensures GCP credentials are available (supporting the
93
- GCP_PROD_DB_ACCESS_CREDENTIALS fallback) before creating the client.
169
+ This function uses get_gcp_credentials() to resolve credentials and passes
170
+ them directly to the client constructor.
94
171
 
95
172
  Returns:
96
173
  A configured SecretManagerServiceClient instance.
97
174
  """
98
- ensure_adc_credentials()
99
- return secretmanager.SecretManagerServiceClient()
175
+ credentials = get_gcp_credentials()
176
+ return secretmanager.SecretManagerServiceClient(credentials=credentials)
177
+
178
+
179
+ def get_logging_client(project: str) -> gcp_logging.Client:
180
+ """Get a Cloud Logging client with proper credential handling.
181
+
182
+ This function uses get_gcp_credentials() to resolve credentials and passes
183
+ them directly to the client constructor.
184
+
185
+ Args:
186
+ project: The GCP project ID to use for logging operations.
187
+
188
+ Returns:
189
+ A configured Cloud Logging Client instance.
190
+ """
191
+ credentials = get_gcp_credentials()
192
+ return gcp_logging.Client(project=project, credentials=credentials)
@@ -24,10 +24,12 @@ from datetime import UTC, datetime, timedelta
24
24
  from enum import StrEnum
25
25
  from typing import Any
26
26
 
27
- from google.cloud import logging
27
+ from google.cloud import logging as gcp_logging
28
28
  from google.cloud.logging_v2 import entries
29
29
  from pydantic import BaseModel, Field
30
30
 
31
+ from airbyte_ops_mcp.gcp_auth import get_logging_client
32
+
31
33
  # Default GCP project for Airbyte Cloud
32
34
  DEFAULT_GCP_PROJECT = "prod-ab-cloud-proj"
33
35
 
@@ -291,14 +293,13 @@ def fetch_error_logs(
291
293
  specified error ID, then fetches related log entries (multi-line stack traces)
292
294
  from the same timestamp and resource.
293
295
  """
294
- client_options = {"quota_project_id": project}
295
- client = logging.Client(project=project, client_options=client_options)
296
+ client = get_logging_client(project)
296
297
 
297
298
  filter_str = _build_filter(error_id, lookback_days, min_severity_filter)
298
299
 
299
300
  entries_iterator = client.list_entries(
300
301
  filter_=filter_str,
301
- order_by=logging.DESCENDING,
302
+ order_by=gcp_logging.DESCENDING,
302
303
  )
303
304
 
304
305
  initial_matches = list(entries_iterator)
@@ -356,7 +357,7 @@ def fetch_error_logs(
356
357
 
357
358
  related_entries = client.list_entries(
358
359
  filter_=related_filter,
359
- order_by=logging.ASCENDING,
360
+ order_by=gcp_logging.ASCENDING,
360
361
  )
361
362
 
362
363
  for entry in related_entries:
@@ -2,12 +2,15 @@
2
2
  """MCP tools for triggering connector pre-release workflows.
3
3
 
4
4
  This module provides MCP tools for triggering the publish-connectors-prerelease
5
- workflow in the airbytehq/airbyte repository via GitHub's workflow dispatch API.
5
+ workflow in the airbytehq/airbyte repository (for OSS connectors) or the
6
+ publish_enterprise_connectors workflow in airbytehq/airbyte-enterprise
7
+ (for enterprise connectors) via GitHub's workflow dispatch API.
6
8
  """
7
9
 
8
10
  from __future__ import annotations
9
11
 
10
12
  import base64
13
+ from enum import StrEnum
11
14
  from typing import Annotated, Literal
12
15
 
13
16
  import requests
@@ -18,12 +21,25 @@ from pydantic import BaseModel, Field
18
21
  from airbyte_ops_mcp.github_actions import GITHUB_API_BASE, resolve_github_token
19
22
  from airbyte_ops_mcp.mcp._mcp_utils import mcp_tool, register_mcp_tools
20
23
 
24
+
25
+ class ConnectorRepo(StrEnum):
26
+ """Repository where connector code is located."""
27
+
28
+ AIRBYTE = "airbyte"
29
+ AIRBYTE_ENTERPRISE = "airbyte-enterprise"
30
+
31
+
21
32
  DEFAULT_REPO_OWNER = "airbytehq"
22
- DEFAULT_REPO_NAME = "airbyte"
33
+ DEFAULT_REPO_NAME = ConnectorRepo.AIRBYTE
23
34
  DEFAULT_BRANCH = "master"
24
35
  PRERELEASE_WORKFLOW_FILE = "publish-connectors-prerelease-command.yml"
25
36
  CONNECTOR_PATH_PREFIX = "airbyte-integrations/connectors"
26
37
 
38
+ # Enterprise repository constants
39
+ ENTERPRISE_REPO_NAME = ConnectorRepo.AIRBYTE_ENTERPRISE
40
+ ENTERPRISE_DEFAULT_BRANCH = "main"
41
+ ENTERPRISE_PRERELEASE_WORKFLOW_FILE = "publish_enterprise_connectors.yml"
42
+
27
43
  # Token env vars for prerelease publishing (in order of preference)
28
44
  PRERELEASE_TOKEN_ENV_VARS = [
29
45
  "GITHUB_CONNECTOR_PUBLISHING_PAT",
@@ -238,6 +254,14 @@ def publish_connector_to_airbyte_registry(
238
254
  int,
239
255
  Field(description="The pull request number containing the connector changes"),
240
256
  ],
257
+ repo: Annotated[
258
+ ConnectorRepo,
259
+ Field(
260
+ default=ConnectorRepo.AIRBYTE,
261
+ description="Repository where the connector PR is located. "
262
+ "Use 'airbyte' for OSS connectors (default) or 'airbyte-enterprise' for enterprise connectors.",
263
+ ),
264
+ ],
241
265
  prerelease: Annotated[
242
266
  Literal[True],
243
267
  Field(
@@ -249,8 +273,10 @@ def publish_connector_to_airbyte_registry(
249
273
  """Publish a connector to the Airbyte registry.
250
274
 
251
275
  Currently only supports pre-release publishing. This tool triggers the
252
- publish-connectors-prerelease workflow in the airbytehq/airbyte repository,
253
- which publishes a pre-release version of the specified connector from the PR branch.
276
+ publish-connectors-prerelease workflow in the airbytehq/airbyte repository
277
+ (for OSS connectors) or the publish_enterprise_connectors workflow in
278
+ airbytehq/airbyte-enterprise (for enterprise connectors), which publishes
279
+ a pre-release version of the specified connector from the PR branch.
254
280
 
255
281
  Pre-release versions are tagged with the format: {version}-preview.{7-char-git-sha}
256
282
  These versions are available for version pinning via the scoped_configuration API.
@@ -267,17 +293,27 @@ def publish_connector_to_airbyte_registry(
267
293
  # Guard: Check for required token
268
294
  token = resolve_github_token(PRERELEASE_TOKEN_ENV_VARS)
269
295
 
296
+ # Determine repo-specific settings
297
+ is_enterprise = repo == ConnectorRepo.AIRBYTE_ENTERPRISE
298
+ target_repo_name = ENTERPRISE_REPO_NAME if is_enterprise else DEFAULT_REPO_NAME
299
+ target_branch = ENTERPRISE_DEFAULT_BRANCH if is_enterprise else DEFAULT_BRANCH
300
+ target_workflow = (
301
+ ENTERPRISE_PRERELEASE_WORKFLOW_FILE
302
+ if is_enterprise
303
+ else PRERELEASE_WORKFLOW_FILE
304
+ )
305
+
270
306
  # Get the PR's head SHA for computing the docker image tag
271
307
  # Note: We no longer pass gitref to the workflow - it derives the ref from PR number
272
308
  head_info = _get_pr_head_info(
273
- DEFAULT_REPO_OWNER, DEFAULT_REPO_NAME, pr_number, token
309
+ DEFAULT_REPO_OWNER, target_repo_name, pr_number, token
274
310
  )
275
311
 
276
312
  # Prepare workflow inputs
277
313
  # The workflow uses refs/pull/{pr}/head directly - no gitref needed
278
314
  # Note: The workflow auto-detects modified connectors from the PR
279
315
  workflow_inputs = {
280
- "repo": f"{DEFAULT_REPO_OWNER}/{DEFAULT_REPO_NAME}",
316
+ "repo": f"{DEFAULT_REPO_OWNER}/{target_repo_name}",
281
317
  "pr": str(pr_number),
282
318
  }
283
319
 
@@ -285,9 +321,9 @@ def publish_connector_to_airbyte_registry(
285
321
  # The workflow will checkout the PR branch via inputs.gitref
286
322
  workflow_url = _trigger_workflow_dispatch(
287
323
  owner=DEFAULT_REPO_OWNER,
288
- repo=DEFAULT_REPO_NAME,
289
- workflow_file=PRERELEASE_WORKFLOW_FILE,
290
- ref=DEFAULT_BRANCH,
324
+ repo=target_repo_name,
325
+ workflow_file=target_workflow,
326
+ ref=target_branch,
291
327
  inputs=workflow_inputs,
292
328
  token=token,
293
329
  )
@@ -297,7 +333,7 @@ def publish_connector_to_airbyte_registry(
297
333
  docker_image_tag: str | None = None
298
334
  metadata = _get_connector_metadata(
299
335
  DEFAULT_REPO_OWNER,
300
- DEFAULT_REPO_NAME,
336
+ target_repo_name,
301
337
  connector_name,
302
338
  head_info.sha,
303
339
  token,
@@ -311,9 +347,10 @@ def publish_connector_to_airbyte_registry(
311
347
  base_version, head_info.sha
312
348
  )
313
349
 
350
+ repo_info = f" from {repo}" if is_enterprise else ""
314
351
  return PrereleaseWorkflowResult(
315
352
  success=True,
316
- message=f"Successfully triggered pre-release workflow for {connector_name} from PR #{pr_number}",
353
+ message=f"Successfully triggered pre-release workflow for {connector_name}{repo_info} from PR #{pr_number}",
317
354
  workflow_url=workflow_url,
318
355
  connector_name=connector_name,
319
356
  pr_number=pr_number,