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'}