Run some example in Kubernetes execution mode in CI

.github/workflows/test.yml

@@ -163,7 +163,6 @@ jobs:
           POSTGRES_DB: postgres
           POSTGRES_SCHEMA: public
           POSTGRES_PORT: 5432
-          SOURCE_RENDERING_BEHAVIOR: all
 
       - name: Upload coverage to Github
         uses: actions/upload-artifact@v2
@@ -235,7 +234,6 @@ jobs:
           POSTGRES_DB: postgres
           POSTGRES_SCHEMA: public
           POSTGRES_PORT: 5432
-          SOURCE_RENDERING_BEHAVIOR: all
 
       - name: Upload coverage to Github
         uses: actions/upload-artifact@v2
@@ -379,7 +377,6 @@ jobs:
           POSTGRES_DB: postgres
           POSTGRES_SCHEMA: public
           POSTGRES_PORT: 5432
-          SOURCE_RENDERING_BEHAVIOR: all
 
       - name: Upload coverage to Github
         uses: actions/upload-artifact@v2
@@ -461,12 +458,82 @@ jobs:
       AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:[email protected]:5432/postgres
       PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH
 
+  Run-Kubernetes-Tests:
+    needs: Authorize
+    runs-on: ubuntu-latest
+    strategy:
+      matrix:
+        python-version: [ "3.11" ]
+        airflow-version: [ "2.9" ]
+    steps:
+      - uses: actions/checkout@v3
+        with:
+          ref: ${{ github.event.pull_request.head.sha || github.ref }}
+      - uses: actions/cache@v3
+        with:
+          path: |
+            ~/.cache/pip
+            .local/share/hatch/
+          key: coverage-integration-kubernetes-test-${{ runner.os }}-${{ matrix.python-version }}-${{ matrix.airflow-version }}-${{ hashFiles('pyproject.toml') }}-${{ hashFiles('cosmos/__init__.py') }}
+
+      - name: Set up Python ${{ matrix.python-version }}
+        uses: actions/setup-python@v4
+        with:
+          python-version: ${{ matrix.python-version }}
+
+      - name: Create KinD cluster
+        uses: container-tools/kind-action@v1
+
+      - name: Install packages and dependencies
+        run: |
+          python -m venv venv
+          source venv/bin/activate
+          pip install --upgrade pip
+          pip install -e ".[tests]"
+          pip install apache-airflow-providers-cncf-kubernetes
+          pip install dbt-postgres==1.8.2 psycopg2==2.9.3 pytz
+          pip install apache-airflow==${{ matrix.airflow-version }}
+          # hatch -e tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }} run pip freeze
+
+      - name: Run kubernetes tests
+        run: |
+          source venv/bin/activate
+          sh ./scripts/test/kubernetes-setup.sh
+          cd dev && sh ../scripts/test/integration-kubernetes.sh
+          # hatch run tests.py${{ matrix.python-version }}-${{ matrix.airflow-version }}:test-kubernetes
+        env:
+          AIRFLOW_HOME: /home/runner/work/astronomer-cosmos/astronomer-cosmos/
+          AIRFLOW_CONN_EXAMPLE_CONN: postgres://postgres:[email protected]:5432/postgres
+          AIRFLOW_CONN_AWS_S3_CONN: ${{ secrets.AIRFLOW_CONN_AWS_S3_CONN }}
+          AIRFLOW_CONN_GCP_GS_CONN: ${{ secrets.AIRFLOW_CONN_GCP_GS_CONN }}
+          AIRFLOW_CONN_AZURE_ABFS_CONN: ${{ secrets.AIRFLOW_CONN_AZURE_ABFS_CONN }}
+          AIRFLOW__CORE__DAGBAG_IMPORT_TIMEOUT: 90.0
+          PYTHONPATH: /home/runner/work/astronomer-cosmos/astronomer-cosmos/:$PYTHONPATH
+          COSMOS_CONN_POSTGRES_PASSWORD: ${{ secrets.COSMOS_CONN_POSTGRES_PASSWORD }}
+          DATABRICKS_CLUSTER_ID: mock
+          DATABRICKS_HOST: mock
+          DATABRICKS_WAREHOUSE_ID: mock
+          DATABRICKS_TOKEN: mock
+          POSTGRES_HOST: localhost
+          POSTGRES_USER: postgres
+          POSTGRES_PASSWORD: postgres
+          POSTGRES_DB: postgres
+          POSTGRES_SCHEMA: public
+          POSTGRES_PORT: 5432
+
+      - name: Upload coverage to Github
+        uses: actions/upload-artifact@v2
+        with:
+          name: coverage-integration-kubernetes-test-${{ matrix.python-version }}-${{ matrix.airflow-version }}
+          path: .coverage
+
   Code-Coverage:
     if: github.event.action != 'labeled'
     needs:
       - Run-Unit-Tests
       - Run-Integration-Tests
       - Run-Integration-Tests-Expensive
