implement review suggestions

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/bump_version/pipeline.py

@@ -153,8 +153,9 @@ def __init__(
 
     async def get_repo_dir(self) -> Directory:
         if not self.repo_dir:
-            self.repo_dir = await self.context.get_repo_dir()
-        return self.repo_dir
+            repo_dir = await self.context.get_repo_dir()
+            self.repo_dir = repo_dir
+        return repo_dir
 
     async def _run(self) -> StepResult:
         result = await self.update_metadata()

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/publish/commands.py

@@ -8,10 +8,10 @@
 from pipelines.airbyte_ci.connectors.pipeline import run_connectors_pipelines
 from pipelines.airbyte_ci.connectors.publish.context import PublishConnectorContext
 from pipelines.airbyte_ci.connectors.publish.pipeline import reorder_contexts, run_connector_publish_pipeline
-from pipelines.cli.airbyte_ci import wrap_gcp_credentials_in_secret, wrap_in_secret
 from pipelines.cli.click_decorators import click_ci_requirements_option
 from pipelines.cli.confirm_prompt import confirm
 from pipelines.cli.dagger_pipeline_command import DaggerPipelineCommand
+from pipelines.cli.secrets import wrap_gcp_credentials_in_secret, wrap_in_secret
 from pipelines.consts import DEFAULT_PYTHON_PACKAGE_REGISTRY_CHECK_URL, DEFAULT_PYTHON_PACKAGE_REGISTRY_URL, ContextState
 from pipelines.helpers.utils import fail_if_missing_docker_hub_creds
 from pipelines.models.secrets import Secret

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/connectors/test/steps/common.py

@@ -28,6 +28,41 @@
 from pipelines.helpers.utils import METADATA_FILE_NAME, get_exec_result
 from pipelines.models.secrets import Secret, SecretNotFoundError, SecretStore
 from pipelines.models.steps import STEP_PARAMS, MountPath, Step, StepResult, StepStatus
+from pydash import find  # type: ignore
+
+
+def _handle_missing_secret_store(
+    secret_info: Dict[str, str | Dict[str, str]], raise_on_missing: bool, logger: Optional[logging.Logger] = None
+) -> None:
+    assert isinstance(secret_info["secretStore"], dict), "The secretStore field must be a dict"
+    message = f"Secret {secret_info['name']} can't be retrieved as {secret_info['secretStore']['alias']} is not available"
+    if raise_on_missing:
+        raise SecretNotFoundError(message)
+    if logger is not None:
+        logger.warn(message)
+
+
+def _process_secret(
+    secret_info: Dict[str, str | Dict[str, str]],
+    secret_stores: Dict[str, SecretStore],
+    raise_on_missing: bool,
+    logger: Optional[logging.Logger] = None,
+) -> Optional[Secret]:
+    assert isinstance(secret_info["secretStore"], dict), "The secretStore field must be a dict"
+    secret_store_alias = secret_info["secretStore"]["alias"]
+    if secret_store_alias not in secret_stores:
+        _handle_missing_secret_store(secret_info, raise_on_missing, logger)
+        return None
+    else:
+        # All these asserts and casting are there to make MyPy happy
+        # The dict structure being nested MyPy can't figure if the values are str or dict
+        assert isinstance(secret_info["name"], str), "The secret name field must be a string"
+        if file_name := secret_info.get("fileName"):
+            assert isinstance(secret_info["fileName"], str), "The secret fileName must be a string"
+            file_name = str(secret_info["fileName"])
+        else:
+            file_name = None
+        return Secret(secret_info["name"], secret_stores[secret_store_alias], file_name=file_name)
 
 
 def get_secrets_from_connector_test_suites_option(
@@ -56,23 +91,12 @@ def get_secrets_from_connector_test_suites_option(
         List[Secret]: List of secrets declared in the connectorTestSuitesOptions for a test suite name.
     """
     secrets: List[Secret] = []
+    enabled_test_suite = find(connector_test_suites_options, lambda x: x["suite"] == suite_name)
 
-    for enabled_test_suite in connector_test_suites_options:
-        if enabled_test_suite["suite"] == suite_name:
-            if enabled_test_suite.get("testSecrets"):
-                assert isinstance(enabled_test_suite["testSecrets"], list)
-                suite_secrets: List[Dict[str, str | Dict[str, str]]] = enabled_test_suite["testSecrets"]
-                for s in suite_secrets:
-                    if s["secretStore"]["alias"] not in secret_stores:
-                        message = f"Secret {s['name']} can't be retrieved as {s['secretStore']['alias']} is not available"
-                        if raise_on_missing_secret_store:
-                            raise SecretNotFoundError(message)
-                        if logger is not None:
-                            logger.warn(message)
-                        continue
-                    secret_store = secret_stores[s["secretStore"]["alias"]]
-                    secret = Secret(s["name"], secret_store, file_name=s.get("fileName"))
-                    secrets.append(secret)
+    if enabled_test_suite and "testSecrets" in enabled_test_suite:
+        for secret_info in enabled_test_suite["testSecrets"]:
+            if secret := _process_secret(secret_info, secret_stores, raise_on_missing_secret_store, logger):
+                secrets.append(secret)
 
     return secrets
 

airbyte-ci/connectors/pipelines/pipelines/airbyte_ci/poetry/publish/commands.py

@@ -12,9 +12,9 @@
 import asyncclick as click
 from packaging import version
 from pipelines.airbyte_ci.steps.python_registry import PublishToPythonRegistry
-from pipelines.cli.airbyte_ci import wrap_in_secret
 from pipelines.cli.confirm_prompt import confirm
 from pipelines.cli.dagger_pipeline_command import DaggerPipelineCommand
+from pipelines.cli.secrets import wrap_in_secret
 from pipelines.consts import DEFAULT_PYTHON_PACKAGE_REGISTRY_CHECK_URL, DEFAULT_PYTHON_PACKAGE_REGISTRY_URL
 from pipelines.models.contexts.click_pipeline_context import ClickPipelineContext, pass_pipeline_context
 from pipelines.models.contexts.python_registry_publish import PythonRegistryPublishContext

airbyte-ci/connectors/pipelines/pipelines/cli/airbyte_ci.py

@@ -17,7 +17,7 @@
 import multiprocessing
 import os
 import sys
-from typing import Any, Optional
+from typing import Optional
 
 import asyncclick as click
 import docker  # type: ignore
@@ -33,14 +33,14 @@
 )
 from pipelines.cli.confirm_prompt import pre_confirm_all_flag
 from pipelines.cli.lazy_group import LazyGroup
+from pipelines.cli.secrets import wrap_gcp_credentials_in_secret, wrap_in_secret
 from pipelines.cli.telemetry import click_track_command
 from pipelines.consts import DAGGER_WRAP_ENV_VAR_NAME, LOCAL_BUILD_PLATFORM, CIContext
 from pipelines.dagger.actions.connector.hooks import get_dagger_sdk_version
 from pipelines.helpers import github
-from pipelines.helpers.gcs import sanitize_gcp_credentials
 from pipelines.helpers.git import get_current_git_branch, get_current_git_revision
 from pipelines.helpers.utils import AIRBYTE_REPO_URL, get_current_epoch_time
-from pipelines.models.secrets import InMemorySecretStore, Secret
+from pipelines.models.secrets import InMemorySecretStore
 
 
 def log_context_info(ctx: click.Context) -> None:
@@ -126,31 +126,6 @@ def is_current_process_wrapped_by_dagger_run() -> bool:
     return called_with_dagger_run
 
 
-def wrap_in_secret(ctx: click.Context, param: click.Option, value: Any) -> Optional[Secret]:  # noqa
-    if value is None:
-        return None
-    assert param.name is not None
-    if not isinstance(value, str):
-        raise click.BadParameter(f"{param.name} value is not a string, only strings can be wrapped in a secret.")
-    ctx.ensure_object(dict)
-    if "secret_stores" not in ctx.obj:
-        ctx.obj["secret_stores"] = {}
-    if "in_memory" not in ctx.obj["secret_stores"]:
-        ctx.obj["secret_stores"]["in_memory"] = InMemorySecretStore()
-
-    ctx.obj["secret_stores"]["in_memory"].add_secret(param.name, value)
-    return Secret(param.name, ctx.obj["secret_stores"]["in_memory"])
-
-
-def wrap_gcp_credentials_in_secret(ctx: click.Context, param: click.Option, value: Any) -> Optional[Secret]:  # noqa
-    if value is None:
-        return None
-    if not isinstance(value, str):
-        raise click.BadParameter(f"{param.name} value is not a string, only strings can be wrapped in a secret.")
-    value = sanitize_gcp_credentials(value)
-    return wrap_in_secret(ctx, param, value)
-
-
 # COMMANDS
 
 

airbyte-ci/connectors/pipelines/pipelines/cli/secrets.py

@@ -0,0 +1,40 @@
+# Copyright (c) 2024 Airbyte, Inc., all rights reserved.
+
+from typing import Any, Optional
+
+import asyncclick as click
+from pipelines.helpers.gcs import sanitize_gcp_credentials
+from pipelines.models.secrets import InMemorySecretStore, Secret
+
+
+def wrap_in_secret(ctx: click.Context, param: click.Option, value: Any) -> Optional[Secret]:  # noqa
+    # Validate callback usage
+    if value is None:
+        return None
+    assert param.name is not None
+    if not isinstance(value, str):
+        raise click.BadParameter(f"{param.name} value is not a string, only strings can be wrapped in a secret.")
+
+    # Make sure the context object is set or set it with an empty dict
+    ctx.ensure_object(dict)
+
+    # Instantiate a global in memory secret store in the context object if it's not yet set
+    if "secret_stores" not in ctx.obj:
+        ctx.obj["secret_stores"] = {}
+    if "in_memory" not in ctx.obj["secret_stores"]:
+        ctx.obj["secret_stores"]["in_memory"] = InMemorySecretStore()
+
+    # Add the CLI option value to the in memory secret store and wrap it in a Secret
+    ctx.obj["secret_stores"]["in_memory"].add_secret(param.name, value)
+    return Secret(param.name, ctx.obj["secret_stores"]["in_memory"])
+
+
+def wrap_gcp_credentials_in_secret(ctx: click.Context, param: click.Option, value: Any) -> Optional[Secret]:  # noqa
+    # Validate callback usage
+    if value is None:
+        return None
+    if not isinstance(value, str):
+        raise click.BadParameter(f"{param.name} value is not a string, only strings can be wrapped in a secret.")
+
+    value = sanitize_gcp_credentials(value)
+    return wrap_in_secret(ctx, param, value)

airbyte-ci/connectors/pipelines/pipelines/models/secrets.py

@@ -54,17 +54,19 @@ def _fetch_secret(self, name: str) -> str:
         request = secretmanager_v1.ListSecretVersionsRequest(
             parent=f"projects/{self.project_id}/secrets/{name}",
         )
-        page_result = self.gsm_client.list_secret_versions(request=request)
-        if not page_result:
+        secret_versions = self.gsm_client.list_secret_versions(request=request)
+        if not secret_versions:
             raise SecretNotFoundError(f"No secret found in GSM for secret {name}")
-        for version in page_result:
+
+        for version in secret_versions:
             # 1 means enabled version
             if version.state == 1:
                 request = secretmanager_v1.AccessSecretVersionRequest(
                     name=version.name,
                 )
                 response = self.gsm_client.access_secret_version(request=request)
                 return response.payload.data.decode()
+
         raise SecretNotFoundError(f"No enabled secret version in GSM found for secret {name}")
 
 

airbyte-ci/connectors/pipelines/pyproject.toml

@@ -37,6 +37,7 @@ ruamel-yaml = "^0.18.6"
 google-cloud-secret-manager = "^2.20.0"
 google-auth = "^2.29.0"
 pygithub = "^2.3.0"
+pydash = "6.0.2"
 
 [tool.poetry.group.dev.dependencies]
 freezegun = "^1.2.2"