Adjust samples

Author: plamut

samples/create_scheduled_query.py

@@ -20,10 +20,10 @@
 
 def sample_create_transfer_config(project_id, dataset_id, authorization_code=""):
     # [START bigquerydatatransfer_create_scheduled_query]
-    from google.cloud import bigquery_datatransfer_v1
+    from google.cloud.bigquery import datatransfer_v1
     import google.protobuf.json_format
 
-    client = bigquery_datatransfer_v1.DataTransferServiceClient()
+    client = datatransfer_v1.DataTransferServiceClient()
 
     # TODO(developer): Set the project_id to the project that contains the
     #                  destination dataset.
@@ -53,7 +53,7 @@ def sample_create_transfer_config(project_id, dataset_id, authorization_code="")
       17 as some_integer
     """
 
-    parent = client.project_path(project_id)
+    parent = f"projects/{project_id}"
 
     transfer_config = google.protobuf.json_format.ParseDict(
         {
@@ -68,11 +68,15 @@ def sample_create_transfer_config(project_id, dataset_id, authorization_code="")
             },
             "schedule": "every 24 hours",
         },
-        bigquery_datatransfer_v1.types.TransferConfig(),
+        datatransfer_v1.types.TransferConfig.pb()(),
     )
 
     response = client.create_transfer_config(
-        parent, transfer_config, authorization_code=authorization_code
+        request={
+            "parent": parent,
+            "transfer_config": transfer_config,
+            "authorization_code": authorization_code,
+        }
     )
 
     print("Created scheduled query '{}'".format(response.name))

samples/snippets/quickstart.py

@@ -15,21 +15,22 @@
 # limitations under the License.
 
 
-def run_quickstart():
+def run_quickstart(project="my-project"):
     # [START bigquerydatatransfer_quickstart]
-    from google.cloud import bigquery_datatransfer
+    from google.cloud.bigquery import datatransfer
 
-    client = bigquery_datatransfer.DataTransferServiceClient()
+    client = datatransfer.DataTransferServiceClient()
 
-    project = 'my-project'  # TODO: Update to your project ID.
+    # TODO: Update to your project ID.
+    # project = "my-project"
 
     # Get the full path to your project.
-    parent = client.project_path(project)
+    parent = f"projects/{project}"
 
     print('Supported Data Sources:')
 
     # Iterate over all possible data sources.
-    for data_source in client.list_data_sources(parent):
+    for data_source in client.list_data_sources(parent=parent):
         print('{}:'.format(data_source.display_name))
         print('\tID: {}'.format(data_source.data_source_id))
         print('\tFull path: {}'.format(data_source.name))

samples/snippets/quickstart_test.py

@@ -14,7 +14,6 @@
 
 import os
 
-import mock
 import pytest
 
 import quickstart
@@ -24,18 +23,13 @@
 
 
 @pytest.fixture
-def mock_project_path():
+def mock_project_id():
     """Mock out project and replace with project from environment."""
-    project_patch = mock.patch(
-        'google.cloud.bigquery_datatransfer.DataTransferServiceClient.'
-        'project_path')
 
-    with project_patch as project_mock:
-        project_mock.return_value = 'projects/{}'.format(PROJECT)
-        yield project_mock
+    return PROJECT
 
 
-def test_quickstart(capsys, mock_project_path):
-    quickstart.run_quickstart()
+def test_quickstart(capsys, mock_project_id):
+    quickstart.run_quickstart(mock_project_id)
     out, _ = capsys.readouterr()
     assert 'Supported Data Sources:' in out

samples/tests/conftest.py

@@ -39,9 +39,9 @@ def credentials():
 
 @pytest.fixture(scope="module")
 def bqdts_client(credentials):
-    from google.cloud import bigquery_datatransfer_v1
+    from google.cloud.bigquery import datatransfer_v1
 
-    return bigquery_datatransfer_v1.DataTransferServiceClient(credentials=credentials)
+    return datatransfer_v1.DataTransferServiceClient(credentials=credentials)
 
 
 @pytest.fixture(scope="module")
@@ -69,6 +69,6 @@ def to_delete(bqdts_client):
 
     for resource_name in doomed:
         try:
-            bqdts_client.delete_transfer_config(resource_name)
+            bqdts_client.delete_transfer_config(name=resource_name)
         except Exception:
             pass

samples/update_transfer_config.py

@@ -20,9 +20,9 @@
 
 def sample_update_transfer_config(config_name, display_name):
     # [START bigquerydatatransfer_update_transfer_config]
-    from google.cloud import bigquery_datatransfer_v1
+    from google.cloud.bigquery import datatransfer_v1
 
-    client = bigquery_datatransfer_v1.DataTransferServiceClient()
+    client = datatransfer_v1.DataTransferServiceClient()
     # TODO(developer): Set the config_name which user wants to update.
     # config_name = "your-created-transfer-config-name"
 
@@ -32,7 +32,9 @@ def sample_update_transfer_config(config_name, display_name):
     transfer_config = client.get_transfer_config(name=config_name)
     transfer_config.display_name = display_name
     field_mask = {"paths": ["display_name"]}
-    response = client.update_transfer_config(transfer_config, field_mask)
+    response = client.update_transfer_config(
+        transfer_config=transfer_config, update_mask=field_mask
+    )
 
     print("Transfer config updated for '{}'".format(response.name))
     # [END bigquerydatatransfer_update_transfer_config]