Add sample to schedule query with BQ DTS. (googleapis#7703)

tswast · web-flow · commit 0b232313a1f8 · 2019-06-06T10:38:23.000-07:00
* Add semi-generated sample for BQ DTS.

* Add authorization_code as parameter.

* Give scheduled query example a shorter filename.

* Add test for create scheduled query sample.

* Create dataset in sample tests.

This ensures the user (or service account) running the tests has owner permissions, which is required to run scheduled queries.

* Add samples nox session.
diff --git a/bigquery_datatransfer/noxfile.py b/bigquery_datatransfer/noxfile.py
@@ -46,7 +46,7 @@ def blacken(session):
     """Run black.
 
     Format code to uniform standard.
-    
+
     This currently uses Python 3.6 due to the automated Kokoro run of synthtool.
     That run uses an image that doesn't have 3.6 installed. Before updating this
     check the state of the `gcp_ubuntu_config` we use for that Kokoro run.
@@ -123,6 +123,25 @@ def system(session):
         session.run("py.test", "--quiet", system_test_folder_path, *session.posargs)
 
 
+@nox.session(python=["2.7", "3.7"])
+def samples(session):
+    requirements_path = os.path.join("samples", "requirements.txt")
+    requirements_exists = os.path.exists(requirements_path)
+
+    # Sanity check: Only run tests if the environment variable is set.
+    if not os.environ.get("GOOGLE_APPLICATION_CREDENTIALS", ""):
+        session.skip("Credentials must be set via environment variable")
+
+    session.install("mock", "pytest")
+    for local_dep in LOCAL_DEPS:
+        session.install("-e", local_dep)
+    if requirements_exists:
+        session.install("-r", requirements_path)
+    session.install("-e", ".")
+
+    session.run("py.test", "--quiet", "samples", *session.posargs)
+
+
 @nox.session(python="3.7")
 def cover(session):
     """Run the final coverage report.
diff --git a/bigquery_datatransfer/samples/__init__.py b/bigquery_datatransfer/samples/__init__.py
diff --git a/bigquery_datatransfer/samples/create_scheduled_query.py b/bigquery_datatransfer/samples/create_scheduled_query.py
@@ -0,0 +1,97 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 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
+#
+#     https://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.
+
+# To install the latest published package dependency, execute the following:
+#   pip install google-cloud-bigquery-datatransfer
+
+
+def sample_create_transfer_config(project_id, dataset_id, authorization_code=""):
+    # [START bigquerydatatransfer_create_scheduled_query]
+    from google.cloud import bigquery_datatransfer_v1
+    import google.protobuf.json_format
+
+    client = bigquery_datatransfer_v1.DataTransferServiceClient()
+
+    # TODO(developer): Set the project_id to the project that contains the
+    #                  destination dataset.
+    # project_id = "your-project-id"
+
+    # TODO(developer): Set the destination dataset. The authorized user must
+    #                  have owner permissions on the dataset.
+    # dataset_id = "your_dataset_id"
+
+    # TODO(developer): The first time you run this sample, set the
+    # authorization code to a value from the URL:
+    # https://www.gstatic.com/bigquerydatatransfer/oauthz/auth?client_id=433065040935-hav5fqnc9p9cht3rqneus9115ias2kn1.apps.googleusercontent.com&scope=https://www.googleapis.com/auth/bigquery%20https://www.googleapis.com/auth/drive&redirect_uri=urn:ietf:wg:oauth:2.0:oob
+    #
+    # authorization_code = "_4/ABCD-EFGHIJKLMNOP-QRSTUVWXYZ"
+    #
+    # You can use an empty string for authorization_code in subsequent runs of
+    # this code sample with the same credentials.
+    #
+    # authorization_code = ""
+
+    # Use standard SQL syntax for the query.
+    query_string = """
+SELECT
+  CURRENT_TIMESTAMP() as current_time,
+  @run_time as intended_run_time,
+  @run_date as intended_run_date,
+  17 as some_integer
+"""
+
+    parent = client.project_path(project_id)
+
+    transfer_config = google.protobuf.json_format.ParseDict(
+        {
+            "destination_dataset_id": dataset_id,
+            "display_name": "Your Scheduled Query Name",
+            "data_source_id": "scheduled_query",
+            "params": {
+                "query": query_string,
+                "destination_table_name_template": "your_table_{run_date}",
+                "write_disposition": "WRITE_TRUNCATE",
+                "partitioning_field": "",
+            },
+            "schedule": "every 24 hours",
+        },
+        bigquery_datatransfer_v1.types.TransferConfig(),
+    )
+
+    response = client.create_transfer_config(
+        parent, transfer_config, authorization_code=authorization_code
+    )
+
+    print("Created scheduled query '{}'".format(response.name))
+    # [END bigquerydatatransfer_create_scheduled_query]
+    # Return the config name for testing purposes, so that it can be deleted.
+    return response.name
+
+
+def main():
+    import argparse
+
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--project_id", type=str, default="your-project-id")
+    parser.add_argument("--dataset_id", type=str, default="your_dataset_id")
+    parser.add_argument("--authorization_code", type=str, default="")
+    args = parser.parse_args()
+
+    sample_create_transfer_config(args.project_id, args.authorization_code)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/bigquery_datatransfer/samples/requirements.txt b/bigquery_datatransfer/samples/requirements.txt
@@ -0,0 +1 @@
+-e ../bigquery
diff --git a/bigquery_datatransfer/samples/tests/__init__.py b/bigquery_datatransfer/samples/tests/__init__.py
diff --git a/bigquery_datatransfer/samples/tests/test_create_scheduled_query.py b/bigquery_datatransfer/samples/tests/test_create_scheduled_query.py
@@ -0,0 +1,82 @@
+# -*- coding: utf-8 -*-
+#
+# Copyright 2019 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
+#
+#     https://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 time
+import os
+
+import google.api_core.exceptions
+import google.auth
+import google.cloud.bigquery
+import pytest
+
+from .. import create_scheduled_query
+
+
+@pytest.fixture
+def project_id():
+    return os.environ["PROJECT_ID"]
+
+
+@pytest.fixture(scope="module")
+def credentials():
+    # If using a service account, the BQ DTS robot associated with your project
+    # requires the roles/iam.serviceAccountShortTermTokenMinter permission to
+    # act on behalf of the account.
+    creds, _ = google.auth.default(["https://www.googleapis.com/auth/cloud-platform"])
+    return creds
+
+
+@pytest.fixture(scope="module")
+def bqdts_client(credentials):
+    from google.cloud import bigquery_datatransfer_v1
+
+    return bigquery_datatransfer_v1.DataTransferServiceClient(credentials=credentials)
+
+
+@pytest.fixture(scope="module")
+def bigquery_client(credentials):
+    return google.cloud.bigquery.Client(credentials=credentials)
+
+
+@pytest.fixture(scope="module")
+def dataset_id(bigquery_client):
+    # Ensure the test account has owner permissions on the dataset by creating
+    # one from scratch.
+    temp_ds_id = "bqdts_{}".format(int(time.clock() * 1000000))
+    bigquery_client.create_dataset(temp_ds_id)
+    yield temp_ds_id
+    bigquery_client.delete_dataset(temp_ds_id)
+
+
+@pytest.fixture
+def to_delete(bqdts_client):
+    doomed = []
+    yield doomed
+
+    for resource_name in doomed:
+        try:
+            bqdts_client.delete_transfer_config(resource_name)
+        except google.api_core.exceptions.NotFound:
+            pass
+
+
+def test_sample(project_id, dataset_id, capsys, to_delete):
+    config_name = create_scheduled_query.sample_create_transfer_config(
+        project_id, dataset_id
+    )
+    to_delete.append(config_name)
+    out, err = capsys.readouterr()
+    assert config_name in out
diff --git a/bigquery_datatransfer/synth.py b/bigquery_datatransfer/synth.py
@@ -62,7 +62,7 @@
 # ----------------------------------------------------------------------------
 # Add templated files
 # ----------------------------------------------------------------------------
-templated_files = common.py_library(unit_cov_level=79, cov_level=79)
+templated_files = common.py_library(unit_cov_level=79, cov_level=79, samples_test=True)
 s.move(templated_files)
 
 s.shell.run(["nox", "-s", "blacken"], hide_output=False)
