samples: Add relationship and per-type bigquery samples to Cloud Asset service (#882)

lvvvvvf · Neenu1995 · gcf-owl-bot[bot] · web-flow · commit 56377d9a42ee · 2021-09-03T13:32:12.000-07:00
* Add relationship and per-type Bigquery samples * remove wait for BatchGetHistory, as we are not creating new bucket now * fix: fix format * fix typo * fix: use diff table name to avoid conflict * fix: fix format and add comments * fix: add detailed exceptions * fix: fix format * 🦉 Updates from OwlBot See https://github.com/googleapis/repo-automation-bots/blob/main/packages/owl-bot/README.md Co-authored-by: Neenu Shaji <Neenu1995@users.noreply.github.com> Co-authored-by: Owl Bot <gcf-owl-bot[bot]@users.noreply.github.com>
diff --git a/asset/src/main/java/com/example/asset/CreateFeedExample.java b/asset/src/main/java/com/example/asset/CreateFeedExample.java
@@ -18,24 +18,29 @@
 
 // [START asset_quickstart_create_feed]
 import com.google.cloud.asset.v1.AssetServiceClient;
+import com.google.cloud.asset.v1.ContentType;
 import com.google.cloud.asset.v1.CreateFeedRequest;
 import com.google.cloud.asset.v1.Feed;
 import com.google.cloud.asset.v1.FeedOutputConfig;
 import com.google.cloud.asset.v1.ProjectName;
 import com.google.cloud.asset.v1.PubsubDestination;
+import java.io.IOException;
 import java.util.Arrays;
 
 public class CreateFeedExample {
   // Create a feed
-  public static void createFeed(String[] assetNames, String feedId, String topic, String projectId)
-      throws Exception {
+  public static void createFeed(
+      String[] assetNames, String feedId, String topic, String projectId, ContentType contentType)
+      throws IOException, IllegalArgumentException {
     // String[] assetNames = {"MY_ASSET_NAME"}
+    // ContentType contentType = contentType
     // String FeedId = "MY_FEED_ID"
     // String topic = "projects/[PROJECT_ID]/topics/[TOPIC_NAME]"
     // String projectID = "MY_PROJECT_ID"
     Feed feed =
         Feed.newBuilder()
             .addAllAssetNames(Arrays.asList(assetNames))
+            .setContentType(contentType)
             .setFeedOutputConfig(
                 FeedOutputConfig.newBuilder()
                     .setPubsubDestination(PubsubDestination.newBuilder().setTopic(topic).build())
@@ -53,7 +58,7 @@ public static void createFeed(String[] assetNames, String feedId, String topic,
     try (AssetServiceClient client = AssetServiceClient.create()) {
       Feed response = client.createFeed(request);
       System.out.println("Feed created successfully: " + response.getName());
-    } catch (Exception e) {
+    } catch (IOException | IllegalArgumentException e) {
       System.out.println("Error during CreateFeed: \n" + e.toString());
     }
   }
diff --git a/asset/src/main/java/com/example/asset/ExportAssetsBigqueryExample.java b/asset/src/main/java/com/example/asset/ExportAssetsBigqueryExample.java
@@ -22,33 +22,58 @@
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.asset.v1.AssetServiceClient;
 import com.google.cloud.asset.v1.BigQueryDestination;
+import com.google.cloud.asset.v1.ContentType;
 import com.google.cloud.asset.v1.ExportAssetsRequest;
 import com.google.cloud.asset.v1.ExportAssetsResponse;
 import com.google.cloud.asset.v1.OutputConfig;
+import com.google.cloud.asset.v1.PartitionSpec;
 import com.google.cloud.asset.v1.ProjectName;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
 
 public class ExportAssetsBigqueryExample {
 
   // Use the default project Id.
   private static final String projectId = ServiceOptions.getDefaultProjectId();
 
-  // Export assets for a project.
-  // @param args path where the results will be exported to.
-  public static void exportBigQuery(String bigqueryDataset, String bigqueryTable) throws Exception {
+  // Export assets to BigQuery for a project.
+  public static void exportBigQuery(
+      String bigqueryDataset, String bigqueryTable, ContentType contentType, boolean isPerType)
+      throws IOException, IllegalArgumentException, InterruptedException, ExecutionException {
     try (AssetServiceClient client = AssetServiceClient.create()) {
       ProjectName parent = ProjectName.of(projectId);
-      OutputConfig outputConfig =
-          OutputConfig.newBuilder()
-              .setBigqueryDestination(
-                  BigQueryDestination.newBuilder()
-                      .setDataset(bigqueryDataset)
-                      .setTable(bigqueryTable)
-                      .setForce(true)
-                      .build())
-              .build();
+      OutputConfig outputConfig;
+      // Outputs to per-type BigQuery table.
+      if (isPerType) {
+        outputConfig =
+            OutputConfig.newBuilder()
+                .setBigqueryDestination(
+                    BigQueryDestination.newBuilder()
+                        .setDataset(bigqueryDataset)
+                        .setTable(bigqueryTable)
+                        .setForce(true)
+                        .setSeparateTablesPerAssetType(true)
+                        .setPartitionSpec(
+                            PartitionSpec.newBuilder()
+                                .setPartitionKey(PartitionSpec.PartitionKey.READ_TIME)
+                                .build())
+                        .build())
+                .build();
+      } else {
+        outputConfig =
+            OutputConfig.newBuilder()
+                .setBigqueryDestination(
+                    BigQueryDestination.newBuilder()
+                        .setDataset(bigqueryDataset)
+                        .setTable(bigqueryTable)
+                        .setForce(true)
+                        .build())
+                .build();
+      }
       ExportAssetsRequest request =
           ExportAssetsRequest.newBuilder()
               .setParent(parent.toString())
+              .setContentType(contentType)
               .setOutputConfig(outputConfig)
               .build();
       ExportAssetsResponse response = client.exportAssetsAsync(request).get();
diff --git a/asset/src/main/java/com/example/asset/ExportAssetsExample.java b/asset/src/main/java/com/example/asset/ExportAssetsExample.java
@@ -21,22 +21,24 @@
 
 import com.google.cloud.ServiceOptions;
 import com.google.cloud.asset.v1.AssetServiceClient;
+import com.google.cloud.asset.v1.ContentType;
 import com.google.cloud.asset.v1.ExportAssetsRequest;
 import com.google.cloud.asset.v1.ExportAssetsResponse;
 import com.google.cloud.asset.v1.GcsDestination;
 import com.google.cloud.asset.v1.OutputConfig;
 import com.google.cloud.asset.v1.ProjectName;
+import java.io.IOException;
+import java.util.concurrent.ExecutionException;
 
 public class ExportAssetsExample {
 
   // Use the default project Id.
   private static final String projectId = ServiceOptions.getDefaultProjectId();
 
   // Export assets for a project.
-  // @param args path where the results will be exported to.
-  public static void main(String... args) throws Exception {
-    // Gcs path, e.g.: "gs://<my_asset_bucket>/<my_asset_dump_file>"
-    String exportPath = args[0];
+  // @param exportPath where the results will be exported to.
+  public static void exportAssets(String exportPath, ContentType contentType)
+      throws IOException, IllegalArgumentException, InterruptedException, ExecutionException {
     try (AssetServiceClient client = AssetServiceClient.create()) {
       ProjectName parent = ProjectName.of(projectId);
       OutputConfig outputConfig =
@@ -47,10 +49,11 @@ public static void main(String... args) throws Exception {
           ExportAssetsRequest.newBuilder()
               .setParent(parent.toString())
               .setOutputConfig(outputConfig)
+              .setContentType(contentType)
               .build();
       ExportAssetsResponse response = client.exportAssetsAsync(request).get();
       System.out.println(response);
     }
   }
 }
-// [END asset_quickstart]
+// [END asset_quickstart_export_assets]
diff --git a/asset/src/main/java/com/example/asset/ListAssetsExample.java b/asset/src/main/java/com/example/asset/ListAssetsExample.java
@@ -21,14 +21,15 @@
 import com.google.cloud.asset.v1.ContentType;
 import com.google.cloud.asset.v1.ListAssetsRequest;
 import com.google.cloud.asset.v1.ProjectName;
+import java.io.IOException;
 import java.util.Arrays;
 
 // [START asset_quickstart_list_assets]
 // Imports the Google Cloud client library
 
 public class ListAssetsExample {
 
-  public static void listAssets() throws Exception {
+  public static void listAssets() throws IOException, IllegalArgumentException {
     // The project id of the asset parent to list.
     String projectId = "YOUR_PROJECT_ID";
     // The asset types to list. E.g.,
@@ -44,7 +45,7 @@ public static void listAssets() throws Exception {
   }
 
   public static void listAssets(String projectId, String[] assetTypes, ContentType contentType)
-      throws Exception {
+      throws IOException, IllegalArgumentException {
     try (AssetServiceClient client = AssetServiceClient.create()) {
       ProjectName parent = ProjectName.of(projectId);
 
diff --git a/asset/src/test/java/com/example/asset/ListAssets.java b/asset/src/test/java/com/example/asset/ListAssets.java
@@ -64,4 +64,18 @@ public void testListAssetsExample() throws Exception {
       assertThat(got).contains("asset");
     }
   }
+
+  @Test
+  public void testListAssetsRelationshipExample() throws Exception {
+    // Use the default project Id (configure it by setting environment variable
+    // "GOOGLE_CLOUD_PROJECT").
+    String projectId = ServiceOptions.getDefaultProjectId();
+    String[] assetTypes = {"compute.googleapis.com/Instance", "compute.googleapis.com/Disk"};
+    ContentType contentType = ContentType.RELATIONSHIP;
+    ListAssetsExample.listAssets(projectId, assetTypes, contentType);
+    String got = bout.toString();
+    if (!got.isEmpty()) {
+      assertThat(got).contains("asset");
+    }
+  }
 }
diff --git a/asset/src/test/java/com/example/asset/QuickStartIT.java b/asset/src/test/java/com/example/asset/QuickStartIT.java
@@ -19,6 +19,7 @@
 import static com.google.common.truth.Truth.assertThat;
 
 import com.google.cloud.ServiceOptions;
+import com.google.cloud.asset.v1.ContentType;
 import com.google.cloud.bigquery.BigQuery;
 import com.google.cloud.bigquery.BigQuery.DatasetDeleteOption;
 import com.google.cloud.bigquery.BigQueryOptions;
@@ -90,25 +91,35 @@ public void tearDown() {
   @Test
   public void testExportAssetExample() throws Exception {
     String assetDumpPath = String.format("gs://%s/%s/my-assets-dump.txt", bucketName, path);
-    ExportAssetsExample.main(assetDumpPath);
+    ExportAssetsExample.exportAssets(assetDumpPath, ContentType.RESOURCE);
     String got = bout.toString();
     assertThat(got).contains(String.format("uri: \"%s\"", assetDumpPath));
   }
 
+  @Test
+  public void testExportAssetBigqueryPerTypeExample() throws Exception {
+    String dataset =
+        String.format("projects/%s/datasets/%s", ServiceOptions.getDefaultProjectId(), datasetName);
+    String table = "java_test_per_type";
+    ExportAssetsBigqueryExample.exportBigQuery(
+        dataset, table, ContentType.RESOURCE, /*perType*/ true);
+    String got = bout.toString();
+    assertThat(got).contains(String.format("dataset: \"%s\"", dataset));
+  }
+
   @Test
   public void testExportAssetBigqueryExample() throws Exception {
     String dataset =
         String.format("projects/%s/datasets/%s", ServiceOptions.getDefaultProjectId(), datasetName);
     String table = "java_test";
-    ExportAssetsBigqueryExample.exportBigQuery(dataset, table);
+    ExportAssetsBigqueryExample.exportBigQuery(
+        dataset, table, ContentType.RELATIONSHIP, /*perType*/ false);
     String got = bout.toString();
     assertThat(got).contains(String.format("dataset: \"%s\"", dataset));
   }
 
   @Test
   public void testBatchGetAssetsHistory() throws Exception {
-    // Wait 10 seconds to let bucket creation event go to CAI
-    Thread.sleep(10000);
     String bucketAssetName = String.format("//storage.googleapis.com/%s", bucketName);
     BatchGetAssetsHistoryExample.main(bucketAssetName);
     String got = bout.toString();
diff --git a/asset/src/test/java/com/example/asset/RealTimeFeed.java b/asset/src/test/java/com/example/asset/RealTimeFeed.java
@@ -18,6 +18,7 @@
 
 import static com.google.common.truth.Truth.assertThat;
 
+import com.google.cloud.asset.v1.ContentType;
 import com.google.cloud.pubsub.v1.TopicAdminClient;
 import com.google.cloud.resourcemanager.ProjectInfo;
 import com.google.cloud.resourcemanager.ResourceManager;
@@ -87,7 +88,20 @@ public void tearDown() {
 
   @Test
   public void test1CreateFeedExample() throws Exception {
-    CreateFeedExample.createFeed(assetNames, feedId, topicName.toString(), projectId);
+    CreateFeedExample.createFeed(
+        assetNames, feedId, topicName.toString(), projectId, ContentType.RESOURCE);
+    String got = bout.toString();
+    assertThat(got).contains("Feed created successfully: " + feedName);
+  }
+
+  @Test
+  public void test1CreateFeedRelationshipExample() throws Exception {
+    CreateFeedExample.createFeed(
+        assetNames,
+        feedId + "relationship",
+        topicName.toString(),
+        projectId,
+        ContentType.RELATIONSHIP);
     String got = bout.toString();
     assertThat(got).contains("Feed created successfully: " + feedName);
   }
@@ -116,6 +130,7 @@ public void test4UpdateFeedExample() throws Exception {
   @Test
   public void test5DeleteFeedExample() throws Exception {
     DeleteFeedExample.deleteFeed(feedName);
+    DeleteFeedExample.deleteFeed(feedName + "relationship");
     String got = bout.toString();
     assertThat(got).contains("Feed deleted");
   }

Original file line number	Diff line number	Diff line change
`@@ -64,4 +64,18 @@ public void testListAssetsExample() throws Exception {`
`64`	`64`	`assertThat(got).contains("asset");`
`65`	`65`	`}`
`66`	`66`	`}`
	`67`	`+`
	`68`	`+ @Test`
	`69`	`+ public void testListAssetsRelationshipExample() throws Exception {`
	`70`	`+ // Use the default project Id (configure it by setting environment variable`
	`71`	`+ // "GOOGLE_CLOUD_PROJECT").`
	`72`	`+ String projectId = ServiceOptions.getDefaultProjectId();`
	`73`	`+ String[] assetTypes = {"compute.googleapis.com/Instance", "compute.googleapis.com/Disk"};`
	`74`	`+ ContentType contentType = ContentType.RELATIONSHIP;`
	`75`	`+ ListAssetsExample.listAssets(projectId, assetTypes, contentType);`
	`76`	`+ String got = bout.toString();`
	`77`	`+ if (!got.isEmpty()) {`
	`78`	`+ assertThat(got).contains("asset");`
	`79`	`+ }`
	`80`	`+ }`
`67`	`81`	`}`