dbt Cloud & airflow: best of friends

.github/workflows/terraform_checkov_security.yml

@@ -21,7 +21,7 @@ jobs:
           pip install checkov
       - name: Security Test-terragrunt_infrastructure_modules, Skip Specific Tests
         run: |
-          checkov -d terragrunt_infrastructure_modules --skip-check CKV_GCP_26,CKV_GCP_32,CKV_GCP_38
+          checkov -d terragrunt_infrastructure_modules --skip-check CKV_GCP_26,CKV_GCP_32,CKV_GCP_38,CKV2_GCP_3
       - name: Security Test-terraform_simple_setup, Skip Specific Tests
         run: |
-          checkov -d terraform_simple_setup --skip-check CKV_GCP_26,CKV_GCP_32,CKV_GCP_38
+          checkov -d terraform_simple_setup --skip-check CKV_GCP_26,CKV_GCP_32,CKV_GCP_38,CKV2_GCP_3

README.md

@@ -299,7 +299,7 @@ extraVolumes: # this will create the volume from the directory
 # leave this terminal open to sustain airflow webserver
 # Set of environment variables
 export ENV="dev"
-export PROJECT_ID="airflow-demo-build"
+export PROJECT_ID="dbt-demos-sung"
 export DOCKER_DBT_IMG="gcr.io/$PROJECT_ID/dbt_docker:$ENV-latest"
 
 source deploy_local_desktop_airflow.sh
@@ -330,7 +330,7 @@ kubectl port-forward --namespace airflow $POD_NAME 8080:8080
 
 ```bash
 # start a remote shell in the airflow worker for ad hoc operations or to run pytests
-kubectl exec -it airflow-worker-0 -- /bin/bash
+kubectl exec -it airflow-local-desktop-worker-0 -- /bin/bash
 ```
 
 - Airflow worker remote shell examples

custom-setup.yaml

@@ -6,7 +6,7 @@ airflow:
   ##
   image:
     repository: apache/airflow
-    tag: 1.10.10-python3.6
+    tag: 2.1.0-python3.8
     ## values: Always or IfNotPresent
     pullPolicy: IfNotPresent
     pullSecret: ""

dags/airflow_utils.py

@@ -5,17 +5,18 @@
 
 # TODO(developer): update for your specific settings
 # GIT_REPO = "git@github.com:sungchun12/airflow-toolkit.git" #placeholder ssh git repo
-GIT_REPO = "github_sungchun12_airflow-toolkit"
-PROJECT_ID = "airflow-demo-build"
+GIT_REPO = "https://github.com/sungchun12/airflow-toolkit.git"
+PROJECT_ID = "dbt-demos-sung"
 DBT_IMAGE = f"gcr.io/{PROJECT_ID}/dbt_docker:dev-latest"
 
 env = os.environ.copy()
 DEPLOYMENT_SETUP = env["DEPLOYMENT_SETUP"]
