airbyte-internal-ops 0.1.5__py3-none-any.whl → 0.1.6__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.
Files changed (36) hide show
  1. {airbyte_internal_ops-0.1.5.dist-info → airbyte_internal_ops-0.1.6.dist-info}/METADATA +70 -1
  2. {airbyte_internal_ops-0.1.5.dist-info → airbyte_internal_ops-0.1.6.dist-info}/RECORD +25 -26
  3. airbyte_ops_mcp/__init__.py +30 -2
  4. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/connectors/pipeline.py +2 -8
  5. airbyte_ops_mcp/airbyte_repo/list_connectors.py +132 -0
  6. airbyte_ops_mcp/cli/registry.py +90 -1
  7. airbyte_ops_mcp/connection_config_retriever/__init__.py +26 -0
  8. airbyte_ops_mcp/{live_tests/_connection_retriever → connection_config_retriever}/audit_logging.py +5 -6
  9. airbyte_ops_mcp/{live_tests/_connection_retriever → connection_config_retriever}/retrieval.py +8 -22
  10. airbyte_ops_mcp/{live_tests/_connection_retriever → connection_config_retriever}/secrets_resolution.py +8 -42
  11. airbyte_ops_mcp/constants.py +35 -0
  12. airbyte_ops_mcp/live_tests/connection_secret_retriever.py +1 -1
  13. airbyte_ops_mcp/mcp/github_repo_ops.py +10 -0
  14. airbyte_ops_mcp/mcp/prod_db_queries.py +357 -0
  15. airbyte_ops_mcp/mcp/server.py +2 -0
  16. airbyte_ops_mcp/prod_db_access/__init__.py +34 -0
  17. airbyte_ops_mcp/prod_db_access/db_engine.py +127 -0
  18. airbyte_ops_mcp/prod_db_access/py.typed +0 -0
  19. airbyte_ops_mcp/prod_db_access/queries.py +272 -0
  20. airbyte_ops_mcp/prod_db_access/sql.py +353 -0
  21. airbyte_ops_mcp/registry/__init__.py +34 -0
  22. airbyte_ops_mcp/registry/models.py +63 -0
  23. airbyte_ops_mcp/registry/publish.py +368 -0
  24. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/connectors/publish/__init__.py +0 -3
  25. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/connectors/publish/commands.py +0 -242
  26. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/connectors/publish/context.py +0 -175
  27. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/connectors/publish/pipeline.py +0 -1056
  28. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/poetry/publish/__init__.py +0 -3
  29. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/poetry/publish/commands.py +0 -127
  30. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/airbyte_ci/steps/python_registry.py +0 -238
  31. airbyte_ops_mcp/_legacy/airbyte_ci/connector_pipelines/models/contexts/python_registry_publish.py +0 -119
  32. airbyte_ops_mcp/live_tests/_connection_retriever/__init__.py +0 -35
  33. airbyte_ops_mcp/live_tests/_connection_retriever/consts.py +0 -33
  34. airbyte_ops_mcp/live_tests/_connection_retriever/db_access.py +0 -82
  35. {airbyte_internal_ops-0.1.5.dist-info → airbyte_internal_ops-0.1.6.dist-info}/WHEEL +0 -0
  36. {airbyte_internal_ops-0.1.5.dist-info → airbyte_internal_ops-0.1.6.dist-info}/entry_points.txt +0 -0
