samples(discoveryengine): Add AlloyDB Import Sample (GoogleCloudPlatform#12032)

holtskinner · dandhlee · web-flow · commit 49ea4a76e849 · 2024-07-30T18:18:19.000+02:00
---------

Co-authored-by: Dan Lee &lt;71398022+dandhlee@users.noreply.github.com&gt;
diff --git a/discoveryengine/documents_sample_test.py b/discoveryengine/documents_sample_test.py
@@ -126,6 +126,21 @@ def test_import_documents_bigtable():
     assert "operations/import-documents" in operation_name
 
 
+def test_import_documents_alloy_db():
+    operation_name = import_documents_sample.import_documents_alloy_db_sample(
+        project_id=project_id,
+        location=location,
+        data_store_id=data_store_id,
+        alloy_db_project_id=project_id,
+        alloy_db_location_id="us-central1",
+        alloy_db_cluster_id="vais-tests",
+        alloy_db_database_id="postgres",
+        alloy_db_table_id="public.vais",
+    )
+
+    assert "operations/import-documents" in operation_name
+
+
 @pytest.mark.skip(reason="Permissions")
 def test_import_documents_healthcare_fhir_sample():
     location = "us"
diff --git a/discoveryengine/import_documents_sample.py b/discoveryengine/import_documents_sample.py
@@ -455,6 +455,81 @@ def import_documents_bigtable_sample(
     return operation.operation.name
 
 
+def import_documents_alloy_db_sample(
+    project_id: str,
+    location: str,
+    data_store_id: str,
+    alloy_db_project_id: str,
+    alloy_db_location_id: str,
+    alloy_db_cluster_id: str,
+    alloy_db_database_id: str,
+    alloy_db_table_id: str,
+) -> str:
+    # [START genappbuilder_import_documents_alloy_db]
+    from google.api_core.client_options import ClientOptions
+    from google.cloud import discoveryengine_v1 as discoveryengine
+
+    # TODO(developer): Uncomment these variables before running the sample.
+    # project_id = "YOUR_PROJECT_ID"
+    # location = "YOUR_LOCATION" # Values: "global"
+    # data_store_id = "YOUR_DATA_STORE_ID"
+    # alloy_db_project_id = "YOUR_ALLOY_DB_PROJECT_ID"
+    # alloy_db_location_id = "YOUR_ALLOY_DB_LOCATION_ID"
+    # alloy_db_cluster_id = "YOUR_ALLOY_DB_CLUSTER_ID"
+    # alloy_db_database_id = "YOUR_ALLOY_DB_DATABASE_ID"
+    # alloy_db_table_id = "YOUR_ALLOY_DB_TABLE_ID"
+
+    # For more information, refer to:
+    # https://cloud.google.com/generative-ai-app-builder/docs/locations#specify_a_multi-region_for_your_data_store
+    client_options = (
+        ClientOptions(api_endpoint=f"{location}-discoveryengine.googleapis.com")
+        if location != "global"
+        else None
+    )
+
+    # Create a client
+    client = discoveryengine.DocumentServiceClient(client_options=client_options)
+
+    # The full resource name of the search engine branch.
+    # e.g. projects/{project}/locations/{location}/dataStores/{data_store_id}/branches/{branch}
+    parent = client.branch_path(
+        project=project_id,
+        location=location,
+        data_store=data_store_id,
+        branch="default_branch",
+    )
+
+    request = discoveryengine.ImportDocumentsRequest(
+        parent=parent,
+        alloy_db_source=discoveryengine.AlloyDbSource(
+            project_id=alloy_db_project_id,
+            location_id=alloy_db_location_id,
+            cluster_id=alloy_db_cluster_id,
+            database_id=alloy_db_database_id,
+            table_id=alloy_db_table_id,
+        ),
+        # Options: `FULL`, `INCREMENTAL`
+        reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
+    )
+
+    # Make the request
+    operation = client.import_documents(request=request)
+
+    print(f"Waiting for operation to complete: {operation.operation.name}")
+    response = operation.result()
+
+    # Once the operation is complete,
+    # get information from operation metadata
+    metadata = discoveryengine.ImportDocumentsMetadata(operation.metadata)
+
+    # Handle the response
+    print(response)
+    print(metadata)
+    # [END genappbuilder_import_documents_alloy_db]
+
+    return operation.operation.name
+
+
 def import_documents_healthcare_fhir_sample(
     project_id: str,
     location: str,
@@ -500,7 +575,12 @@ def import_documents_healthcare_fhir_sample(
     request = discoveryengine.ImportDocumentsRequest(
         parent=parent,
         fhir_store_source=discoveryengine.FhirStoreSource(
-            fhir_store=f"projects/{healthcare_project_id}/locations/{healthcare_location}/datasets/{healthcare_dataset_id}/fhirStores/{healthcare_fihr_store_id}",
+            fhir_store=client.fhir_store_path(
+                healthcare_project_id,
+                healthcare_location,
+                healthcare_dataset_id,
+                healthcare_fihr_store_id,
+            ),
         ),
         # Options: `FULL`, `INCREMENTAL`
         reconciliation_mode=discoveryengine.ImportDocumentsRequest.ReconciliationMode.INCREMENTAL,
diff --git a/discoveryengine/requirements.txt b/discoveryengine/requirements.txt
@@ -1,2 +1,2 @@
-google-cloud-discoveryengine==0.11.13
+google-cloud-discoveryengine==0.11.14
 google-api-core==2.19.0

Original file line number	Diff line number	Diff line change
`@@ -1,2 +1,2 @@`
`1`		`-google-cloud-discoveryengine==0.11.13`
	`1`	`+google-cloud-discoveryengine==0.11.14`
`2`	`2`	`google-api-core==2.19.0`