feat: add relationship samples (#293)

* Add export to BigQuery

* Add comments to variables

* fix: fix lint

* fix: fix the test

* feat: add relationship samples

* fix: fix build errors in main

* fix: fix feed tests

samples/snippets/conftest.py

@@ -56,13 +56,15 @@ def another_topic():
 
 @pytest.fixture(scope="module")
 def test_feed(test_topic):
+    from google.cloud import asset_v1
+
     feed_id = f"feed-{uuid.uuid4().hex}"
     asset_name = f"assets-{uuid.uuid4().hex}"
 
     @backoff.on_exception(backoff.expo, InternalServerError, max_time=60)
     def create_feed():
         return quickstart_createfeed.create_feed(
-            PROJECT, feed_id, [asset_name], test_topic.name
+            PROJECT, feed_id, [asset_name], test_topic.name, asset_v1.ContentType.RESOURCE
         )
 
     feed = create_feed()

samples/snippets/quickstart_createfeed.py

@@ -18,20 +18,22 @@
 import argparse
 
 
-def create_feed(project_id, feed_id, asset_names, topic):
+def create_feed(project_id, feed_id, asset_names, topic, content_type):
     # [START asset_quickstart_create_feed]
     from google.cloud import asset_v1
 
     # TODO project_id = 'Your Google Cloud Project ID'
     # TODO feed_id = 'Feed ID you want to create'
     # TODO asset_names = 'List of asset names the feed listen to'
     # TODO topic = "Topic name of the feed"
+    # TODO content_type ="Content type of the feed"
 
     client = asset_v1.AssetServiceClient()
     parent = "projects/{}".format(project_id)
     feed = asset_v1.Feed()
     feed.asset_names.extend(asset_names)
     feed.feed_output_config.pubsub_destination.topic = topic
+    feed.content_type = content_type
     response = client.create_feed(
         request={"parent": parent, "feed_id": feed_id, "feed": feed}
     )
@@ -48,5 +50,6 @@ def create_feed(project_id, feed_id, asset_names, topic):
     parser.add_argument("feed_id", help="Feed ID you want to create")
     parser.add_argument("asset_names", help="List of asset names the feed listen to")
     parser.add_argument("topic", help="Topic name of the feed")
+    parser.add_argument("content_type", help="Content type of the feed")
     args = parser.parse_args()
-    create_feed(args.project_id, args.feed_id, args.asset_names, args.topic)
+    create_feed(args.project_id, args.feed_id, args.asset_names, args.topic, args.content_type)

samples/snippets/quickstart_createfeed_test.py

@@ -23,12 +23,22 @@
 PROJECT = os.environ["GOOGLE_CLOUD_PROJECT"]
 ASSET_NAME = "assets-{}".format(uuid.uuid4().hex)
 FEED_ID = "feed-{}".format(uuid.uuid4().hex)
+FEED_ID_R = "feed-{}".format(uuid.uuid4().hex)
 
 
 def test_create_feed(capsys, test_topic, deleter):
+    from google.cloud import asset_v1
+
     feed = quickstart_createfeed.create_feed(
-        PROJECT, FEED_ID, [ASSET_NAME], test_topic.name
-    )
+        PROJECT, FEED_ID, [ASSET_NAME], test_topic.name,
+        asset_v1.ContentType.RESOURCE)
     deleter.append(feed.name)
     out, _ = capsys.readouterr()
     assert "feed" in out
+
+    feed_r = quickstart_createfeed.create_feed(
+        PROJECT, FEED_ID_R, [ASSET_NAME], test_topic.name,
+        asset_v1.ContentType.RELATIONSHIP)
+    deleter.append(feed_r.name)
+    out_r, _ = capsys.readouterr()
+    assert "feed" in out_r

samples/snippets/quickstart_exportassets.py

@@ -36,17 +36,17 @@ def export_assets(project_id, dump_file_path):
     # [END asset_quickstart_export_assets]
 
 
