fix bigquery

Signed-off-by: Achal Shah <achals@gmail.com>
feast-dev · Jul 7, 2022 · 7fbc119 · 7fbc119
1 parent dcfba45
commit 7fbc119
Show file tree

Hide file tree

Showing 2 changed files with 6 additions and 2 deletions.
diff --git a/sdk/python/feast/infra/offline_stores/bigquery.py b/sdk/python/feast/infra/offline_stores/bigquery.py
@@ -513,9 +513,13 @@ def to_remote_storage(self) -> List[str]:
         )
         extract_job.result()
 
-        storage_client = StorageClient(project=self.config.offline_store.project_id)
+        bucket: str
+        prefix: str
+        storage_client = StorageClient(project=self.client.project)
         bucket, prefix = self._gcs_path[len("gs://") :].split("/", 1)
         prefix = prefix.rsplit("/", 1)[0]
+        if prefix.startswith("/"):
+            prefix = prefix[1:]
 
         blobs = storage_client.list_blobs(bucket, prefix=prefix)
         results = []

diff --git a/sdk/python/tests/integration/offline_store/test_universal_historical_retrieval.py b/sdk/python/tests/integration/offline_store/test_universal_historical_retrieval.py
@@ -373,7 +373,7 @@ def test_historical_features(environment, universal_data_sources, full_feature_n
     if job_from_df.supports_remote_storage_export():
         files = job_from_df.to_remote_storage()
         print(files)
-        assert len(files) > 0  # This should be way more detailed
+        assert len(files) > 0  # This test should be way more detailed
 
     start_time = datetime.utcnow()
     actual_df_from_df_entities = job_from_df.to_df()