migrate code from googleapis/python-bigquery-datatransfer

.github/CODEOWNERS

@@ -98,6 +98,7 @@
 /bigquery/remote_function/**/*         @autoerr @GoogleCloudPlatform/python-samples-reviewers
 /cloud-media-livestream/**/*           @GoogleCloudPlatform/cloud-media-team @GoogleCloudPlatform/python-samples-reviewers
 /bigquery-connection/**/*              @GoogleCloudPlatform/api-bigquery @GoogleCloudPlatform/python-samples-reviewers
+/bigquery-datatransfer/**/*            @GoogleCloudPlatform/api-bigquery @GoogleCloudPlatform/python-samples-reviewers
 /dlp/**/*                              @GoogleCloudPlatform/googleapis-dlp @GoogleCloudPlatform/python-samples-reviewers
 /functions/spanner/*                   @GoogleCloudPlatform/api-spanner-python @GoogleCloudPlatform/functions-framework-google @GoogleCloudPlatform/python-samples-reviewers
 /healthcare/**/*                       @GoogleCloudPlatform/healthcare-life-sciences @GoogleCloudPlatform/python-samples-reviewers

bigquery-datatransfer/snippets/__init__.py

@@ -0,0 +1,13 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.

bigquery-datatransfer/snippets/conftest.py

@@ -0,0 +1,156 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import os
+import random
+import uuid
+
+from google.api_core import client_options
+import google.api_core.exceptions
+import google.auth
+from google.cloud import bigquery, bigquery_datatransfer, pubsub_v1
+import pytest
+
+RESOURCE_PREFIX = "python_bigquery_datatransfer_samples_snippets"
+RESOURCE_DATE_FORMAT = "%Y%m%d%H%M%S"
+RESOURCE_DATE_LENGTH = 4 + 2 + 2 + 2 + 2 + 2
+
+
+def resource_prefix() -> str:
+    timestamp = datetime.datetime.utcnow().strftime(RESOURCE_DATE_FORMAT)
+    random_string = hex(random.randrange(1000000))[2:]
+    return f"{RESOURCE_PREFIX}_{timestamp}_{random_string}"
+
+
+def resource_name_to_date(resource_name: str):
+    start_date = len(RESOURCE_PREFIX) + 1
+    date_string = resource_name[start_date : start_date + RESOURCE_DATE_LENGTH]
+    parsed_date = datetime.datetime.strptime(date_string, RESOURCE_DATE_FORMAT)
+    return parsed_date
+
+
+@pytest.fixture(scope="session", autouse=True)
+def cleanup_pubsub_topics(pubsub_client: pubsub_v1.PublisherClient, project_id):
+    yesterday = datetime.datetime.utcnow() - datetime.timedelta(days=1)
+    for topic in pubsub_client.list_topics(project=f"projects/{project_id}"):
+        topic_id = topic.name.split("/")[-1]
+        if (
+            topic_id.startswith(RESOURCE_PREFIX)
+            and resource_name_to_date(topic_id) < yesterday
+        ):
+            pubsub_client.delete_topic(topic=topic.name)
+
+
+def temp_suffix():
+    now = datetime.datetime.now()
+    return f"{now.strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:8]}"
+
+
+@pytest.fixture(scope="session")
+def bigquery_client(default_credentials):
+    credentials, project_id = default_credentials
+    return bigquery.Client(credentials=credentials, project=project_id)
+
+
+@pytest.fixture(scope="session")
+def pubsub_client(default_credentials):
+    credentials, _ = default_credentials
+    return pubsub_v1.PublisherClient(credentials=credentials)
+
+
+@pytest.fixture(scope="session")
+def pubsub_topic(pubsub_client: pubsub_v1.PublisherClient, project_id):
+    topic_id = resource_prefix()
+    topic_path = pubsub_v1.PublisherClient.topic_path(project_id, topic_id)
+    pubsub_client.create_topic(name=topic_path)
+    yield topic_path
+    pubsub_client.delete_topic(topic=topic_path)
+
+
+@pytest.fixture(scope="session")
+def dataset_id(bigquery_client, project_id):
+    dataset_id = f"bqdts_{temp_suffix()}"
+    bigquery_client.create_dataset(f"{project_id}.{dataset_id}")
+    yield dataset_id
+    bigquery_client.delete_dataset(dataset_id, delete_contents=True)
+
+
+@pytest.fixture(scope="session")
+def default_credentials():
+    return google.auth.default(["https://www.googleapis.com/auth/cloud-platform"])
+
+
+@pytest.fixture(scope="session")
+def project_id():
+    return os.environ["GOOGLE_CLOUD_PROJECT"]
+
+
+@pytest.fixture(scope="session")
+def service_account_name(default_credentials):
+    credentials, _ = default_credentials
+    # The service_account_email attribute is not available when running with
+    # user account credentials, but should be available when running from our
+    # continuous integration tests.
+    return getattr(credentials, "service_account_email", None)
+
+
+@pytest.fixture(scope="session")
+def transfer_client(default_credentials, project_id):
+    credentials, _ = default_credentials
+    options = client_options.ClientOptions(quota_project_id=project_id)
+
+    transfer_client = bigquery_datatransfer.DataTransferServiceClient(
+        credentials=credentials, client_options=options
+    )
+
+    # Ensure quota is always attributed to the correct project.
+    bigquery_datatransfer.DataTransferServiceClient = lambda: transfer_client
+
+    return transfer_client
+
+
+@pytest.fixture(scope="session")
+def transfer_config_name(transfer_client, project_id, dataset_id, service_account_name):
+    from . import manage_transfer_configs, scheduled_query
+
+    # Use the transfer_client fixture so we know quota is attributed to the
+    # correct project.
+    assert transfer_client is not None
+
+    # To conserve limited BQ-DTS quota, this fixture creates only one transfer
+    # config for a whole session and is used to test the scheduled_query.py and
+    # the delete operation in manage_transfer_configs.py.
+    transfer_config = scheduled_query.create_scheduled_query(
+        {
+            "project_id": project_id,
+            "dataset_id": dataset_id,
+            "service_account_name": service_account_name,
+        }
+    )
+    yield transfer_config.name
+    manage_transfer_configs.delete_config(
+        {"transfer_config_name": transfer_config.name}
+    )
+
+
+@pytest.fixture
+def to_delete_configs(transfer_client):
+    to_delete = []
+    yield to_delete
+    for config_name in to_delete:
+        try:
+            transfer_client.delete_transfer_config(name=config_name)
+        except google.api_core.exceptions.GoogleAPICallError:
+            pass

