气流：Snowflakeperator：最后查询ID返回XCOM

Question

我有一个带有如下查询的雪花文件，如果我有返回，请在雪花操作员中，以便可以将XCOM传递到下一个任务。

如何仅获取XCOM返回的最后一个查询ID？基本上，我需要将Snowflake的最后查询ID到

SQL文件中的XCOM：

select columns from tableA ;
select last_query_id();

错误：不支持单个API调用中的多个SQL语句；改用每个语句使用一个API调用。

还是有一种方法可以在返回XCOM的查询ID以下获得

代码：

  class LastQueryId(SnowflakeOperator):

    def execute(self, context: Any) -> None:
        """Run query on snowflake"""
        self.log.info('Executing: %s', self.sql)
        hook = SnowflakeHook(snowflake_conn_id=self.snowflake_conn_id,
                             warehouse=self.warehouse, database=self.database,
                             role=self.role, schema=self.schema, authenticator=self.authenticator)
        result = hook.run(self.sql, autocommit=self.autocommit, parameters=self.parameters)
        self.query_ids = hook.query_ids

        if self.do_xcom_push and len(self.query_ids) > 0:
            return self.query_ids[-1]

更新了：我能够使用上述代码获得雪花的查询ID，但是在日志中，我也看到了查询的结果，我如何避免日志中的那些P>

[2022-06-23, 20:43:39 UTC] {cursor.py:696} INFO - query: [SELECT  modifieddate, documentdate...]
[2022-06-23, 20:43:40 UTC] {cursor.py:720} INFO - query execution done
[2022-06-23, 20:43:56 UTC] {snowflake.py:307} INFO - Statement execution info - {'MODIFIEDDATE': datetime.datetime(2022, 6, 23, 11, 42, 34, 233000), 'DOCUMENTDATE': datetime.datetime(2015, 10, 1, 0, 0)...}
[2022-06-23, 20:43:56 UTC] {snowflake.py:307} INFO - Statement execution info - {'MODIFIEDDATE': datetime.datetime(2022, 6, 23, 13, 50, 45, 377000), 'DOCUMENTDATE': datetime.datetime(2021, 7, 1, 0, 0)...}
[2022-06-23, 20:43:56 UTC] {snowflake.py:307} INFO - Statement execution info - {'MODIFIEDDATE': datetime.datetime(2022, 6, 23, 11, 36, 51, 583000), 'DOCUMENTDATE': datetime.datetime(2015, 8, 31, 0, 0)...}
....
....
....
[2022-06-23, 20:43:56 UTC] {snowflake.py:311} INFO - Rows affected: 22116
[2022-06-23, 20:43:56 UTC] {snowflake.py:312} INFO - Snowflake query id: 01a5259b-0501-98f3-0251-830144baa623
[2022-06-23, 20:43:56 UTC] {connection.py:509} INFO - closed

Answer 1

SnowflakeOperator already store the query_ids但它不会将它们推到XCOM。

您可以创建一个自定义操作员：

class MySnowflakeOperator(SnowflakeOperator):

    def execute(self, context: Any) -> None:
        """Run query on snowflake"""
        self.log.info('Executing: %s', self.sql)
        hook = self.get_db_hook()
        execution_info = hook.run(self.sql, autocommit=self.autocommit, parameters=self.parameters)
        self.query_ids = hook.query_ids

        if self.do_xcom_push and len(self.query_ids) > 0:
            return self.query_ids[-1]  # last query_id

如果要维护原始运算符功能，则可以执行：

class MySnowflakeOperator(SnowflakeOperator):

    def execute(self, context: Any) -> None:
        parent_return_value = super().execute(context)
        if self.do_xcom_push and len(self.query_ids) > 0:
            self.xcom_push(
                context,
                key="last_query_id",
                value=self.query_ids[-1],
            )   
        return parent_return_value