+      - Run-Kubernetes-Tests
     runs-on: ubuntu-latest
     steps:
       - uses: actions/checkout@v3

dev/Dockerfile.postgres_profile_docker_k8s

@@ -0,0 +1,18 @@
+FROM python:3.11
+
+RUN pip install dbt-postgres==1.8.2 psycopg2==2.9.3 pytz
+
+ENV POSTGRES_DATABASE=postgres
+ENV POSTGRES_DB=postgres
+ENV POSTGRES_HOST=postgres.default.svc.cluster.local
+ENV POSTGRES_PASSWORD=postgres
+ENV POSTGRES_PORT=5432
+ENV POSTGRES_SCHEMA=public
+ENV POSTGRES_USER=postgres
+
+RUN mkdir /root/.dbt
+COPY dags/dbt/jaffle_shop/profiles.yml /root/.dbt/profiles.yml
+
+RUN mkdir dags
+COPY dags dags
+RUN rm dags/dbt/jaffle_shop/packages.yml

dev/dags/dbt/jaffle_shop/profiles.yml

@@ -10,3 +10,15 @@ default:
       dbname: "{{ env_var('POSTGRES_DB') }}"
       schema: "{{ env_var('POSTGRES_SCHEMA') }}"
       threads: 4
+
+postgres_profile:
+  target: dev
+  outputs:
+    dev:
+      type: postgres
+      dbname: "{{ env_var('POSTGRES_DATABASE') }}"
+      host: "{{ env_var('POSTGRES_HOST') }}"
+      pass: "{{ env_var('POSTGRES_PASSWORD') }}"
+      port: 5432 # "{{ env_var('POSTGRES_PORT') | as_number }}"
+      schema: "{{ env_var('POSTGRES_SCHEMA') }}"
+      user: "{{ env_var('POSTGRES_USER') }}"

dev/dags/jaffle_shop_kubernetes.py