-GIT_BRANCH = "master" # TODO: make this an env var: env["GIT_BRANCH"]
+GIT_BRANCH = "master"  # TOsDO: make this an env var: env["GIT_BRANCH"]
+
 
 def get_secret(project_name, secret_name):
     """
-        Returns the value of a secret in Secret Manager for use in DAGs
+    Returns the value of a secret in Secret Manager for use in DAGs
     """
     secrets = secretmanager.SecretManagerServiceClient()
     secret_value = (
@@ -85,11 +86,13 @@ def set_google_app_credentials(deployment_setup, dag_name):
 git_clone_cmds = f"""
     /entrypoint.sh &&
     gcloud auth activate-service-account --key-file=account.json &&
-    gcloud source repos clone {GIT_REPO} --project={PROJECT_ID}"""
+    git clone {GIT_REPO} """
 
 dbt_setup_cmds = f"""
     {git_clone_cmds} &&
-    cd {GIT_REPO}/dbt_bigquery_example &&
+    pwd &&
+    ls &&
+    cd airflow-toolkit/dbt_bigquery_example &&
     git checkout {GIT_BRANCH} &&
     export PROJECT_ID={PROJECT_ID} &&
     export DBT_PROFILES_DIR=$(pwd) &&

dags/dbt_cloud_utils.py

@@ -0,0 +1,130 @@
+import time
+
+import requests
+from dataclasses import dataclass
+from airflow.models import Variable
+
+# TODO: MANUALLY create a dbt Cloud job: https://docs.getdbt.com/docs/dbt-cloud/cloud-quickstart#create-a-new-job
+# Example dbt Cloud job URL
+# https://cloud.getdbt.com/#/accounts/4238/projects/12220/jobs/12389/
+@dataclass(frozen=True)  # make attributes immutable
+class dbt_cloud_job_vars:
+    """Basic dbt Cloud job configurations. Set attributes based on the ids within dbt Cloud URL"""
+
+    # add type hints
+    account_id: int
+    project_id: int
+    job_id: int
+    cause: str
+    dbt_cloud_api_key: str = Variable.get(
+        "dbt_cloud_api_key"
+    )  # TODO: manually set this in the airflow variables UI
+
+
+@dataclass(frozen=True)
+class dbt_job_run_status:
+    """Different dbt Cloud API status responses in integer format"""
+
+    QUEUED: int = 1
+    STARTING: int = 2
+    RUNNING: int = 3
+    SUCCESS: int = 10
+    ERROR: int = 20
+    CANCELLED: int = 30
+
+
+class dbt_cloud_job_runner(dbt_cloud_job_vars, dbt_job_run_status):
+    """Utility to run dbt Cloud jobs.
+
+    Inherits dbt_cloud_job_vars(dataclass) and dbt_job_run_status(dataclass)
+
+    Parameters
+    ----------
+        account_id(int): dbt Cloud account id
+        project_id(int): dbt Cloud project id
+        job_id(int): dbt Cloud job id
+        cause(str): dbt Cloud cause(ex: name of DAG)
+        dbt_cloud_api_key(str)(OPTIONAL): dbt Cloud api key to authorize programmatically interacting with dbt Cloud
+    """
+
+    def _trigger_job(self) -> int:
+        """Trigger the dbt Cloud job asynchronously.
+
+        Verifies dbt_cloud_job_vars match response payload from dbt Cloud api.
+
+        Returns
+        ----------
+            job_run_id(int): specific dbt Cloud job run id invoked
+        """
+        url = f"https://cloud.getdbt.com/api/v2/accounts/{self.account_id}/jobs/{self.job_id}/run/"
+        headers = {"Authorization": f"Token {self.dbt_cloud_api_key}"}
+        res = requests.post(
+            url=url,
+            headers=headers,
+            data={
+                "cause": f"{self.cause}",  # name of the python file invoking this
+            },
+        )
+
+        try:
+            res.raise_for_status()
+        except:
+            print(f"API token (last four): ...{self.dbt_cloud_api_key[-4:]}")
+            raise
+
+        response_payload = res.json()
+        # Verify the dbt Cloud job matches the arguments passed
+        assert self.account_id == response_payload["data"]["account_id"]
+        assert self.project_id == response_payload["data"]["project_id"]
+        assert self.job_id == response_payload["data"]["job_definition_id"]
+
+        job_run_id = response_payload["data"]["id"]
+        return job_run_id
+
+    # to be used in a while loop to check on job status
+    def _get_job_run_status(self, job_run_id) -> int:
+        """Trigger the dbt Cloud job asynchronously.
+
+        Verifies dbt_cloud_job_vars match response payload from dbt Cloud api.
+
+        Parameters
+        ----------
+            job_run_id(int): specific job run id invoked
+
+        Returns
+        ----------
+            job_run_status(int): status of the dbt Cloud job run
+        """
+        url = f"https://cloud.getdbt.com/api/v2/accounts/{self.account_id}/runs/{job_run_id}/"
+        headers = {"Authorization": f"Token {self.dbt_cloud_api_key}"}
+        res = requests.get(url=url, headers=headers)
+
+        res.raise_for_status()
+        response_payload = res.json()
+        job_run_status = response_payload["data"]["status"]
+        return job_run_status
+
+    # main function operator to trigger the job and a while loop to wait for success or error
+    def run_job(self) -> None:
+        """Main handler method to run the dbt Cloud job and track the job run status"""
+        job_run_id = self._trigger_job()
+        print(f"job_run_id = {job_run_id}")
+
+        visit_url = f"https://cloud.getdbt.com/#/accounts/{self.account_id}/projects/{self.project_id}/runs/{job_run_id}/"
+        print(f"Check the dbt Cloud job status! Visit URL:{visit_url}")
+
+        while True:
+            time.sleep(1)  # make an api call every 1 second
+
+            job_run_status = self._get_job_run_status(job_run_id)
+
+            print(f"job_run_status = {job_run_status}")
+
+            if job_run_status == dbt_job_run_status.SUCCESS:
+                print(f"Success! Visit URL: {visit_url}")
+                break
+            elif (
+                job_run_status == dbt_job_run_status.ERROR
+                or job_run_status == dbt_job_run_status.CANCELLED
+            ):
+                raise Exception(f"Failure! Visit URL: {visit_url}")

dags/examples/add_gcp_connections.py

@@ -38,7 +38,7 @@
 
 # TODO(developer): update for your specific naming conventions
 CONN_PARAMS_DICT = {
-    "gcp_project": "airflow-demo-build",
+    "gcp_project": "dbt-demos-sung",
     "gcp_conn_id": "my_gcp_connection",
     "gcr_conn_id": "gcr_docker_connection",
     "secret_name": "airflow-conn-secret",

dags/examples/bigquery_connection_check.py

@@ -16,7 +16,7 @@
 # TODO(developer): update for your specific settings
 TASK_PARAMS_DICT = {
     "dataset_id": "dbt_bq_example",
-    "project_id": "airflow-demo-build",
+    "project_id": "dbt-demos-sung",
     "gcp_conn_id": "my_gcp_connection",
 }
 

dags/examples/dbt_cloud_example.py

@@ -0,0 +1,45 @@
+from airflow import DAG
+from datetime import datetime
+from airflow.operators.dummy import DummyOperator
+from airflow.operators.python_operator import PythonOperator
+from dbt_cloud_utils import dbt_cloud_job_runner
+
+
+dag_file_name = __file__
+
+# TODO: MANUALLY create a dbt Cloud job: https://docs.getdbt.com/docs/dbt-cloud/cloud-quickstart#create-a-new-job
+# Example dbt Cloud job URL
+# https://cloud.getdbt.com/#/accounts/4238/projects/12220/jobs/12389/
+# example dbt Cloud job config
+dbt_cloud_job_runner_config = dbt_cloud_job_runner(
+    account_id=4238, project_id=12220, job_id=12389, cause=dag_file_name
+)
+
+
+default_args = {
+    "owner": "airflow",
+    "email": ["airflow@example.com"],
+    "depends_on_past": False,
+    "start_date": datetime(2001, 1, 1),
+    "email_on_failure": False,
+    "email_on_retry": False,
+    "retries": 0,
+    "priority_weight": 1000,
+}
+
+
+with DAG(
+    "dbt_cloud_example", default_args=default_args, schedule_interval="@once"
+) as dag:
+    # have a separate extract and load process(think: FivetranOperator and/or custom gcs load to bigquery tasks)
+    extract = DummyOperator(task_id="extract")
+    load = DummyOperator(task_id="load")
+
+    # Single task to execute dbt Cloud job and track status over time
+    transform = PythonOperator(
+        task_id="run-dbt-cloud-job",
+        python_callable=dbt_cloud_job_runner_config.run_job,
+        provide_context=True,
+    )
+
+    extract >> load >> transform

dags/examples/kubernetes_sample.py

@@ -1,5 +1,6 @@
 from airflow import DAG
-from datetime import datetime, timedelta
+from datetime import datetime
+from airflow.operators.bash_operator import BashOperator
 from airflow.contrib.operators.kubernetes_pod_operator import KubernetesPodOperator
 from airflow.operators.dummy_operator import DummyOperator
 
@@ -17,7 +18,11 @@
     # "retry_delay": timedelta(minutes=5),
 }
 
-dag = DAG("kubernetes_sample", default_args=default_args, schedule_interval="@once",)
+dag = DAG(
+    "kubernetes_sample",
+    default_args=default_args,
+    schedule_interval="@once",
+)
 
 
 start = DummyOperator(task_id="run_this_first", dag=dag)
@@ -40,11 +45,12 @@
     namespace="default",
     image="ubuntu:16.04",
     cmds=["/bin/bash", "-cx"],
-    arguments=["echo hello world"],
+    arguments=["mkdir -p /airflow/xcom/;echo '[1,2,3,4]' > /airflow/xcom/return.json"],
     labels={"foo": "bar"},
-    name="fail",
+    name="passing-bash-task",
     task_id="passing-bash-task",
     get_logs=True,
+    do_xcom_push=True,
     dag=dag,
 )
 
