From d4a0e20b17fc30d9b7f5e28c712b4cfea9d2b058 Mon Sep 17 00:00:00 2001
From: ismail simsek <6005685+ismailsimsek@users.noreply.github.com>
Date: Tue, 12 Nov 2024 20:06:42 +0100
Subject: [PATCH] Use __main__.py (#27)

---
 .github/workflows/tests.yml                   |  6 ++--
 opendbt/__init__.py                           | 11 -------
 opendbt/__main__.py                           | 13 ++++++++
 setup.py                                      |  4 +--
 .../models/example/my_first_dbt_model.sql     |  2 +-
 tests/test_opendbt_airflow.py                 | 33 +++++++++++++++++++
 6 files changed, 53 insertions(+), 16 deletions(-)
 create mode 100644 opendbt/__main__.py
 create mode 100644 tests/test_opendbt_airflow.py

diff --git a/.github/workflows/tests.yml b/.github/workflows/tests.yml
index 2f9ccf9..fc7c795 100644
--- a/.github/workflows/tests.yml
+++ b/.github/workflows/tests.yml
@@ -31,9 +31,11 @@ jobs:
           dbt-version: ${{ matrix.dbt-version }}
       - name: Build & Install
         run: |
-          pip install -q .[test]
           pip install -q coverage pylint
-          pip install -q dbt-core==${{ matrix.dbt-version }}.* dbt-duckdb==${{ matrix.dbt-version }}.* --force-reinstall
+          pip install -q dbt-core==${{ matrix.dbt-version }}.* dbt-duckdb==${{ matrix.dbt-version }}.* --force-reinstall --upgrade
+          # FIX for protobuf issue: https://github.com/dbt-labs/dbt-core/issues/9759
+          pip install -q "apache-airflow" "protobuf>=4.25.3,<5.0.0" "opentelemetry-proto<1.28.0" --prefer-binary
+          pip install -q .[test] --prefer-binary
           python --version
           python -c "from dbt.version import get_installed_version as get_dbt_version;print(f'dbt version={get_dbt_version()}')"
           python -m compileall -f opendbt setup.py
diff --git a/opendbt/__init__.py b/opendbt/__init__.py
index d244da0..16d5440 100644
--- a/opendbt/__init__.py
+++ b/opendbt/__init__.py
@@ -1,4 +1,3 @@
-import argparse
 import logging
 import os
 import sys
@@ -117,13 +116,3 @@ def manifest(self, partial_parse=True, no_write_manifest=True) -> Manifest:
     def generate_docs(self, args: list = None):
         _args = ["generate"] + args if args else []
         self.run(command="docs", args=_args)
-
-
-def main():
-    p = argparse.ArgumentParser()
-    _, args = p.parse_known_args()
-    OpenDbtCli.run(args=args)
-
-
-if __name__ == "__main__":
-    main()
diff --git a/opendbt/__main__.py b/opendbt/__main__.py
new file mode 100644
index 0000000..844bf9e
--- /dev/null
+++ b/opendbt/__main__.py
@@ -0,0 +1,13 @@
+import argparse
+
+from opendbt import OpenDbtCli
+
+
+def main():
+    p = argparse.ArgumentParser()
+    _, args = p.parse_known_args()
+    OpenDbtCli.run(args=args)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/setup.py b/setup.py
index 57b1db1..2eb06a5 100644
--- a/setup.py
+++ b/setup.py
@@ -10,7 +10,7 @@
     name='opendbt',
     entry_points={
         'console_scripts': [
-            'opendbt = opendbt:main',
+            'opendbt = opendbt.__main__:main',
         ],
     },
     version="0.7.0",
@@ -27,7 +27,7 @@
     install_requires=["dbt-duckdb>=1.6", "sqlfluff", "sqlfluff-templater-dbt"],
     extras_require={
         "airflow": ["apache-airflow"],
-        "test": ["testcontainers>=3.7,<4.9"],
+        "test": ["testcontainers>=3.7,<4.9", "apache-airflow"],
     },
     python_requires='>=3.8'
 )
diff --git a/tests/resources/dbttest/models/example/my_first_dbt_model.sql b/tests/resources/dbttest/models/example/my_first_dbt_model.sql
index f31a12d..c7dfb32 100644
--- a/tests/resources/dbttest/models/example/my_first_dbt_model.sql
+++ b/tests/resources/dbttest/models/example/my_first_dbt_model.sql
@@ -13,7 +13,7 @@ with source_data as (
 
     select 1 as id
     union all
-    select null as id
+    select 2 as id
 
 )
 
diff --git a/tests/test_opendbt_airflow.py b/tests/test_opendbt_airflow.py
new file mode 100644
index 0000000..428045b
--- /dev/null
+++ b/tests/test_opendbt_airflow.py
@@ -0,0 +1,33 @@
+from pathlib import Path
+from unittest import TestCase
+
+from airflow import DAG
+from airflow.operators.empty import EmptyOperator
+from airflow.utils.dates import days_ago
+
+from opendbt.airflow import OpenDbtAirflowProject
+
+
+class TestOpenDbtProject(TestCase):
+    RESOURCES_DIR = Path(__file__).parent.joinpath("resources")
+    DBTTEST_DIR = RESOURCES_DIR.joinpath("dbttest")
+    TEST_DAG = DAG(
+        dag_id='dbt_test_workflow',
+        schedule_interval=None,
+        start_date=days_ago(3),
+        catchup=False,
+        max_active_runs=1
+    )
+
+    def test_run_dbt_as_airflow_task(self):
+        with self.TEST_DAG as dag:
+            # load dbt jobs to airflow dag
+            start = EmptyOperator(task_id="start")
+            end = EmptyOperator(task_id="end")
+            p = OpenDbtAirflowProject(project_dir=self.DBTTEST_DIR, profiles_dir=self.DBTTEST_DIR, target='dev')
+            p.load_dbt_tasks(dag=dag, start_node=start, end_node=end, include_singular_tests=True,
+                             include_dbt_seeds=True)
+
+            for j in dag.tasks:
+                # run all dbt tasks trough airflow operator
+                j.execute({})