VFagundes commented on a change in pull request #5268: [AIRFLOW-4477] Add
models methods to BQHook
URL: https://github.com/apache/airflow/pull/5268#discussion_r284322025
##########
File path: airflow/contrib/hooks/bigquery_hook.py
##########
@@ -1776,6 +1776,178 @@ def insert_all(self, project_id, dataset_id, table_id,
'BigQuery job failed. Error was: {}'.format(err.content)
)
+ def list_model(self, project_id, dataset_id):
+ """
+ Method to list all models in the specified dataset.
+
+ .. seealso::
+ For more information, see:
+
https://cloud.google.com/bigquery/docs/reference/rest/v2/models/list
+
+ :param project_id: The project ID of the models to list.
+ :type project_id: str
+ :param dataset_id: The dataset ID of the models to list.
+ :type dataset_id: str
+ """
+ project_id = project_id if project_id else self.project_id
+
+ self.log.info('Listing BigQuery models from project: %s Dataset:%s',
+ project_id, dataset_id)
+
+ try:
+ models = self.service.models().list(
+ projectId=project_id,
+ datasetId=dataset_id).execute(num_retries=self.num_retries)
+ self.log.info('BigQuery models listed successfully: In project %s '
+ 'Dataset %s', project_id, dataset_id)
+ return models
+ except HttpError as err:
+ raise AirflowException(
+ 'BigQuery models failed to list. Error was:
{}'.format(err.content)
+ )
+
+ def get_model(self, project_id, dataset_id, model_id):
+ """
+ Method to get the specified model resource by model_id.
+
+ .. seealso::
+ For more information, see:
+ https://cloud.google.com/bigquery/docs/reference/rest/v2/models/get
+
+ :param project_id: The project ID of the requested model.
+ :type project_id: str
+ :param dataset_id: The dataset ID of the requested model.
+ :type dataset_id: str
+ :param model_id: The model ID of the requested model.
+ :type model_id: str
+ """
+ if model_id is None:
+ raise ValueError('model_id cannot be None.')
+
+ project_id = project_id if project_id else self.project_id
+
+ self.log.info('Retrieving BigQuery model %s from project: %s
Dataset:%s',
+ model_id, project_id, dataset_id)
+
+ try:
+ model = self.service.models().get(
Review comment:
how about instantiate the model outside of your `try/except`, to isolate
your model data from your execution, then just leave the execute instruction
inside of it?
something like:
```
model = self.service.models().get( projectId=project_id,
datasetId=dataset_id,
modelId=model_id)
try:
big_query_model = model.execute(num_retiries=self.num_retries)
self.log.info('BigQuery model retrieved successfully: In project %s '
'Dataset %s Model %s', project_id, dataset_id,
model_id)
return big_query_model
except HttpError as err:
raise AirflowException(
'BigQuery models failed to retrieve. Error was:
{}'.format(err.content)
)
```
what do you think? :)
----------------------------------------------------------------
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
For queries about this service, please contact Infrastructure at:
[email protected]
With regards,
Apache Git Services