@@ -62,4 +68,26 @@
     dag=dag,
 )
 
+bash_task_xcom_result = BashOperator(
+    bash_command="echo \"{{ task_instance.xcom_pull('passing-bash-task')[0] }}\"",
+    task_id="bash_task_xcom_result",
+    dag=dag,
+)
+
+# xcom pull follows a different syntax compared to the BashOperator
+pod_task_xcom_result = KubernetesPodOperator(
+    namespace="default",
+    image="ubuntu:16.04",
+    cmds=["/bin/bash", "-c"],
+    arguments=["echo \"{{ task_instance.xcom_pull(task_ids='passing-bash-task', key='return_value')[0] }}\""],
+    labels={"foo": "bar"},
+    name="pod_task_xcom_result",
+    task_id="pod_task_xcom_result",
+    get_logs=True,
+    dag=dag,
+)
+
+
 start >> [passing_python, passing_bash, private_gcr_passing]
+
+passing_bash >> [bash_task_xcom_result, pod_task_xcom_result]

dags/kube_secrets.py

@@ -1,4 +1,4 @@
-from airflow.contrib.kubernetes.secret import Secret
+from airflow.kubernetes.secret import Secret
 
 
 # GCP service account for dbt operations with BigQuery
@@ -17,5 +17,8 @@
 
 # This is included as a placeholder based on the note in `airflow_utils.py`
 GIT_SECRET_ID_RSA_PRIVATE = Secret(
-    deploy_type="volume", deploy_target="/dbt/.ssh/", secret="ssh-key-secret", key="id_rsa",
+    deploy_type="volume",
+    deploy_target="/dbt/.ssh/",
+    secret="ssh-key-secret",
+    key="id_rsa",
 )

deploy_local_desktop_airflow.sh

@@ -68,11 +68,11 @@ echo "***********************"
 # install airflow helm chart
 # https://helm.sh/docs/helm/helm_install/
 helm install \
-airflow-local-desktop \
-airflow-stable/airflow \
---version 7.16.0 \
---namespace "airflow" \
---values ./custom-setup.yaml
+  airflow-local-desktop \
+  airflow-stable/airflow \
+  --version 8.3.1 \
+  --namespace "airflow" \
+  --values ./custom-setup.yaml
 
 echo "***********************"
 echo "Wait for the Kubernetes Cluster to settle"

utils/checkov/requirements.txt

@@ -1 +1 @@
-checkov==1.0.501
+checkov==2.0.259