如何在 python 可调用中使用 postgres 连接 id 连接到 postgres

Question

我正在使用 Airflow 的 python 运算符来调用 python 函数。错误发生在 try/ except 块中。

def python_callable_new():
    print("Inside python callable ...")

    import psycopg2

    try:
        print("attempting database connection from python method.. ")
        conn = psycopg2.connect('postgres_defined_connection')
        print("success. ")
    except Exception as error:
        print("failed: ")
        print (error)
    return 'End of callable. '

    

with dag:
    start_task  = DummyOperator(  task_id= "start" )
    stop_task   = DummyOperator(  task_id= "stop"  )
    
    do_python_task = PythonOperator(
        task_id = 'do-py-operation',
        python_callable= python_callable_new,
    )

    extract_source_data = PostgresOperator(
        task_id='extract-cb-source-data',
        postgres_conn_id='postgres_defined_connection',
        sql='./sql_scripts/extract_csv_data.sql'
    )

    # csv_to_postgres

start_task >> do_python_task >> extract_source_data >> stop_task

基本上，我的问题是

• 如何使用“postgres_define_connection”连接到 python 函数内部的 postgres？
• 当我使用 PostgresOperator 时，它连接得很好，如 extract_source_data 任务中所示，但我需要在可调用函数中使用它。
• 出现的错误是 无效的 dsn：连接信息字符串中的“postgres_define_connection”后缺少“=”

（仅供参考 - 我将 postgres_define_connection 存储在使用 sqlalchemy 引擎和 PostgresHook 的单独的connections.py 中）

Answer 1

psycopg2.connect需要连接参数。如果将连接参数格式化为以空格分隔的键/值对，则可以向它们传递单个字符串。这就是为什么它给你错误消息缺少“=”。

请参阅psycopg 文档以获取更多信息。

要连接到 Airflow 中的 Postgres 数据库，只要您已创建连接，就可以利用PostgresHook 。

from airflow.hooks.postgres_hook import PostgresHook
 
def execute_query_with_conn_obj(query):
    hook = PostgresHook(postgres_conn_id='my_connection')
    conn = hook.get_conn()
    cur = conn.cursor()
    cur.execute(query)

def execute_query_with_hook(query):
    hook = PostgresHook(postgres_conn_id='my_connection')
    hook.run(sql=query)

您也可以使用纯 Python 代码来完成此操作。

def execute_query_with_psycopg(query):
    conn_args = dict(
        host='myhost',
        user='admin',
        password='password',
        dbname='my_schema',
        port=5432)
    conn = psycopg2.connect(**conn_args)
    cur = conn.cursor()
    cur.execute(query)

def execute_query_with_psycopg_string(query):
    conn = psycopg2.connect("dbname=test user=postgres password=secret")
    cur = conn.cursor()
    cur.execute(query)