if __name__ == "__main__":
    dag.test()

from airflow import DAG
from airflow.operators.python import PythonOperator
from datetime import datetime, timedelta
from pathlib import Path
from kedro.framework.startup import bootstrap_project
from kedro.framework.session import KedroSession
from kedro.framework.project import configure_project

# Define default_args dictionary to pass to the DAG
default_args = {
    'owner': 'airflow',
    'depends_on_past': False,
    'start_date': datetime(2024, 1, 1),
    'email_on_failure': False,
    'email_on_retry': False,
    'retries': 1,
    'retry_delay': timedelta(minutes=5)
}

# Instantiate a DAG
dag = DAG(
    'simple_python_operator_example',
    default_args=default_args,
    description='A simple example DAG with two PythonOperators',
    schedule_interval='@hourly',  # Set to run once an hour
    catchup=False
)


# Define Python functions for the two tasks
def task_one_function():
    print("Executing Task One-------")
    env = "airflow"
    pipeline_name = "__default__"
    project_path = Path.cwd()

    bootstrap_project(Path.cwd())
    with KedroSession.create(project_path=project_path,
                             env=env) as session:
        session.run(pipeline_name)


def task_two_function():
    print("Executing Task Two-------")
    # Add your Task Two logic here


# Instantiate PythonOperators for the two tasks
task_one = PythonOperator(
    task_id='task_one',
    python_callable=task_one_function,
    dag=dag,
)

task_two = PythonOperator(
    task_id='task_two',
    python_callable=task_two_function,
    dag=dag,
)

# Set the task dependencies
task_one >> task_two

if __name__ == "__main__":
    dag.test()

KedroOperator

class KedroOperator(BaseOperator):
    @apply_defaults
    def __init__(
        self,
        package_name: str,
        pipeline_name: str,
        node_name: str,
        project_path: str | Path,
        env: str,
        *args, **kwargs
    ) -> None:
        super().__init__(*args, **kwargs)
        self.package_name = package_name
        self.pipeline_name = pipeline_name
        self.node_name = node_name
        self.project_path = project_path
        self.env = env

    def execute(self, context):
        configure_project(self.package_name)
        with KedroSession.create(project_path=self.project_path,
                                 env=self.env) as session:
            session.run(self.pipeline_name, node_names=[self.node_name])

PythonOperator

configure_project(self.package_name)
        with KedroSession.create(project_path=self.project_path,
                                 env=self.env) as session:
            session.run(self.pipeline_name, node_names=[self.node_name])