“AttributeError: 'GoogleBaseHook' object 没有属性 'get_records'” 尝试为 BigQuery 运行自定义 SqlSensor 时出错

Question

我正在将一些 DAG 从“Cloud Composer-1 Airflow-1”迁移到云“Composer-2 Airflow-2”。

我创建了一个名为BigQuerySqlSensor的自定义模块，其作用类似于 BigQuery 的SqlSensor 。

我通过创建另一个名为MySqlSensor的自定义模块来完成此操作，因为 Airflow 删除了BaseHook的一些支持属性，并建议改用DbApiHook 。

我基本上使用了与 SqlSensor 相同的源代码（ 可以在此处找到）并将BaseHook模块替换为DbApiHook 。

然后我使用MySqlSensor创建自定义BigQuerySqlSensor ，如下所示：

class BigQuerySqlSensor(MySqlSensor):
    def _get_hook(self):
        hook = super()._get_hook()
        hook.use_legacy_sql = False
        return hook

我调用BigQuerySqlSensor的任务是：

task = BigQuerySqlSensor(
        task_id='my_task',
        conn_id='bigquery_default',
        timeout= <Some Numbers>
        poke_interval= <Some Numbers>
        mode='reschedule',
        sql='my_sql.sql'
    )

但是，我收到如下错误：

        records = hook.get_records(self.sql, self.parameters)
AttributeError: 'GoogleBaseHook' object has no attribute 'get_records'

我不确定从哪里获取GoogleBaseHook 。 我猜这与 Airflow 中的连接有关。

在我之前的 Cloud Composer 实例中，我使用相同的配置和连接运行相同的代码。 在这一点上我没有想法。 任何想法，将不胜感激！

Answer 1

问题是您使用的是GoogleBaseHook而不是 BigQueryHook （可能您的连接类型是google_cloud_platform ）。

似乎可以通过定义gcpbigquery连接来解决 - 但基础SqlSensor class 检查挂钩类型并且不允许它（它看起来像 airflow 错误 - 似乎它已经在 airflow main分支上解决）

我通过这样定义我自己的传感器 class 解决了这个问题：

class BigQuerySqlSensor(SqlSensor):
    def _get_hook(self):
        return BigQueryHook(gcp_conn_id=self.conn_id, use_legacy_sql=False, location="us")

（conn_id 是google_cloud_platform连接的标识符）。 这样您还可以添加其他挂钩参数 - 似乎至少需要“位置”。

它适用于 Airflow 2.2.3