awscli to boto3

Author: jx2lee

cosmos/operators/kubernetes.py

@@ -1,5 +1,8 @@
 from __future__ import annotations
 
+import json
+import shlex
+import textwrap
 from abc import ABC, abstractmethod
 from collections.abc import Callable, Sequence
 from typing import TYPE_CHECKING, Any
@@ -28,11 +31,6 @@
     DbtTestMixin,
 )
 
-try:
-    from airflow.sdk.bases.hook import BaseHook
-except ImportError:  # Since Airflow 3.1, BaseHook is in the airflow.sdk.bases.hook module
-    from airflow.hooks.base import BaseHook
-
 
 class DbtKubernetesBaseOperator(AbstractDbtBase, KubernetesPodOperator):  # type: ignore[misc]
     """
@@ -218,14 +216,22 @@ def __init__(self, **kwargs: Any) -> None:
     @abstractmethod
     def build_upload_shell_command(self, docs_target: str) -> str:
         """
-        Build the shell command that will upload the generated docs from
-        `docs_target` to cloud storage. Implemented by subclasses.
+        Build the shell command that uploads generated docs from `docs_target`
+        to cloud storage inside the Kubernetes Pod.
         """
 
     @abstractmethod
     def get_upload_env_vars(self) -> dict[str, str]:
         """Return env vars required by the upload command."""
 
+    @staticmethod
+    def _command_parts(command: Any) -> list[str]:
+        if not command:
+            return []
+        if isinstance(command, (list, tuple)):
+            return [str(part) for part in command]
+        return [str(command)]
+
     def build_and_run_cmd(
         self,
         context: Context,
@@ -244,25 +250,20 @@ def build_and_run_cmd(
         # the leading "dbt" is not dropped when folded into the bash -c string below (see PR #2488).
         cmds: Any = self.cmds  # type: ignore[has-type]
         arguments: Any = self.arguments  # type: ignore[has-type]
-        cmd_parts = list(cmds) if isinstance(cmds, (list, tuple)) else [cmds]
-        if isinstance(arguments, (list, tuple)):
-            cmd_parts.extend(arguments)
-        else:
-            cmd_parts.append(arguments)
-
-        dbt_cmd_str = " ".join([str(part) for part in cmd_parts])
+        cmd_parts = self._command_parts(cmds) + self._command_parts(arguments)
+        dbt_cmd_str = shlex.join(cmd_parts)
         docs_target = f"{self.project_dir}/target"
 
         upload_cmd = self.build_upload_shell_command(docs_target)
         shell_cmd = f"{dbt_cmd_str} && {upload_cmd}"
 
-        # Override container command and arguments
         self.cmds = ["/bin/bash", "-c"]
         self.arguments = [shell_cmd]
 
         self.log.info("Running command in Kubernetes Pod: %s", self.arguments)
         result = KubernetesPodOperator.execute(self, context)
         self.log.info(result)
+
         return result
 
     def inject_upload_env_vars(self, env_vars: dict[str, str]) -> None:
@@ -286,7 +287,7 @@ def inject_upload_env_vars(self, env_vars: dict[str, str]) -> None:
 class DbtDocsS3KubernetesOperator(DbtDocsCloudKubernetesOperator):
     """
     Executes `dbt docs generate` inside a Kubernetes Pod and uploads the generated
