From ea775feb40580e26358af9249effafcbe40fcc72 Mon Sep 17 00:00:00 2001
From: Pankaj <pankaj.singh@astronomer.io>
Date: Mon, 22 Jul 2024 16:51:30 +0530
Subject: [PATCH] Fix

---
 dev/Dockerfile.postgres_profile_docker_k8s | 16 ++++
 dev/dags/jaffle_shop_kubernetes.py         | 90 ++++++++++++++++++++++
 scripts/test/kubernetes-test.sh            |  2 +
 3 files changed, 108 insertions(+)
 create mode 100644 dev/Dockerfile.postgres_profile_docker_k8s
 create mode 100644 dev/dags/jaffle_shop_kubernetes.py

diff --git a/dev/Dockerfile.postgres_profile_docker_k8s b/dev/Dockerfile.postgres_profile_docker_k8s
new file mode 100644
index 000000000..929700e35
--- /dev/null
+++ b/dev/Dockerfile.postgres_profile_docker_k8s
@@ -0,0 +1,16 @@
+FROM python:3.9
+
+RUN pip install dbt-postgres==1.3.1 psycopg2==2.9.3 pytz
+
+ENV POSTGRES_DATABASE=postgres
+ENV POSTGRES_HOST=host.docker.internal
+ENV POSTGRES_PASSWORD=<postgres_password>
+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
+
+COPY dags dags
+COPY dags/dbt/jaffle_shop/dbt_project.yml dags/dbt/jaffle_shop/dbt_project.yml
diff --git a/dev/dags/jaffle_shop_kubernetes.py b/dev/dags/jaffle_shop_kubernetes.py
new file mode 100644
index 000000000..5b5d0fdd1
--- /dev/null
+++ b/dev/dags/jaffle_shop_kubernetes.py
@@ -0,0 +1,90 @@
+"""
+## 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 pathlib import Path
+
+from airflow import DAG
+from airflow.providers.cncf.kubernetes.secret import Secret
+from pendulum import datetime
+
+from cosmos import (
+    ProjectConfig,
+    ProfileConfig,
+    ExecutionConfig,
+    ExecutionMode,
+    DbtSeedKubernetesOperator,
+    DbtTaskGroup,
+)
+from cosmos.profiles import PostgresUserPasswordProfileMapping
+
+
+PROJECT_DIR = Path("dags/dbt/jaffle_shop/")
+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:
+    load_seeds = DbtSeedKubernetesOperator(
+        task_id="load_seeds",
+        project_dir=PROJECT_DIR,
+        get_logs=True,
+        schema="public",
+        conn_id="postgres_default",
+        image=DBT_IMAGE,
+        is_delete_operator_pod=False,
+        secrets=[postgres_password_secret, postgres_host_secret],
+    )
+
+    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],
+        },
+    )
+
+    load_seeds >> run_models
diff --git a/scripts/test/kubernetes-test.sh b/scripts/test/kubernetes-test.sh
index 1ad2c1d5a..7c5352a46 100644
--- a/scripts/test/kubernetes-test.sh
+++ b/scripts/test/kubernetes-test.sh
@@ -11,3 +11,5 @@ export POSTGRES_PASSWORD
 sleep 300
 kubectl port-forward --namespace default postgres-postgresql-0  5432:5432 &
 kubectl create secret generic postgres-secrets --from-literal=host=postgres-postgresql.default.svc.cluster.local --from-literal=password=$POSTGRES_PASSWORD
+cd dev && docker build -t dbt-jaffle-shop:1.0.0 -f Dockerfile.postgres_profile_docker_k8s .
+kind load docker-image dbt-jaffle-shop:1.0.0