bigquery-datatransfer/snippets/copy_dataset.py

@@ -0,0 +1,54 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+def copy_dataset(override_values={}):
+    # [START bigquerydatatransfer_copy_dataset]
+    from google.cloud import bigquery_datatransfer
+
+    transfer_client = bigquery_datatransfer.DataTransferServiceClient()
+
+    destination_project_id = "my-destination-project"
+    destination_dataset_id = "my_destination_dataset"
+    source_project_id = "my-source-project"
+    source_dataset_id = "my_source_dataset"
+    # [END bigquerydatatransfer_copy_dataset]
+    # To facilitate testing, we replace values with alternatives
+    # provided by the testing harness.
+    destination_project_id = override_values.get(
+        "destination_project_id", destination_project_id
+    )
+    destination_dataset_id = override_values.get(
+        "destination_dataset_id", destination_dataset_id
+    )
+    source_project_id = override_values.get("source_project_id", source_project_id)
+    source_dataset_id = override_values.get("source_dataset_id", source_dataset_id)
+    # [START bigquerydatatransfer_copy_dataset]
+    transfer_config = bigquery_datatransfer.TransferConfig(
+        destination_dataset_id=destination_dataset_id,
+        display_name="Your Dataset Copy Name",
+        data_source_id="cross_region_copy",
+        params={
+            "source_project_id": source_project_id,
+            "source_dataset_id": source_dataset_id,
+        },
+        schedule="every 24 hours",
+    )
+    transfer_config = transfer_client.create_transfer_config(
+        parent=transfer_client.common_project_path(destination_project_id),
+        transfer_config=transfer_config,
+    )
+    print(f"Created transfer config: {transfer_config.name}")
+    # [END bigquerydatatransfer_copy_dataset]
+    return transfer_config

bigquery-datatransfer/snippets/copy_dataset_test.py

@@ -0,0 +1,66 @@
+# Copyright 2020 Google LLC
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import datetime
+import uuid
+
+import pytest
+
+from . import copy_dataset
+
+
+def temp_suffix():
+    now = datetime.datetime.now()
+    return f"{now.strftime('%Y%m%d%H%M%S')}_{uuid.uuid4().hex[:8]}"
+
+
+@pytest.fixture(scope="module")
+def destination_dataset_id(bigquery_client, project_id):
+    dataset_id = f"bqdts_dest_{temp_suffix()}"
+    bigquery_client.create_dataset(f"{project_id}.{dataset_id}")
+    yield dataset_id
+    bigquery_client.delete_dataset(dataset_id, delete_contents=True)
+
+
+@pytest.fixture(scope="module")
+def source_dataset_id(bigquery_client, project_id):
+    dataset_id = f"bqdts_src_{temp_suffix()}"
+    bigquery_client.create_dataset(f"{project_id}.{dataset_id}")
+    yield dataset_id
+    bigquery_client.delete_dataset(dataset_id, delete_contents=True)
+
+
+def test_copy_dataset(
+    capsys,
+    transfer_client,
+    project_id,
+    destination_dataset_id,
+    source_dataset_id,
+    to_delete_configs,
+):
+    # Use the transfer_client fixture so we know quota is attributed to the
+    # correct project.
+    assert transfer_client is not None
+
+    transfer_config = copy_dataset.copy_dataset(
+        {
+            "destination_project_id": project_id,
+            "destination_dataset_id": destination_dataset_id,
+            "source_project_id": project_id,
+            "source_dataset_id": source_dataset_id,
+        }
+    )
+    to_delete_configs.append(transfer_config.name)
+    out, _ = capsys.readouterr()
+    assert transfer_config.name in out