-def export_assets_bigquery(project_id, dataset, table):
+def export_assets_bigquery(project_id, dataset, table, content_type):
     # [START asset_quickstart_export_assets_bigquery]
     from google.cloud import asset_v1
 
     # TODO project_id = 'Your Google Cloud Project ID'
     # TODO dataset = 'Your BigQuery dataset path'
     # TODO table = 'Your BigQuery table name'
+    # TODO content_type ="Content type to export"
 
     client = asset_v1.AssetServiceClient()
     parent = "projects/{}".format(project_id)
-    content_type = asset_v1.ContentType.RESOURCE
     output_config = asset_v1.OutputConfig()
     output_config.bigquery_destination.dataset = dataset
     output_config.bigquery_destination.table = table

samples/snippets/quickstart_exportassets_test.py

@@ -17,6 +17,7 @@
 import os
 import uuid
 
+from google.cloud import asset_v1
 from google.cloud import bigquery
 from google.cloud import storage
 import pytest
@@ -70,8 +71,15 @@ def test_export_assets(asset_bucket, dataset, capsys):
     out, _ = capsys.readouterr()
     assert dump_file_path in out
 
+    content_type = asset_v1.ContentType.RESOURCE
     dataset_id = "projects/{}/datasets/{}".format(PROJECT, dataset)
     quickstart_exportassets.export_assets_bigquery(
-        PROJECT, dataset_id, "assettable")
+        PROJECT, dataset_id, "assettable", content_type)
     out, _ = capsys.readouterr()
     assert dataset_id in out
+
+    content_type_r = asset_v1.ContentType.RELATIONSHIP
+    quickstart_exportassets.export_assets_bigquery(
+        PROJECT, dataset_id, "assettable", content_type_r)
+    out_r, _ = capsys.readouterr()
+    assert dataset_id in out_r

samples/snippets/quickstart_listassets.py

@@ -18,7 +18,7 @@
 import argparse
 
 
-def list_assets(project_id, asset_types, page_size):
+def list_assets(project_id, asset_types, page_size, content_type):
     # [START asset_quickstart_list_assets]
     from google.cloud import asset_v1
 
@@ -27,9 +27,9 @@ def list_assets(project_id, asset_types, page_size):
     # ["storage.googleapis.com/Bucket","bigquery.googleapis.com/Table"]'
     # TODO page_size = 'Num of assets in one page, which must be between 1 and
     # 1000 (both inclusively)'
+    # TODO content_type ="Content type to list"
 
     project_resource = "projects/{}".format(project_id)
-    content_type = asset_v1.ContentType.RESOURCE
     client = asset_v1.AssetServiceClient()
 
     # Call ListAssets v1 to list assets.
@@ -63,9 +63,10 @@ def list_assets(project_id, asset_types, page_size):
         help="Num of assets in one page, which must be between 1 and 1000 "
         "(both inclusively)",
     )
+    parser.add_argument("content_type", help="Content type to list")
 
     args = parser.parse_args()
 
     asset_type_list = args.asset_types.split(",")
 
-    list_assets(args.project_id, asset_type_list, int(args.page_size))
+    list_assets(args.project_id, asset_type_list, int(args.page_size), args.content_type)

samples/snippets/quickstart_listassets_test.py

@@ -22,8 +22,15 @@
 
 
 def test_list_assets(capsys):
+    from google.cloud import asset_v1
+
     quickstart_listassets.list_assets(
-        project_id=PROJECT, asset_types=["iam.googleapis.com/Role"], page_size=10
-    )
+        project_id=PROJECT, asset_types=["iam.googleapis.com/Role"], page_size=10,
+        content_type=asset_v1.ContentType.RESOURCE)
     out, _ = capsys.readouterr()
     assert "asset" in out
+
+    quickstart_listassets.list_assets(
+        project_id=PROJECT, asset_types=[], page_size=10, content_type=asset_v1.ContentType.RELATIONSHIP)
+    out_r, _ = capsys.readouterr()
+    assert "asset" in out_r