@@ -0,0 +1,98 @@
+"""
+## Jaffle Shop DAG
+[Jaffle Shop](https://github.com/dbt-labs/jaffle_shop) is a fictional eCommerce store. This dbt project originates from
+dbt labs as an example project with dummy data to demonstrate a working dbt core project. This DAG uses the cosmos dbt
+parser to generate an Airflow TaskGroup from the dbt project folder.
+
+
+The step-by-step to run this DAG are described in:
+https://astronomer.github.io/astronomer-cosmos/getting_started/kubernetes.html#kubernetes
+
+"""
+
+from airflow import DAG
+from airflow.providers.cncf.kubernetes.secret import Secret
+from pendulum import datetime
+
+from cosmos import (
+    DbtSeedKubernetesOperator,
+    DbtTaskGroup,
+    ExecutionConfig,
+    ExecutionMode,
+    ProfileConfig,
+    ProjectConfig,
+)
+from cosmos.profiles import PostgresUserPasswordProfileMapping
+
+DBT_IMAGE = "dbt-jaffle-shop:1.0.0"
+
+project_seeds = [{"project": "jaffle_shop", "seeds": ["raw_customers", "raw_payments", "raw_orders"]}]
+
+postgres_password_secret = Secret(
+    deploy_type="env",
+    deploy_target="POSTGRES_PASSWORD",
+    secret="postgres-secrets",
+    key="password",
+)
+
+postgres_host_secret = Secret(
+    deploy_type="env",
+    deploy_target="POSTGRES_HOST",
+    secret="postgres-secrets",
+    key="host",
+)
+
+with DAG(
+    dag_id="jaffle_shop_kubernetes",
+    start_date=datetime(2022, 11, 27),
+    doc_md=__doc__,
+    catchup=False,
+) as dag:
+    # [START kubernetes_seed_example]
+    load_seeds = DbtSeedKubernetesOperator(
+        task_id="load_seeds",
+        project_dir="dags/dbt/jaffle_shop",
+        get_logs=True,
+        schema="public",
+        image=DBT_IMAGE,
+        is_delete_operator_pod=False,
+        secrets=[postgres_password_secret, postgres_host_secret],
+        profile_config=ProfileConfig(
+            profile_name="postgres_profile",
+            target_name="dev",
+            profile_mapping=PostgresUserPasswordProfileMapping(
+                conn_id="postgres_default",
+                profile_args={
+                    "schema": "public",
+                },
+            ),
+        ),
+    )
+    # [END kubernetes_seed_example]
+
+    # [START kubernetes_tg_example]
+    run_models = DbtTaskGroup(
+        profile_config=ProfileConfig(
+            profile_name="postgres_profile",
+            target_name="dev",
+            profile_mapping=PostgresUserPasswordProfileMapping(
+                conn_id="postgres_default",
+                profile_args={
+                    "schema": "public",
+                },
+            ),
+        ),
+        project_config=ProjectConfig(dbt_project_path="dags/dbt/jaffle_shop"),
+        execution_config=ExecutionConfig(
+            execution_mode=ExecutionMode.KUBERNETES,
+        ),
+        operator_args={
+            "image": DBT_IMAGE,
+            "get_logs": True,
+            "is_delete_operator_pod": False,
+            "secrets": [postgres_password_secret, postgres_host_secret],
+        },
+    )
+    # [END kubernetes_tg_example]
+
+    load_seeds >> run_models

docs/getting_started/execution-modes.rst

@@ -144,27 +144,11 @@ Check the step-by-step guide on using the ``kubernetes`` execution mode at :ref:
 
 Example DAG:
 
-.. code-block:: python
-
-    postgres_password_secret = Secret(
-        deploy_type="env",
-        deploy_target="POSTGRES_PASSWORD",
-        secret="postgres-secrets",
-        key="password",
-    )
+.. literalinclude:: ../../dev/dags/jaffle_shop_kubernetes.py
+   :language: python
+   :start-after: [START kubernetes_seed_example]
+   :end-before: [END kubernetes_seed_example]
 
-    docker_cosmos_dag = DbtDag(
-        # ...
-        execution_config=ExecutionConfig(
-            execution_mode=ExecutionMode.KUBERNETES,
-        ),
-        operator_args={
-            "image": "dbt-jaffle-shop:1.0.0",
-            "get_logs": True,
-            "is_delete_operator_pod": False,
-            "secrets": [postgres_password_secret],
-        },
-    )
 AWS_EKS
 ----------
 

docs/getting_started/kubernetes.rst

@@ -28,30 +28,10 @@ Additional KubernetesPodOperator parameters can be added on the operator_args pa
 
 For instance,
 
-.. code-block:: python
-
-    run_models = DbtTaskGroup(
-        profile_config=ProfileConfig(
-            profile_name="postgres_profile",
-            target_name="dev",
-            profile_mapping=PostgresUserPasswordProfileMapping(
-                conn_id="postgres_default",
-                profile_args={
-                    "schema": "public",
-                },
-            ),
-        ),
-        project_config=ProjectConfig(PROJECT_DIR),
-        execution_config=ExecutionConfig(
-            execution_mode=ExecutionMode.KUBERNETES,
-        ),
-        operator_args={
-            "image": DBT_IMAGE,
-            "get_logs": True,
-            "is_delete_operator_pod": False,
-            "secrets": [postgres_password_secret, postgres_host_secret],
-        },
-    )
+.. literalinclude:: ../../dev/dags/jaffle_shop_kubernetes.py
+   :language: python
+   :start-after: [START kubernetes_tg_example]
+   :end-before: [END kubernetes_tg_example]
 
 Step-by-step instructions
 +++++++++++++++++++++++++

pyproject.toml

@@ -168,6 +168,7 @@ freeze = "pip freeze"
 test = 'sh scripts/test/unit.sh'
 test-cov = 'sh scripts/test/unit-cov.sh'
 test-integration = 'sh scripts/test/integration.sh'
