From e8d1ff24ccbbef79bb3685e27d48d8ae3be2cd6e Mon Sep 17 00:00:00 2001 From: Patrick Marlow Date: Mon, 5 Feb 2024 16:28:06 -0600 Subject: [PATCH 1/2] fix: correct data store api endpoint build for DRZ --- src/dfcx_scrapi/core/data_stores.py | 19 ++++++++++++++++--- 1 file changed, 16 insertions(+), 3 deletions(-) diff --git a/src/dfcx_scrapi/core/data_stores.py b/src/dfcx_scrapi/core/data_stores.py index 54618a77..20d6b564 100644 --- a/src/dfcx_scrapi/core/data_stores.py +++ b/src/dfcx_scrapi/core/data_stores.py @@ -82,7 +82,12 @@ def list_data_stores( self, location: str = "global") -> List[DataStore]: """List all data stores for a given project and location.""" parent = self._build_data_store_parent(location) - client = discoveryengine_v1alpha.DataStoreServiceClient() + client_options = self._client_options_discovery_engine( + f"projects/{self.project_id}/locations/{location}" + ) + client = discoveryengine_v1alpha.DataStoreServiceClient( + credentials=self.creds, client_options=client_options + ) request = discoveryengine_v1alpha.ListDataStoresRequest(parent=parent) page_result = client.list_data_stores(request=request) @@ -94,7 +99,10 @@ def list_data_stores( def get_data_store(self, data_store_id: str) -> DataStore: """Get a single Data Store by specified ID.""" - client = discoveryengine_v1alpha.DataStoreServiceClient() + client_options = self._client_options_discovery_engine(data_store_id) + client = discoveryengine_v1alpha.DataStoreServiceClient( + credentials=self.creds, client_options=client_options + ) request = discoveryengine_v1alpha.GetDataStoreRequest( name=data_store_id ) @@ -122,7 +130,12 @@ def create_datastore( location: the GCP region to create the Data Store in """ parent = self._build_data_store_parent(location) - client = discoveryengine_v1alpha.DataStoreServiceClient() + client_options = self._client_options_discovery_engine( + f"projects/{self.project_id}/locations/{location}" + ) + client = discoveryengine_v1alpha.DataStoreServiceClient( + credentials=self.creds, client_options=client_options + ) data_store = discoveryengine_v1alpha.DataStore() data_store.display_name = display_name data_store.industry_vertical = 1 From 8899bf412572636843e8061557821494537dc3a0 Mon Sep 17 00:00:00 2001 From: Patrick Marlow Date: Wed, 7 Feb 2024 21:32:27 -0600 Subject: [PATCH 2/2] chore: modify method name --- src/dfcx_scrapi/core/data_stores.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/src/dfcx_scrapi/core/data_stores.py b/src/dfcx_scrapi/core/data_stores.py index 20d6b564..4fd944f1 100644 --- a/src/dfcx_scrapi/core/data_stores.py +++ b/src/dfcx_scrapi/core/data_stores.py @@ -50,7 +50,7 @@ def __init__( self.project_id = project_id @staticmethod - def __get_content_config(content_type: str) -> DataStore.ContentConfig: + def _get_content_config(content_type: str) -> DataStore.ContentConfig: """Build the Content Config used for the Data Store.""" cc_map = { "unstructured": 1, @@ -140,7 +140,7 @@ def create_datastore( data_store.display_name = display_name data_store.industry_vertical = 1 data_store.solution_types = [self._get_solution_type(solution_type)] - data_store.content_config = self.__get_content_config(datastore_type) + data_store.content_config = self._get_content_config(datastore_type) request = discoveryengine_v1alpha.CreateDataStoreRequest( parent=parent,