Airflow中的PostgresOperator获得超时

Question

我在Postgres中创建了一个函数，它有以下声明：

FUNCTION
        SET statement_timeout TO "3600s"
     SELECT * FROM schema.table_name
        END 
FUNCTION

在Airflow中我使用PostgresOperator来执行此功能，但是[2018-06-01 00:00:01,066] {models.py:1595} ERROR - canceling statement due to statement timeout ，我收到消息[2018-06-01 00:00:01,066] {models.py:1595} ERROR - canceling statement due to statement timeout 。

我看到PostgresOperator使用postgres_hook ，而postgres_hook使用psycopg2作为连接器。

正如我所见，我可以通过cli应用程序超时，而不是数据库中的超时。

我想知道如何解决这个问题？ 我是否需要在Airflow中配置Psycopg，还是可以使用一些环境变量来设置超时以避免此问题？

Answer 1

您可以通过连接上的Airflow extras属性将连接参数传递到psycopg2库。 在撰写本文时，postgres_hook支持以下参数

['sslmode', 'sslcert', 'sslkey','sslrootcert', 'sslcrl', 'application_name', 'keepalives_idle']

为了将statement_timeout参数传递给PostgresHook，您需要覆盖PostgresHook的get_conn以接受您想要的参数。

防爆。 类方法覆盖

class NewPostgresHook(PostgresHook):
    def __init__(self, *args, **kwargs):
        super(NewPostgresHook, self).__init__(*args, **kwargs)

    def get_conn(self):
        conn = self.get_connection(self.postgres_conn_id)
        conn_args = dict(
            host=conn.host,
            user=conn.login,
            password=conn.password,
            dbname=self.schema or conn.schema,
            port=conn.port)
        # check for ssl parameters in conn.extra
        for arg_name, arg_val in conn.extra_dejson.items():
            if arg_name in ['sslmode', 'sslcert', 'sslkey',
                            'sslrootcert', 'sslcrl', 'application_name',
                            'keepalives_idle', 'statement_timeout']:
                conn_args[arg_name] = arg_val

        self.conn = psycopg2.connect(**conn_args)
        return self.conn

然后，您可以以JSON字符串的形式在connection extras字段中指定此参数。

防爆。 连接附加字段中的JSON字符串

{'statement_timeout': '3600s'}