+test-kubernetes = "sh scripts/test/integration-kubernetes.sh"
 test-integration-dbt-1-5-4 = 'sh scripts/test/integration-dbt-1-5-4.sh'
 test-integration-expensive = 'sh scripts/test/integration-expensive.sh'
 test-integration-setup = 'sh scripts/test/integration-setup.sh'

scripts/test/integration-dbt-1-5-4.sh

@@ -10,4 +10,5 @@ pytest -vv \
     --durations=0 \
     -m integration  \
     --ignore=tests/perf \
+    --ignore=tests/test_example_k8s_dags.py \
     -k 'basic_cosmos_task_group'

scripts/test/integration-expensive.sh

@@ -6,4 +6,5 @@ pytest -vv \
     --durations=0 \
     -m integration  \
     --ignore=tests/perf \
+    --ignore=tests/test_example_k8s_dags.py \
     -k 'example_cosmos_python_models or example_virtualenv'

scripts/test/integration-kubernetes.sh

@@ -0,0 +1,16 @@
+#!/bin/bash
+
+set -x
+set -e
+
+# Reset the Airflow database to its initial state
+airflow db reset -y
+
+# Run tests using pytest
+pytest -vv \
+    --cov=cosmos \
+    --cov-report=term-missing \
+    --cov-report=xml \
+    --durations=0 \
+    -m integration \
+    ../tests/test_example_k8s_dags.py

scripts/test/integration-sqlite.sh

@@ -5,4 +5,5 @@ pytest -vv \
     --durations=0 \
     -m integration  \
     --ignore=tests/perf \
+    --ignore=tests/test_example_k8s_dags.py \
     -k 'example_cosmos_sources or sqlite'

scripts/test/integration.sh

@@ -20,4 +20,5 @@ pytest -vv \
     --durations=0 \
     -m integration  \
     --ignore=tests/perf \
-    -k 'not (sqlite or example_cosmos_sources or example_cosmos_python_models or example_virtualenv)'
+    --ignore=tests/test_example_k8s_dags.py \
+    -k 'not (sqlite or example_cosmos_sources or example_cosmos_python_models or example_virtualenv or jaffle_shop_kubernetes)'

scripts/test/kubernetes-setup.sh

@@ -0,0 +1,34 @@
+#!/bin/bash
+
+# Print each command before executing it
+# Exit the script immediately if any command exits with a non-zero status (for debugging purposes)
+set -x
+set -e
+
+# Create a Kubernetes secret named 'postgres-secrets' with the specified literals for host and password
+kubectl create secret generic postgres-secrets \
+  --from-literal=host=postgres-postgresql.default.svc.cluster.local \
+  --from-literal=password=postgres
+
+# Apply the PostgreSQL deployment configuration from the specified YAML file
+kubectl apply -f scripts/test/postgres-deployment.yaml
+
+# Build the Docker image with tag 'dbt-jaffle-shop:1.0.0' using the specified Dockerfile
+cd dev && docker build --progress=plain --no-cache -t dbt-jaffle-shop:1.0.0 -f Dockerfile.postgres_profile_docker_k8s .
+
+# Load the Docker image into the local KIND cluster
+kind load docker-image dbt-jaffle-shop:1.0.0
+
+# Retrieve the name of the PostgreSQL pod using the label selector 'app=postgres'
+# The output is filtered to get the first pod's name
+POD_NAME=$(kubectl get pods -n default -l app=postgres -o jsonpath='{.items[0].metadata.name}')
+
+# Print the name of the PostgreSQL pod
+echo "$POD_NAME"
+
+# Forward port 5432 from the PostgreSQL pod to the local machine's port 5432
+# This allows local access to the PostgreSQL instance running in the pod
+kubectl port-forward --namespace default "$POD_NAME" 5432:5432 &
+
+# List all pods in the default namespace to verify the status of pods
+kubectl get pod

scripts/test/performance.sh

@@ -2,4 +2,5 @@ pytest -vv \
     -s \
     -m 'perf' \
     --ignore=tests/test_example_dags.py \
-    --ignore=tests/test_example_dags_no_connections.py
+    --ignore=tests/test_example_dags_no_connections.py \
+    --ignore=tests/test_example_k8s_dags.py