From 651337b2e3f050549c4ed9cb1cc77465bfd13eac Mon Sep 17 00:00:00 2001 From: Kevin Loftis Date: Fri, 1 Oct 2021 17:45:42 -0700 Subject: [PATCH 1/2] update bq feature store configs and client configs Signed-off-by: Kevin Loftis --- .../feast/infra/offline_stores/bigquery.py | 17 +++++++++++++---- 1 file changed, 13 insertions(+), 4 deletions(-) diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py index e1e53d4128..80b1eebd4e 100644 --- a/sdk/python/feast/infra/offline_stores/bigquery.py +++ b/sdk/python/feast/infra/offline_stores/bigquery.py @@ -52,6 +52,9 @@ class BigQueryOfflineStoreConfig(FeastConfigBaseModel): project_id: Optional[StrictStr] = None """ (optional) GCP project name used for the BigQuery offline store """ + location: Optional[StrictStr] = None + """ (optional) GCP location name used for the BigQuery offline store """ + class BigQueryOfflineStore(OfflineStore): @staticmethod @@ -79,7 +82,10 @@ def pull_latest_from_table_or_query( timestamp_desc_string = " DESC, ".join(timestamps) + " DESC" field_string = ", ".join(join_key_columns + feature_name_columns + timestamps) - client = _get_bigquery_client(project=config.offline_store.project_id) + client = _get_bigquery_client( + project=config.offline_store.project_id, + location=config.offline_store.location, + ) query = f""" SELECT {field_string} @@ -115,7 +121,10 @@ def get_historical_features( # TODO: Add entity_df validation in order to fail before interacting with BigQuery assert isinstance(config.offline_store, BigQueryOfflineStoreConfig) - client = _get_bigquery_client(project=config.offline_store.project_id) + client = _get_bigquery_client( + project=config.offline_store.project_id, + location=config.offline_store.location, + ) assert isinstance(config.offline_store, BigQueryOfflineStoreConfig) @@ -367,9 +376,9 @@ def _upload_entity_df_and_get_entity_schema( return entity_schema -def _get_bigquery_client(project: Optional[str] = None): +def _get_bigquery_client(project: Optional[str] = None, location: Optional[str] = None): try: - client = bigquery.Client(project=project) + client = bigquery.Client(project=project, location=location) except DefaultCredentialsError as e: raise FeastProviderLoginError( str(e) From 55fd2a1345d631935fd34b0a838ae4251fb9c1e7 Mon Sep 17 00:00:00 2001 From: Kevin Loftis Date: Sat, 2 Oct 2021 12:29:26 -0700 Subject: [PATCH 2/2] add additional docs Signed-off-by: Kevin Loftis --- sdk/python/feast/infra/offline_stores/bigquery.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py index 80b1eebd4e..aa4d43047c 100644 --- a/sdk/python/feast/infra/offline_stores/bigquery.py +++ b/sdk/python/feast/infra/offline_stores/bigquery.py @@ -53,7 +53,11 @@ class BigQueryOfflineStoreConfig(FeastConfigBaseModel): """ (optional) GCP project name used for the BigQuery offline store """ location: Optional[StrictStr] = None - """ (optional) GCP location name used for the BigQuery offline store """ + """ (optional) GCP location name used for the BigQuery offline store. + Examples of location names include ``US``, ``EU``, ``us-central1``, ``us-west4``. + If a location is not specified, the location defaults to the ``US`` multi-regional location. + For more information on BigQuery data locations see: https://cloud.google.com/bigquery/docs/locations + """ class BigQueryOfflineStore(OfflineStore):