Composer: Move trigger response DAG to GitHub. (#1645)

* Composer: Move trigger response DAG to GitHub.

Sample originally was posted at https://cloud.google.com/composer/docs/how-to/using/triggering-with-gcf#wzxhzdk15gcs_response_dagpywzxhzdk16

* Composer: Add snippet to get the client ID for a composer env.

Author: tswast

composer/rest/get_client_id.py

@@ -0,0 +1,70 @@
+# Copyright 2018 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.
+
+"""Get the client ID associated with a Cloud Composer environment."""
+
+import argparse
+
+
+def get_client_id(project_id, location, composer_environment):
+    # [START composer_get_environment_client_id]
+    import google.auth
+    import google.auth.transport.requests
+    import requests
+    import six.moves.urllib.parse
+
+    # Authenticate with Google Cloud.
+    # See: https://cloud.google.com/docs/authentication/getting-started
+    credentials, _ = google.auth.default(
+        scopes=['https://www.googleapis.com/auth/cloud-platform'])
+    authed_session = google.auth.transport.requests.AuthorizedSession(
+        credentials)
+
+    # project_id = 'YOUR_PROJECT_ID'
+    # location = 'us-central1'
+    # composer_environment = 'YOUR_COMPOSER_ENVIRONMENT_NAME'
+
+    environment_url = (
+        'https://composer.googleapis.com/v1beta1/projects/{}/locations/{}'
+        '/environments/{}').format(project_id, location, composer_environment)
+    composer_response = authed_session.request('GET', environment_url)
+    environment_data = composer_response.json()
+    airflow_uri = environment_data['config']['airflowUri']
+
+    # The Composer environment response does not include the IAP client ID.
+    # Make a second, unauthenticated HTTP request to the web server to get the
+    # redirect URI.
+    redirect_response = requests.get(airflow_uri, allow_redirects=False)
+    redirect_location = redirect_response.headers['location']
+
+    # Extract the client_id query parameter from the redirect.
+    parsed = six.moves.urllib.parse.urlparse(redirect_location)
+    query_string = six.moves.urllib.parse.parse_qs(parsed.query)
+    print(query_string['client_id'][0])
+    # [END composer_get_environment_client_id]
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser(
+        description=__doc__,
+        formatter_class=argparse.RawDescriptionHelpFormatter)
+    parser.add_argument('project_id', help='Your Project ID.')
+    parser.add_argument(
+        'location', help='Region of the Cloud Composer environent.')
+    parser.add_argument(
+        'composer_environment', help='Name of the Cloud Composer environent.')
+
+    args = parser.parse_args()
+    get_client_id(
+        args.project_id, args.location, args.composer_environment)

composer/rest/get_client_id_test.py

@@ -0,0 +1,28 @@
+# Copyright 2018 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 os
+
+from .get_client_id import get_client_id
+
+
+PROJECT = os.environ['GOOGLE_CLOUD_PROJECT']
+COMPOSER_LOCATION = os.environ['COMPOSER_LOCATION']
+COMPOSER_ENVIRONMENT = os.environ['COMPOSER_ENVIRONMENT']
+
+
+def test_get_client_id(capsys):
+    get_client_id(PROJECT, COMPOSER_LOCATION, COMPOSER_ENVIRONMENT)
+    out, _ = capsys.readouterr()
+    assert '.apps.googleusercontent.com' in out

composer/rest/requirements.txt

@@ -1,2 +1,3 @@
 google-auth==1.4.1
-requests==2.18.4
+requests==2.18.4
+six==1.11.0

composer/workflows/trigger_response_dag.py

@@ -0,0 +1,45 @@
+# Copyright 2018 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.
+
+"""DAG running in response to a Cloud Storage bucket change."""
+
+# [START composer_trigger_response_dag]
+import datetime
+
+import airflow
+from airflow.operators import bash_operator
+
+
+default_args = {
+    'owner': 'Composer Example',
+    'depends_on_past': False,
+    'email': [''],
+    'email_on_failure': False,
+    'email_on_retry': False,
+    'retries': 1,
+    'retry_delay': datetime.timedelta(minutes=5),
+    'start_date': datetime.datetime(2017, 1, 1),
+}
+
+with airflow.DAG(
+        'composer_sample_trigger_response_dag',
+        default_args=default_args,
+        # Not scheduled, trigger only
+        schedule_interval=None) as dag:
+
+    # Print the dag_run's configuration, which includes information about the
+    # Cloud Storage object change.
+    print_gcs_info = bash_operator.BashOperator(
+        task_id='print_gcs_info', bash_command='echo {{ dag_run.conf }}')
+# [END composer_trigger_response_dag]

composer/workflows/trigger_response_dag_test.py

@@ -0,0 +1,23 @@
+# Copyright 2018 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.
+
+
+def test_dag_import():
+    """Test that the DAG file can be successfully imported.
+
+    This tests that the DAG can be parsed, but does not run it in an Airflow
+    environment. This is a recommended sanity check by the official Airflow
+    docs: https://airflow.incubator.apache.org/tutorial.html#testing
+    """
+    from . import trigger_response_dag  # noqa