@@ -1,3 +0,0 @@
1
- #
2
- # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
3
- #
@@ -1,127 +0,0 @@
1
- #
2
- # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
3
- #
4
-
5
- """
6
- Module exposing the format commands.
7
- """
8
-
9
- from __future__ import annotations
10
-
11
- from typing import Optional
12
-
13
- import asyncclick as click
14
- from consts import (
15
- DEFAULT_PYTHON_PACKAGE_REGISTRY_CHECK_URL,
16
- DEFAULT_PYTHON_PACKAGE_REGISTRY_URL,
17
- )
18
- from packaging import version
19
- from pipelines.airbyte_ci.steps.python_registry import PublishToPythonRegistry
20
- from pipelines.cli.confirm_prompt import confirm
21
- from pipelines.cli.dagger_pipeline_command import DaggerPipelineCommand
22
- from pipelines.cli.secrets import wrap_in_secret
23
- from pipelines.models.contexts.click_pipeline_context import (
24
- ClickPipelineContext,
25
- pass_pipeline_context,
26
- )
27
- from pipelines.models.contexts.python_registry_publish import (
28
- PythonRegistryPublishContext,
29
- )
30
- from pipelines.models.secrets import Secret
31
- from pipelines.models.steps import StepStatus
32
-
33
-
34
- async def _has_metadata_yaml(context: PythonRegistryPublishContext) -> bool:
35
- dir_to_publish = context.get_repo_dir(context.package_path)
36
- return "metadata.yaml" in await dir_to_publish.entries()
37
-
38
-
39
- def _validate_python_version(
40
- _ctx: dict, _param: dict, value: Optional[str]
41
- ) -> Optional[str]:
42
- """
43
- Check if an given version is valid.
44
- """
45
- if value is None:
46
- return value
47
- try:
48
- version.Version(value)
49
- return value
50
- except version.InvalidVersion:
51
- raise click.BadParameter(f"Version {value} is not a valid version.")
52
-
53
-
54
- @click.command(
55
- cls=DaggerPipelineCommand,
56
- name="publish",
57
- help="Publish a Python package to a registry.",
58
- )
59
- @click.option(
60
- "--python-registry-token",
61
- help="Access token",
62
- type=click.STRING,
63
- required=True,
64
- envvar="PYTHON_REGISTRY_TOKEN",
65
- callback=wrap_in_secret,
66
- )
67
- @click.option(
68
- "--python-registry-url",
69
- help="Which registry to publish to. If not set, the default pypi is used. For test pypi, use https://test.pypi.org/legacy/",
70
- type=click.STRING,
71
- default=DEFAULT_PYTHON_PACKAGE_REGISTRY_URL,
72
- envvar="PYTHON_REGISTRY_URL",
73
- )
74
- @click.option(
75
- "--publish-name",
76
- help="The name of the package to publish. If not set, the name will be inferred from the pyproject.toml file of the package.",
77
- type=click.STRING,
78
- )
79
- @click.option(
80
- "--publish-version",
81
- help="The version of the package to publish. If not set, the version will be inferred from the pyproject.toml file of the package.",
82
- type=click.STRING,
83
- callback=_validate_python_version,
84
- )
85
- @pass_pipeline_context
86
- @click.pass_context
87
- async def publish(
88
- ctx: click.Context,
89
- click_pipeline_context: ClickPipelineContext,
90
- python_registry_token: Secret,
91
- python_registry_url: str,
92
- publish_name: Optional[str],
93
- publish_version: Optional[str],
94
- ) -> bool:
95
- context = PythonRegistryPublishContext(
96
- is_local=ctx.obj["is_local"],
97
- git_branch=ctx.obj["git_branch"],
98
- git_revision=ctx.obj["git_revision"],
99
- diffed_branch=ctx.obj["diffed_branch"],
100
- git_repo_url=ctx.obj["git_repo_url"],
101
- ci_report_bucket=ctx.obj["ci_report_bucket_name"],
102
- report_output_prefix=ctx.obj["report_output_prefix"],
103
- gha_workflow_run_url=ctx.obj.get("gha_workflow_run_url"),
104
- dagger_logs_url=ctx.obj.get("dagger_logs_url"),
105
- pipeline_start_timestamp=ctx.obj.get("pipeline_start_timestamp"),
106
- ci_context=ctx.obj.get("ci_context"),
107
- ci_gcp_credentials=ctx.obj["ci_gcp_credentials"],
108
- python_registry_token=python_registry_token,
109
- registry=python_registry_url,
110
- registry_check_url=DEFAULT_PYTHON_PACKAGE_REGISTRY_CHECK_URL,
111
- package_path=ctx.obj["package_path"],
112
- package_name=publish_name,
113
- version=publish_version,
114
- )
115
-
116
- dagger_client = await click_pipeline_context.get_dagger_client()
117
- context.dagger_client = dagger_client
118
-
119
- if await _has_metadata_yaml(context):
120
- confirm(
121
- "It looks like you are trying to publish a connector. In most cases, the `connectors` command group should be used instead. Do you want to continue?",
122
- abort=True,
123
- )
124
-
125
- publish_result = await PublishToPythonRegistry(context).run()
126
-
127
- return publish_result.status is StepStatus.SUCCESS
@@ -1,238 +0,0 @@
1
- #
2
- # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
3
- #
4
-
5
- import configparser
6
- import io
7
- import uuid
8
- from enum import Enum, auto
9
- from typing import Dict, Optional
10
-
11
- import tomli_w
12
- import tomllib
13
- from consts import PYPROJECT_TOML_FILE_PATH, SETUP_PY_FILE_PATH
14
- from dagger import Container, Directory
15
- from pipelines.dagger.actions.python.poetry import with_poetry
16
- from pipelines.helpers.utils import sh_dash_c
17
- from pipelines.models.contexts.python_registry_publish import (
18
- PythonPackageMetadata,
19
- PythonRegistryPublishContext,
20
- )
21
- from pipelines.models.steps import Step, StepResult
22
-
23
-
24
- class PackageType(Enum):
25
- POETRY = auto()
26
- PIP = auto()
27
-
28
-
29
- class PublishToPythonRegistry(Step):
30
- context: PythonRegistryPublishContext
31
- title = "Publish package to python registry"
32
- max_retries = 3
33
-
34
- def _get_base_container(self) -> Container:
35
- return with_poetry(self.context)
36
-
37
- async def _get_package_metadata_from_pyproject_toml(
38
- self, package_dir_to_publish: Directory
39
- ) -> Optional[PythonPackageMetadata]:
40
- pyproject_toml = package_dir_to_publish.file(PYPROJECT_TOML_FILE_PATH)
41
- pyproject_toml_content = await pyproject_toml.contents()
42
- contents = tomllib.loads(pyproject_toml_content)
43
- try:
44
- return PythonPackageMetadata(
45
- contents["tool"]["poetry"]["name"],
46
- contents["tool"]["poetry"]["version"],
47
- )
48
- except KeyError:
49
- return None
50
-
51
- async def _get_package_type(
52
- self, package_dir_to_publish: Directory
53
- ) -> Optional[PackageType]:
54
- files = await package_dir_to_publish.entries()
55
- has_pyproject_toml = PYPROJECT_TOML_FILE_PATH in files
56
- has_setup_py = SETUP_PY_FILE_PATH in files
57
- if has_pyproject_toml:
58
- return PackageType.POETRY
59
- elif has_setup_py:
60
- return PackageType.PIP
61
- else:
62
- return None
63
-
64
- async def _run(self) -> StepResult:
65
- package_dir_to_publish = await self.context.get_repo_dir(
66
- self.context.package_path
67
- )
68
- package_type = await self._get_package_type(package_dir_to_publish)
69
-
70
- if not package_type:
71
- return self.skip(
72
- "Connector does not have a pyproject.toml file or setup.py file, skipping."
73
- )
74
-
75
- result = await self._ensure_package_name_and_version(
76
- package_dir_to_publish, package_type
77
- )
78
- if result:
79
- return result
80
-
81
- self.logger.info(
82
- f"Uploading package {self.context.package_metadata.name} version {self.context.package_metadata.version} to {self.context.registry}..."
83
- )
84
-
85
- return await self._publish(package_dir_to_publish, package_type)
86
-
87
- async def _ensure_package_name_and_version(
88
- self, package_dir_to_publish: Directory, package_type: PackageType
89
- ) -> Optional[StepResult]:
90
- """
91
- Try to infer package name and version from the pyproject.toml file. If it is not present, we need to have the package name and version set.
92
- Setup.py packages need to set package name and version as parameter.
93
-
94
- Returns None if package name and version are set, otherwise a StepResult with a skip message.
95
- """
96
- if self.context.package_metadata.name and self.context.package_metadata.version:
97
- return None
98
-
99
- if package_type is not PackageType.POETRY:
100
- return self.skip(
101
- "Connector does not have a pyproject.toml file and version and package name is not set otherwise, skipping."
102
- )
103
-
104
- inferred_package_metadata = (
105
- await self._get_package_metadata_from_pyproject_toml(package_dir_to_publish)
106
- )
107
-
108
- if not inferred_package_metadata:
109
- return self.skip(
110
- "Connector does not have a pyproject.toml file which specifies package name and version and they are not set otherwise, skipping."
111
- )
112
-
113
- if not self.context.package_metadata.name:
114
- self.context.package_metadata.name = inferred_package_metadata.name
115
- if not self.context.package_metadata.version:
116
- self.context.package_metadata.version = inferred_package_metadata.version
117
-
118
- return None
119
-
120
- async def _publish(
121
- self, package_dir_to_publish: Directory, package_type: PackageType
122
- ) -> StepResult:
123
- if package_type is PackageType.PIP:
124
- return await self._pip_publish(package_dir_to_publish)
125
- else:
126
- return await self._poetry_publish(package_dir_to_publish)
127
-
128
- async def _poetry_publish(self, package_dir_to_publish: Directory) -> StepResult:
129
- pyproject_toml = package_dir_to_publish.file(PYPROJECT_TOML_FILE_PATH)
130
- pyproject_toml_content = await pyproject_toml.contents()
131
- contents = tomllib.loads(pyproject_toml_content)
132
- # make sure package name and version are set to the configured one
133
- contents["tool"]["poetry"]["name"] = self.context.package_metadata.name
134
- contents["tool"]["poetry"]["version"] = self.context.package_metadata.version
135
- # enforce consistent author
136
- contents["tool"]["poetry"]["authors"] = ["Airbyte <contact@airbyte.io>"]
137
- poetry_publish = (
138
- self._get_base_container()
139
- .with_secret_variable(
140
- "PYTHON_REGISTRY_TOKEN",
141
- self.context.python_registry_token.as_dagger_secret(self.dagger_client),
142
- )
143
- .with_directory("package", package_dir_to_publish)
144
- .with_workdir("package")
145
- .with_new_file(PYPROJECT_TOML_FILE_PATH, contents=tomli_w.dumps(contents))
146
- # Make sure these steps are always executed and not cached as they are triggering a side-effect (calling the registry)
147
- # Env var setting needs to be in this block as well to make sure a change of the env var will be propagated correctly
148
- .with_env_variable("CACHEBUSTER", str(uuid.uuid4()))
149
- .with_exec(
150
- ["poetry", "config", "repositories.mypypi", self.context.registry],
151
- use_entrypoint=True,
152
- )
153
- .with_exec(
154
- sh_dash_c(["poetry config pypi-token.mypypi $PYTHON_REGISTRY_TOKEN"])
155
- )
156
- # Default timeout is set to 15 seconds
157
- # We sometime face 443 HTTP read timeout responses from PyPi
158
- # Setting it to 60 seconds to avoid transient publish failures
159
- .with_env_variable("POETRY_REQUESTS_TIMEOUT", "60")
160
- .with_exec(
161
- sh_dash_c(
162
- ["poetry publish --build --repository mypypi -vvv --no-interaction"]
163
- )
164
- )
165
- )
166
-
167
- return await self.get_step_result(poetry_publish)
168
-
169
- async def _pip_publish(self, package_dir_to_publish: Directory) -> StepResult:
170
- files = await package_dir_to_publish.entries()
171
- metadata: Dict[str, str] = {
172
- "name": str(self.context.package_metadata.name),
173
- "version": str(self.context.package_metadata.version),
174
- # Enforce consistent author
175
- "author": "Airbyte",
176
- "author_email": "contact@airbyte.io",
177
- }
178
- if "README.md" in files:
179
- metadata["long_description"] = await package_dir_to_publish.file(
180
- "README.md"
181
- ).contents()
182
- metadata["long_description_content_type"] = "text/markdown"
183
-
184
- config = configparser.ConfigParser()
185
- config["metadata"] = metadata
186
-
187
- setup_cfg_io = io.StringIO()
188
- config.write(setup_cfg_io)
189
- setup_cfg = setup_cfg_io.getvalue()
190
-
191
- twine_upload = (
192
- self._get_base_container()
193
- .with_exec(sh_dash_c(["apt-get update", "apt-get install -y twine"]))
194
- .with_directory("package", package_dir_to_publish)
195
- .with_workdir("package")
196
- .with_exec(
197
- [
198
- "sed",
199
- "-i",
200
- "/name=/d; /author=/d; /author_email=/d; /version=/d",
201
- SETUP_PY_FILE_PATH,
202
- ],
203
- use_entrypoint=True,
204
- )
205
- .with_new_file("setup.cfg", contents=setup_cfg)
206
- .with_exec(
207
- ["pip", "install", "--upgrade", "setuptools", "wheel"],
208
- use_entrypoint=True,
209
- )
210
- .with_exec(
211
- ["python", SETUP_PY_FILE_PATH, "sdist", "bdist_wheel"],
212
- use_entrypoint=True,
213
- )
214
- # Make sure these steps are always executed and not cached as they are triggering a side-effect (calling the registry)
215
- # Env var setting needs to be in this block as well to make sure a change of the env var will be propagated correctly
216
- .with_env_variable("CACHEBUSTER", str(uuid.uuid4()))
217
- .with_secret_variable(
218
- "TWINE_USERNAME",
219
- self.context.dagger_client.set_secret("pypi_username", "__token__"),
220
- )
221
- .with_secret_variable(
222
- "TWINE_PASSWORD",
223
- self.context.python_registry_token.as_dagger_secret(self.dagger_client),
224
- )
225
- .with_exec(
226
- [
227
- "twine",
228
- "upload",
229
- "--verbose",
230
- "--repository-url",
231
- self.context.registry,
232
- "dist/*",
233
- ],
234
- use_entrypoint=True,
235
- )
236
- )
237
-
238
- return await self.get_step_result(twine_upload)
@@ -1,119 +0,0 @@
1
- #
2
- # Copyright (c) 2023 Airbyte, Inc., all rights reserved.
3
- #
4
-
5
- from dataclasses import dataclass
6
- from datetime import datetime
7
- from typing import Optional, Type
8
-
9
- from consts import DEFAULT_PYTHON_PACKAGE_REGISTRY_URL
10
- from pipelines.airbyte_ci.connectors.context import PipelineContext
11
- from pipelines.airbyte_ci.connectors.publish.context import PublishConnectorContext
12
- from pipelines.models.secrets import Secret
13
-
14
-
15
- @dataclass
16
- class PythonPackageMetadata:
17
- name: Optional[str]
18
- version: Optional[str]
19
-
20
-
21
- class PythonRegistryPublishContext(PipelineContext):
22
- def __init__(
23
- self,
24
- python_registry_token: Secret,
25
- registry_check_url: str,
26
- package_path: str,
27
- report_output_prefix: str,
28
- is_local: bool,
29
- git_branch: str,
30
- git_revision: str,
31
- diffed_branch: str,
32
- git_repo_url: str,
33
- ci_report_bucket: Optional[str] = None,
34
- registry: str = DEFAULT_PYTHON_PACKAGE_REGISTRY_URL,
35
- gha_workflow_run_url: Optional[str] = None,
36
- dagger_logs_url: Optional[str] = None,
37
- pipeline_start_timestamp: Optional[int] = None,
38
- ci_context: Optional[str] = None,
39
- ci_gcp_credentials: Optional[Secret] = None,
40
- package_name: Optional[str] = None,
41
- version: Optional[str] = None,
42
- ) -> None:
43
- self.python_registry_token = python_registry_token
44
- self.registry = registry
45
- self.registry_check_url = registry_check_url
46
- self.package_path = package_path
47
- self.package_metadata = PythonPackageMetadata(package_name, version)
48
-
49
- pipeline_name = f"Publish PyPI {package_path}"
50
-
51
- super().__init__(
52
- pipeline_name=pipeline_name,
53
- report_output_prefix=report_output_prefix,
54
- ci_report_bucket=ci_report_bucket,
55
- is_local=is_local,
56
- git_branch=git_branch,
57
- git_revision=git_revision,
58
- diffed_branch=diffed_branch,
59
- git_repo_url=git_repo_url,
60
- gha_workflow_run_url=gha_workflow_run_url,
61
- dagger_logs_url=dagger_logs_url,
62
- pipeline_start_timestamp=pipeline_start_timestamp,
63
- ci_context=ci_context,
64
- ci_gcp_credentials=ci_gcp_credentials,
65
- )
66
-
67
- @classmethod
68
- async def from_publish_connector_context(
69
- cls: Type["PythonRegistryPublishContext"],
70
- connector_context: PublishConnectorContext,
71
- ) -> Optional["PythonRegistryPublishContext"]:
72
- """
73
- Create a PythonRegistryPublishContext from a ConnectorContext.
74
-
75
- The metadata of the connector is read from the current workdir to capture changes that are not yet published.
76
- If pypi is not enabled, this will return None.
77
- """
78
-
79
- current_metadata = connector_context.connector.metadata
80
- connector_context.logger.info(f"Current metadata: {current_metadata!s}")
81
- if (
82
- "remoteRegistries" not in current_metadata
83
- or "pypi" not in current_metadata["remoteRegistries"]
84
- or not current_metadata["remoteRegistries"]["pypi"]["enabled"]
85
- ):
86
- return None
87
-
88
- version = current_metadata["dockerImageTag"]
89
- if connector_context.pre_release:
90
- # use current date as pre-release version
91
- # we can't use the git revision because not all python registries allow local version identifiers. Public version identifiers must conform to PEP 440 and only allow digits.
92
- release_candidate_tag = datetime.now().strftime("%Y%m%d%H%M")
93
- version = f"{version}.dev{release_candidate_tag}"
94
-
95
- assert connector_context.python_registry_token is not None, (
96
- "The connector context must have python_registry_token Secret attribute"
97
- )
98
- pypi_context = cls(
99
- python_registry_token=connector_context.python_registry_token,
100
- registry=str(connector_context.python_registry_url),
101
- registry_check_url=str(connector_context.python_registry_check_url),
102
- package_path=str(connector_context.connector.code_directory),
103
- package_name=current_metadata["remoteRegistries"]["pypi"]["packageName"],
104
- version=version,
105
- ci_report_bucket=connector_context.ci_report_bucket,
106
- report_output_prefix=connector_context.report_output_prefix,
107
- is_local=connector_context.is_local,
108
- git_branch=connector_context.git_branch,
109
- git_revision=connector_context.git_revision,
110
- diffed_branch=connector_context.diffed_branch,
111
- git_repo_url=connector_context.git_repo_url,
112
- gha_workflow_run_url=connector_context.gha_workflow_run_url,
113
- dagger_logs_url=connector_context.dagger_logs_url,
114
- pipeline_start_timestamp=connector_context.pipeline_start_timestamp,
115
- ci_context=connector_context.ci_context,
116
- ci_gcp_credentials=connector_context.ci_gcp_credentials,
117
- )
118
- pypi_context.dagger_client = connector_context.dagger_client
119
- return pypi_context
@@ -1,35 +0,0 @@
1
- # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
2
- """Vendored subset of connection-retriever from airbyte-platform-internal.
3
-
4
- This module contains a minimal subset of the connection-retriever tool,
5
- vendored to avoid depending on an unpublished internal package. It provides
6
- functionality to retrieve unmasked source configuration from Airbyte Cloud's
7
- internal database.
8
-
9
- Original source: airbyte-platform-internal/tools/connection-retriever
10
- Vendored: 2025-01-XX
11
-
12
- Only the following functionality is included:
13
- - Retrieve unmasked source config for a given connection ID
14
- - Secret resolution from GCP Secret Manager
15
- - Audit logging to GCP Cloud Logging
16
-
17
- NOT included (see issue #91 for future work):
18
- - retrieve_testing_candidates() - BigQuery-based candidate discovery
19
- - Destination config retrieval
20
- - CLI interface
21
- """
22
-
23
- from airbyte_ops_mcp.live_tests._connection_retriever.consts import (
24
- ConnectionObject,
25
- )
26
- from airbyte_ops_mcp.live_tests._connection_retriever.retrieval import (
27
- TestingCandidate,
28
- retrieve_objects,
29
- )
30
-
31
- __all__ = [
32
- "ConnectionObject",
33
- "TestingCandidate",
34
- "retrieve_objects",
35
- ]
@@ -1,33 +0,0 @@
1
- # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
2
- """Constants for vendored connection-retriever.
3
-
4
- Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/consts.py
5
- """
6
-
7
- from enum import Enum
8
-
9
- GCP_PROJECT_NAME = "prod-ab-cloud-proj"
10
-
11
- CLOUD_REGISTRY_URL = (
12
- "https://connectors.airbyte.com/files/registries/v0/cloud_registry.json"
13
- )
14
-
15
- CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID = (
16
- "projects/587336813068/secrets/CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS"
17
- )
18
-
19
-
20
- class ConnectionObject(Enum):
21
- """Types of connection objects that can be retrieved."""
22
-
23
- CONNECTION = "connection"
24
- SOURCE_ID = "source-id"
25
- DESTINATION_ID = "destination-id"
26
- DESTINATION_CONFIG = "destination-config"
27
- SOURCE_CONFIG = "source-config"
28
- CATALOG = "catalog"
29
- CONFIGURED_CATALOG = "configured-catalog"
30
- STATE = "state"
31
- WORKSPACE_ID = "workspace-id"
32
- DESTINATION_DOCKER_IMAGE = "destination-docker-image"
33
- SOURCE_DOCKER_IMAGE = "source-docker-image"
@@ -1,82 +0,0 @@
1
- # Copyright (c) 2025 Airbyte, Inc., all rights reserved.
2
- """Database access for vendored connection-retriever.
3
-
4
- Vendored from: airbyte-platform-internal/tools/connection-retriever/src/connection_retriever/db_access.py
5
- """
6
-
7
- from __future__ import annotations
8
-
9
- import json
10
- import os
11
- import traceback
12
- from typing import Any, Callable
13
-
14
- import sqlalchemy
15
- from google.cloud import secretmanager
16
- from google.cloud.sql.connector import Connector
17
- from google.cloud.sql.connector.enums import IPTypes
18
-
19
- from airbyte_ops_mcp.live_tests._connection_retriever.consts import (
20
- CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID,
21
- )
22
- from airbyte_ops_mcp.live_tests._connection_retriever.secrets_resolution import (
23
- get_secret_value,
24
- )
25
-
26
- PG_DRIVER = "pg8000"
27
-
28
- # Lazy-initialized to avoid import-time GCP auth
29
- _connector: Connector | None = None
30
-
31
-
32
- def _get_connector() -> Connector:
33
- """Get the Cloud SQL connector, initializing lazily on first use."""
34
- global _connector
35
- if _connector is None:
36
- _connector = Connector()
37
- return _connector
38
-
39
-
40
- def get_database_creator(pg_connection_details: dict) -> Callable:
41
- """Create a database connection creator function."""
42
-
43
- def creator() -> Any:
44
- return _get_connector().connect(
45
- pg_connection_details["database_address"],
46
- PG_DRIVER,
47
- user=pg_connection_details["pg_user"],
48
- password=pg_connection_details["pg_password"],
49
- db=pg_connection_details["database_name"],
50
- ip_type=IPTypes.PRIVATE,
51
- )
52
-
53
- return creator
54
-
55
-
56
- def get_pool(
57
- secret_manager_client: secretmanager.SecretManagerServiceClient,
58
- ) -> sqlalchemy.Engine:
59
- """Get a SQLAlchemy connection pool for the Airbyte Cloud database."""
60
- pg_connection_details = json.loads(
61
- get_secret_value(
62
- secret_manager_client, CONNECTION_RETRIEVER_PG_CONNECTION_DETAILS_SECRET_ID
63
- )
64
- )
65
-
66
- if os.getenv("CI"):
67
- # In CI we connect via Cloud SQL Auth Proxy, running on localhost
68
- host = "127.0.0.1"
69
- try:
70
- return sqlalchemy.create_engine(
71
- f"postgresql+{PG_DRIVER}://{pg_connection_details['pg_user']}:{pg_connection_details['pg_password']}@127.0.0.1/{pg_connection_details['database_name']}",
72
- )
73
- except Exception as e:
74
- raise AssertionError(
75
- f"sqlalchemy.create_engine exception; could not connect to the proxy at {host}. "
76
- f"Error: {traceback.format_exception(e)}"
77
- ) from e
78
- else:
79
- return sqlalchemy.create_engine(
80
- f"postgresql+{PG_DRIVER}://",
81
- creator=get_database_creator(pg_connection_details),
82
- )