带有提交的气流+熊猫read_sql_query()

时间:2018-12-12 22:42:37

标签: sql pandas commit airflow postgresql-8.2

问题

我可以使用read_sql()将SQL事务提交到数据库吗?

用例和背景

我有一个用例,我想允许用户执行一些预定义的SQL并返回熊猫数据框。在某些情况下,此SQL将需要查询预先填充的表,而在其他情况下,此SQL将执行将写入表的函数,然后将查询该表。 该逻辑当前包含在Airflow DAG的方法内部,以利用PostgresHook来利用Airflow可以访问的数据库连接信息-该方法最终在PythonOperator任务中调用。通过测试,我了解到PostgresHook创建了一个psycopg2连接对象。

代码

from airflow.hooks.postgres_hook import PostgresHook
import pandas as pd 

def create_df(job_id,other_unrelated_inputs):
    conn = job_type_to_connection(job_type) # method that helps choose a database
    sql = open('/sql_files/job_id_{}.sql'.format(job_id)) #chooses arbitrary SQL  
    sql_template = sql.read() 
    hook = PostgresHook(postgres_conn_id=conn) #connection information for alias is predefined elsewhere within Airflow


    try:
        hook_conn_obj = hook.get_conn()
        print(type(hook_conn_obj)) # <class 'psycopg2.extensions.connection'>
        # Runs SQL template with variables, but does not commit. Alternatively, have used hook.get_pandas_df(sql_template)
        df = pd.io.sql.read_sql(sql_template, con = hook_conn_obj) 
    except:
        #catches some errors#
    return df

问题

当前,在执行SQL函数时,此代码会生成一个数据帧,但不会提交在SQL函数中进行的任何数据库更改。例如,更确切地说,如果SQL函数将一行插入表中,则该事务将不会提交,并且该行也不会出现在表中。

尝试

我尝试了一些修复程序,但被卡住了。我的最新工作是更改read_sql使用的psycopg2连接的autocommit属性,以自动提交事务。

我将承认我无法弄清楚连接的属性何时会影响SQL的执行。

我认识到,另一种方法是复制PostgresHook.run()中的某些逻辑以进行提交,然后添加一些代码以将结果推送到数据帧中,但是对于将来的支持来说,使用这些方法似乎更加省事和容易。已创建(如果可能)。

我能找到的最类似的SO问题是this one,但我对与气流无关的解决方案很感兴趣。

编辑

...
    try:
        hook_conn_obj = hook.get_conn()
        print(type(hook_conn_obj)) # <class 'psycopg2.extensions.connection'>
        hook_conn_obj.autocommit = True
        df = pd.io.sql.read_sql(sql_template, con = hook_conn_obj) # Runs SQL template with variables, but does not commit
    except:
        #catches some errors#
    return df

这似乎有效。如果有人对实现此目标的更好方法有任何评论或想法,我仍然有兴趣从讨论中学习。

谢谢!

2 个答案:

答案 0 :(得分:0)

read_sql将不会提交,因为正如该方法名称所暗示的那样,目标是读取数据而不是写入。 pandas是不错的设计选择。这一点很重要,因为它可以防止意外写入,并允许运行程序,读取其作用之类的有趣场景,但没有任何持久性。 read_sql的目的是阅读,而不是写作。直接表达意图是金标准原则。

表达意图的更明确的方法是在execute之前显式地fetchall(带有提交)。但是,由于pandas没有提供简单的方法来读取cursor对象,因此您将失去read_sql提供的便利,而必须自己创建DataFrame。

因此,所有解决方案都很好,通过设置autocommit=True表示您的数据库交互将继续执行,因此不会发生任何意外。读起来有点奇怪,但是如果您将sql_template变量命名为write_then_read_sql或在文档字符串中进行解释,则意图会更清楚。

答案 1 :(得分:0)

我有一个类似的用例-用Pandas将数据加载到SQL Server中,调用繁重的存储过程并写入表,然后将结果集捕获到新的DataFrame中。

我通过使用上下文管理器并显式提交事务来解决它:

# Connect to SQL Server
engine = sqlalchemy.create_engine('db_string')
with engine.connect() as connection:
    # Write dataframe to table with replace
    df.to_sql(name='myTable', con=connection, if_exists='replace')

    with connection.begin() as transaction:
        # Execute verification routine and capture results
        df_processed = pandas.read_sql(sql='exec sproc', con=connection)
        transaction.commit()