How to pass a database connection into Airflow KubernetesPodOperator

user_6516240 0
Viewed 753

I'm having a confusion with KubernetesPodOperator from Airflow, and I'm wondering how to pass the load_users_into_table() function that it has a conn_id parameter stored in connection of Airflow in the Pod ?

In the official doc proposes to put the conn_id in Secret but I don't understand how can I pass it in my function load_users_into_table() after that.

https://airflow.apache.org/docs/stable/kubernetes.html

the function (task) to be executed in the pod:

def load_users_into_table(postgres_hook, schema, path):
  gdf = read_csv(path)
  gdf.to_sql('users', con=postgres_hook.get_sqlalchemy_engine(), schema=schema)

the dag:

_pg_hook = PostgresHook(postgres_conn_id = _conn_id)

with dag: 
test = KubernetesPodOperator(
        namespace=namespace,
        image=image_name,
        cmds=["python", "-c"],
        arguments=[load_users_into_table],
        labels={"dag-id": dag.dag_id},
        name="airflow-test-pod",
        task_id="task-1",
        is_delete_operator_pod=True,
        in_cluster=in_cluster,
        get_logs=True,
        config_file=config_file,
        executor_config={
            "KubernetesExecutor": {"request_memory": "512Mi",
                                   "limit_memory": "1024Mi",
                                   "request_cpu": "1",
                                   "limit_cpu": "2"}
        }
    )
kubernetes-pod airflow kubernetes
1 Answers

Assuming you want to run with K8sPodOperator, you can use argparse and add arguments to the docker cmd. Something in these lines should do the job:

import argparse
    

def f(arg):
    print(arg)


parser = argparse.ArgumentParser()
parser.add_argument('--foo', help='foo help')
args = parser.parse_args()
    

if __name__ == '__main__':
    f(args.foo)

Dockerfile:

FROM python:3
COPY main.py main.py
CMD ["python", "main.py", "--foo", "somebar"]

There are other ways to solve this such as using secrets, configMaps or even Airflow Variables, but this should get you moving forward.

user_41391120
Related