"AttributeError: 'GoogleBaseHook' object has no attribute 'get_records'" Error while trying to run Custom SqlSensor for BigQuery

user_18153304 0
Viewed 243

I am migrating some DAG's from "Cloud Composer-1 Airflow-1" to Cloud "Composer-2 Airflow-2".

I have created a custom module called BigQuerySqlSensor acting like a SqlSensor for BigQuery.

I have done this by creating another custom module called MySqlSensor because Airflow has removed some of the supported attributes for BaseHook and recommended to use DbApiHook instead.

I have basically used the same source code of SqlSensor (can be found here) and replaced BaseHook module with DbApiHook.

Then I used the MySqlSensor to create custom BigQuerySqlSensor like below:

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

The task where I am calling the BigQuerySqlSensor is:

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

However, I am getting an error such as:

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

I am not sure where GoogleBaseHook is being picked up. I am guessing it is something to do with the connection in Airflow.

In my previous Cloud Composer instance, I am running the same code with same configurations and connections. I am out of ideas at this point. Any ideas would be appreciated!

google-cloud-composer airflow google-bigquery sensors
1 Answers

The problem is you are using GoogleBaseHook instead of BigQueryHook (probably your connection type is google_cloud_platform).

It seems it can be solved by defining gcpbigquery connection - but base SqlSensor class checks type of hook and doesn't allow it (it looks like airflow bug - and it seems it is already solved on airflow main branch)

I solved it by defining my own sensor class this way:

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

(conn_id is identifier of google_cloud_platform connection). This way you can also add additional hook parameters - it seems at least 'location' is required.

It works for me on Airflow 2.2.3

user_293710
Related