Как запустить функцию asyn c в Airflow? - PullRequest
Новая
       6

Как запустить функцию asyn c в Airflow?

3 голосов
/ 17 февраля 2020

Я пишу задачу воздушного потока, чтобы прочитать большой CSV и сохранить его в postgresql базе данных. Я нашел этот пакет asyncpg, который имеет функцию копирования, которая работает намного быстрее, чем любые другие пакеты. Тем не менее, это asyn c, и я не знаю, как включить его в Airflow. Вот пример кода: 

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime, timedelta
from pandas import DataFrame
import asyncpg

async def to_sql(dataframe, table_name, schema_name='public', timeout=None, truncate=False):
    connection = await asyncpg.connect(user='postgres', host='host.docker.internal', database='quantaxis', password='123456')
    result = await connection.copy_records_to_table(
        table_name,
        records=dataframe.values.tolist(),
        columns=shared_columns,
        schema_name=schema_name,
        timeout=timeout)
    await connection.close()
    return result


default_args = {
    'owner': 'Airflow',
    'depends_on_past': False,
    'start_date': datetime(2020, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=1),
}

dag = DAG('pythonexp2123', default_args=default_args, schedule_interval=timedelta(days=1))

async def save_file_to_database(ds):
    df = pd.read_csv("data{0}.csv".format(ds))
    r = await to_sql(df, 'test')
    return r

t1 = PythonOperator(
    task_id='pushing_task',
    provide_context=True,
    python_callable=save_file_to_database,
    dag=dag
    )

t1

Когда я запускаю его, он возвращает ошибку: 

Can't Pickle Object <Corountine>

Как я мог изменить функцию, чтобы этот Dag работал? Я все еще хочу использовать пакет asyncpg из-за его скорости.

1 Ответ

1 голос
/ 18 февраля 2020

Вы можете попробовать запустить функцию asyn c в eventl oop, используя asyncio. Если вы используете python 3.7>, вы можете просто вызвать asyncio.run (async_function ())

https://docs.python.org/3/library/asyncio-task.html

from airflow import DAG
from airflow.operators.python_operator import PythonOperator
from datetime import datetime, timedelta
from pandas import DataFrame
import asyncpg
import ayncio

async def to_sql(dataframe, table_name, schema_name='public', timeout=None, truncate=False):
    connection = await asyncpg.connect(user='postgres', host='host.docker.internal', database='quantaxis', password='123456')
    result = await connection.copy_records_to_table(
        table_name,
        records=dataframe.values.tolist(),
        columns=shared_columns,
        schema_name=schema_name,
        timeout=timeout)
    await connection.close()
    return result



default_args = {
    'owner': 'Airflow',
    'depends_on_past': False,
    'start_date': datetime(2020, 1, 1),
    'retries': 1,
    'retry_delay': timedelta(minutes=1),
}

dag = DAG('pythonexp2123', default_args=default_args, schedule_interval=timedelta(days=1))

async def save_file_to_database(ds):
    df = pd.read_csv("data{0}.csv".format(ds))
    r = await to_sql(df, 'test')
    return r

def run_async(ds):
   loop = asyncio.get_event_loop()
   result = loop.run_until_complete(save_file_to_database(ds))
   return result

t1 = PythonOperator(
    task_id='pushing_task',
    provide_context=True,
    python_callable=run_async,
    dag=dag
    )

t1
...