From d2d9a042417e5056daa8c5455bc9efea13199c4e Mon Sep 17 00:00:00 2001
From: Bruno Antonellini <antonellinibruno@gmail.com>
Date: Fri, 1 Mar 2024 15:02:04 -0300
Subject: [PATCH] Test sample_script calling on new Operator

---
 orchestrate/dags/coldbore.py                | 41 ---------------------
 orchestrate/dags/sample_script.py           | 27 ++++++++++++++
 orchestrate/python_scripts/sample_script.py | 17 +--------
 orchestrate/some_module/__init__.py         |  0
 orchestrate/some_module/main.py             |  3 ++
 5 files changed, 32 insertions(+), 56 deletions(-)
 delete mode 100644 orchestrate/dags/coldbore.py
 create mode 100644 orchestrate/dags/sample_script.py
 create mode 100644 orchestrate/some_module/__init__.py
 create mode 100644 orchestrate/some_module/main.py

diff --git a/orchestrate/dags/coldbore.py b/orchestrate/dags/coldbore.py
deleted file mode 100644
index 85000347..00000000
--- a/orchestrate/dags/coldbore.py
+++ /dev/null
@@ -1,41 +0,0 @@
-from airflow.decorators import dag
-from kubernetes.client import models as k8s
-from operators.datacoves.bash import DatacovesBashOperator
-
-# Replace with your docker image repo path
-IMAGE_REPO = "python"
-# Replace with your docker image repo tag, or use "latest"
-IMAGE_TAG = "3.8.18-slim-bullseye"
-
-default_args = {
-    "owner": "airflow",
-    "email": "some_user@exanple.com",
-    "email_on_failure": True,
-    "description": "Sample python dag",
-}
-
-TRANSFORM_CONFIG = {
-    "pod_override": k8s.V1Pod(
-        spec=k8s.V1PodSpec(
-            containers=[k8s.V1Container(name="base", image=f"{IMAGE_REPO}:{IMAGE_TAG}")]
-        )
-    ),
-}
-
-
-@dag(
-    default_args={"start_date": "2021-01"},
-    description="k8s_executor",
-    schedule_interval="0 0 1 */12 *",
-    tags=["version_1"],
-    catchup=False,
-)
-def k8s_executor():
-    k8s_executor = DatacovesBashOperator(
-        task_id="k8s_executor",
-        executor_config=TRANSFORM_CONFIG,
-        bash_command=f"echo Hello",
-    )
-
-
-dag = k8s_executor()
diff --git a/orchestrate/dags/sample_script.py b/orchestrate/dags/sample_script.py
new file mode 100644
index 00000000..418874d3
--- /dev/null
+++ b/orchestrate/dags/sample_script.py
@@ -0,0 +1,27 @@
+from airflow.decorators import dag
+from kubernetes.client import models as k8s
+from operators.datacoves.bash import DatacovesBashOperator
+
+default_args = {
+    "owner": "airflow",
+    "email": "some_user@exanple.com",
+    "email_on_failure": True,
+    "description": "Sample python dag",
+}
+
+
+@dag(
+    default_args={"start_date": "2021-01"},
+    description="sample_script",
+    schedule_interval="0 0 1 */12 *",
+    tags=["version_2"],
+    catchup=False,
+)
+def sample_script():
+    sample_script = DatacovesBashOperator(
+        task_id="sample_script",
+        bash_command=f"python orchestrate/python_scripts/sample_script.py",
+    )
+
+
+dag = sample_script()
diff --git a/orchestrate/python_scripts/sample_script.py b/orchestrate/python_scripts/sample_script.py
index 5ac1e831..77ef188f 100644
--- a/orchestrate/python_scripts/sample_script.py
+++ b/orchestrate/python_scripts/sample_script.py
@@ -1,16 +1,3 @@
-import pandas as pd
+from orchestrate.some_module.main import MainRunner
 
-def print_sample_dataframe():
-
-    # Creating a simple DataFrame
-    data = {'Name': ['Alice', 'Bob', 'Charlie'],
-            'Age': [25, 30, 35],
-            'City': ['New York', 'San Francisco', 'Los Angeles']}
-
-    df = pd.DataFrame(data)
-
-    # Displaying the DataFrame
-    print("DataFrame created using Pandas:")
-    print(df)
-
-print_sample_dataframe()
+runner = MainRunner()
diff --git a/orchestrate/some_module/__init__.py b/orchestrate/some_module/__init__.py
new file mode 100644
index 00000000..e69de29b
diff --git a/orchestrate/some_module/main.py b/orchestrate/some_module/main.py
new file mode 100644
index 00000000..36f2be39
--- /dev/null
+++ b/orchestrate/some_module/main.py
@@ -0,0 +1,3 @@
+class MainRunner:
+    def __init__(self) -> None:
+        print("SUCCESS")