-    documentation to S3 *also inside that Pod* using `aws s3 sync`.
+    documentation to S3 also inside that Pod using ``boto3``.
         - The Kubernetes Pod receives AWS credentials resolved from the supplied
           Airflow `connection_id`.
     """
@@ -306,45 +307,97 @@ def __init__(
         self.folder_dir = folder_dir
 
     def build_upload_shell_command(self, docs_target: str) -> str:
-        if self.folder_dir:
-            s3_prefix = f"s3://{self.bucket_name}/{self.folder_dir}".rstrip("/")
-        else:
-            s3_prefix = f"s3://{self.bucket_name}"
-
-        return f"aws s3 sync {docs_target} {s3_prefix}"
+        folder_dir = self.folder_dir.rstrip("/") if self.folder_dir else ""
+        upload_script = textwrap.dedent(f"""
+            import json
+            import mimetypes
+            import os
+            from pathlib import Path
+
+            try:
+                import boto3
+            except ImportError as exc:
+                raise SystemExit("boto3 is required in the Kubernetes image to upload dbt docs to S3.") from exc
+
+            target_dir = Path({json.dumps(docs_target)})
+            bucket_name = {json.dumps(self.bucket_name)}
+            folder_dir = {json.dumps(folder_dir)}
+
+            client_kwargs = dict()
+            endpoint_url = os.environ.get("AWS_ENDPOINT_URL_S3")
+            if endpoint_url:
+                client_kwargs["endpoint_url"] = endpoint_url
+
+            client_config = json.loads(os.environ.get("COSMOS_AWS_CLIENT_CONFIG", "{{}}"))
+            if "verify" in client_config:
+                client_kwargs["verify"] = client_config["verify"]
+
+            config_kwargs = client_config.get("config_kwargs")
+            if config_kwargs:
+                from botocore.config import Config
+
+                client_kwargs["config"] = Config(**config_kwargs)
+
+            s3 = boto3.client("s3", **client_kwargs)
+            for file_path in target_dir.rglob("*"):
+                if not file_path.is_file():
+                    continue
+
+                relative_path = file_path.relative_to(target_dir).as_posix()
+                key = f"{{folder_dir}}/{{relative_path}}" if folder_dir else relative_path
+                content_type, _ = mimetypes.guess_type(str(file_path))
+                extra_args = {{"ContentType": content_type}} if content_type else None
+                print(f"Uploading {{file_path}} to s3://{{bucket_name}}/{{key}}")
+                if extra_args:
+                    s3.upload_file(str(file_path), bucket_name, key, ExtraArgs=extra_args)
+                else:
+                    s3.upload_file(str(file_path), bucket_name, key)
+            """).strip()
+        return f"$(command -v python3 || command -v python) - <<'PY'\n{upload_script}\nPY"
 
     def get_upload_env_vars(self) -> dict[str, str]:
         return self.aws_env_vars_from_connection(self.connection_id)
 
     def aws_env_vars_from_connection(self, connection_id: str) -> dict[str, str]:
-        conn = BaseHook.get_connection(connection_id)
-        conn_extra = conn.extra_dejson
-
-        access_key = conn.login or conn_extra.get("aws_access_key_id")
-        secret_key = conn.password or conn_extra.get("aws_secret_access_key")
-        session_token = conn_extra.get("aws_session_token") or conn_extra.get("session_token")
-
-        session_kwargs = conn_extra.get("session_kwargs", {})
-        config_kwargs = conn_extra.get("config_kwargs", {})
-        if not isinstance(session_kwargs, dict):
-            session_kwargs = {}
-        if not isinstance(config_kwargs, dict):
-            config_kwargs = {}
-        region_name = (
-            conn_extra.get("region_name")
-            or conn_extra.get("region")
-            or session_kwargs.get("region_name")
-            or config_kwargs.get("region_name")
-        )
+        try:
+            from airflow.providers.amazon.aws.hooks.base_aws import AwsBaseHook
+        except ImportError:
+            from cosmos.operators.lazy_load import MissingPackage
+
+            AwsBaseHook = MissingPackage(
+                "airflow.providers.amazon.aws.hooks.base_aws.AwsBaseHook",
+                "amazon",
+            )
+
+        hook = AwsBaseHook(aws_conn_id=connection_id, client_type="s3")
+        conn_config = hook.conn_config
+        conn_extra = conn_config.extra_config
+
+        config_kwargs = conn_extra.get("config_kwargs") or {}
+
+        region_name = hook.region_name or conn_extra.get("region") or config_kwargs.get("region_name")
+        endpoint_url = conn_config.get_service_endpoint_url("s3")
+        verify = hook.verify
 
         env_vars = {}
-        if access_key:
-            env_vars["AWS_ACCESS_KEY_ID"] = access_key
-        if secret_key:
-            env_vars["AWS_SECRET_ACCESS_KEY"] = secret_key
-        if session_token:
-            env_vars["AWS_SESSION_TOKEN"] = session_token
+        if conn_config.aws_access_key_id:
+            env_vars["AWS_ACCESS_KEY_ID"] = conn_config.aws_access_key_id
+        if conn_config.aws_secret_access_key:
+            env_vars["AWS_SECRET_ACCESS_KEY"] = conn_config.aws_secret_access_key
+        if conn_config.aws_session_token:
+            env_vars["AWS_SESSION_TOKEN"] = conn_config.aws_session_token
+        if endpoint_url:
+            env_vars["AWS_ENDPOINT_URL_S3"] = endpoint_url
         if region_name:
             env_vars["AWS_DEFAULT_REGION"] = region_name
+            env_vars["AWS_REGION"] = region_name
+
+        client_config = {}
+        if verify is not None:
+            client_config["verify"] = verify
+        if config_kwargs:
+            client_config["config_kwargs"] = config_kwargs
+        if client_config:
+            env_vars["COSMOS_AWS_CLIENT_CONFIG"] = json.dumps(client_config)
 
         return env_vars

dev/Dockerfile.postgres_profile_docker_k8s

@@ -2,7 +2,7 @@ FROM python:3.12
 
 RUN pip install --force-reinstall 'dbt-postgres>=1.8' 'dbt-core<2.0'
 RUN pip install --force-reinstall dbt-adapters
-RUN pip install --force-reinstall awscli
+RUN pip install --force-reinstall boto3
 
 ENV POSTGRES_DATABASE=postgres
 ENV POSTGRES_DB=postgres

dev/Dockerfile.watcher_failing_tests_k8s

@@ -2,7 +2,7 @@ FROM python:3.12
 
 RUN pip install --force-reinstall 'dbt-postgres>=1.8' 'dbt-core<2.0'
 RUN pip install --force-reinstall dbt-adapters
-RUN pip install --force-reinstall awscli
+RUN pip install --force-reinstall boto3
 
 ENV POSTGRES_DATABASE=postgres
 ENV POSTGRES_DB=postgres

docs/guides/dbt_docs/generating-docs.rst

@@ -139,7 +139,7 @@ Requirements specific to Kubernetes:
 
 - The container image must include your dbt project files.
 - The container image or mounted files must include a ``profiles.yml`` file, because Kubernetes execution mode does not support :ref:`use-profile-mapping`.
-- The container image must have the AWS CLI available because Cosmos uploads the generated docs with ``aws s3 sync``.
+- The container image must have ``boto3`` available because Cosmos uploads the generated docs from inside the Pod with a small Python uploader.
 - The Pod still needs the database credentials and any other secrets required to run ``dbt docs generate``.
 
 The following example extends the Kubernetes example DAG and uploads the generated docs to S3:
@@ -149,7 +149,7 @@ The following example extends the Kubernetes example DAG and uploads the generat
    :start-after: [START kubernetes_docs_to_s3_example]
    :end-before: [END kubernetes_docs_to_s3_example]
 
-The ``connection_id`` is resolved from Airflow and translated into AWS environment variables that are injected into the Pod before ``aws s3 sync`` runs.
+The ``connection_id`` is resolved from Airflow with ``AwsBaseHook`` and translated into AWS environment variables that are injected into the Pod before the ``boto3`` uploader runs.
 
 .. note::
     This Kubernetes integration currently supports S3 only. If you need another storage backend, use one of the local operators or extend Cosmos with another Kubernetes docs operator.

docs/guides/run_dbt/container/kubernetes.rst

@@ -38,7 +38,7 @@ At the moment, the user is expected to add to the Docker image both:
 - The dbt Profile, which contains the information for dbt to access the database while parsing the project from Apache Airflow nodes
 - Handle secrets
 
-If you plan to generate dbt docs and upload them to S3 from Kubernetes, the image also needs the AWS CLI because Cosmos performs the upload from inside the Pod.
+If you plan to generate dbt docs and upload them to S3 from Kubernetes, the Pod image must have ``boto3`` available, either installed directly or provided by packages such as ``apache-airflow-providers-amazon``.
 
 Additional KubernetesPodOperator parameters can be added to the ``operator_args`` parameter of the ``DbtKubernetesOperator``.