|
| 1 | +# Copyright 2019 Google LLC |
| 2 | + |
| 3 | +# Licensed under the Apache License, Version 2.0 (the "License"); |
| 4 | +# you may not use this file except in compliance with the License. |
| 5 | +# You may obtain a copy of the License at |
| 6 | + |
| 7 | +# https://www.apache.org/licenses/LICENSE-2.0 |
| 8 | + |
| 9 | +# Unless required by applicable law or agreed to in writing, software |
| 10 | +# distributed under the License is distributed on an "AS IS" BASIS, |
| 11 | +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. |
| 12 | +# See the License for the specific language governing permissions and |
| 13 | +# limitations under the License. |
| 14 | + |
| 15 | +# [START composer_trigger] |
| 16 | + |
| 17 | +import google.auth |
| 18 | +import google.auth.compute_engine.credentials |
| 19 | +import google.auth.iam |
| 20 | +from google.auth.transport.requests import Request |
| 21 | +import google.oauth2.credentials |
| 22 | +import google.oauth2.service_account |
| 23 | +import requests |
| 24 | + |
| 25 | + |
| 26 | +IAM_SCOPE = 'https://www.googleapis.com/auth/iam' |
| 27 | +OAUTH_TOKEN_URI = 'https://www.googleapis.com/oauth2/v4/token' |
| 28 | + |
| 29 | + |
| 30 | +def trigger_dag(data, context=None): |
| 31 | + """Makes a POST request to the Composer DAG Trigger API |
| 32 | +
|
| 33 | + When called via Google Cloud Functions (GCF), |
| 34 | + data and context are Background function parameters. |
| 35 | +
|
| 36 | + For more info, refer to |
| 37 | + https://cloud.google.com/functions/docs/writing/background#functions_background_parameters-python |
| 38 | +
|
| 39 | + To call this function from a Python script, omit the ``context`` argument |
| 40 | + and pass in a non-null value for the ``data`` argument. |
| 41 | + """ |
| 42 | + |
| 43 | + # Fill in with your Composer info here |
| 44 | + # Navigate to your webserver's login page and get this from the URL |
| 45 | + # Or use the script found at |
| 46 | + # https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/composer/rest/get_client_id.py |
| 47 | + client_id = 'YOUR-CLIENT-ID' |
| 48 | + # This should be part of your webserver's URL: |
| 49 | + # {tenant-project-id}.appspot.com |
| 50 | + webserver_id = 'YOUR-TENANT-PROJECT' |
| 51 | + # The name of the DAG you wish to trigger |
| 52 | + dag_name = 'composer_sample_trigger_response_dag' |
| 53 | + webserver_url = ( |
| 54 | + 'https://' |
| 55 | + + webserver_id |
| 56 | + + '.appspot.com/api/experimental/dags/' |
| 57 | + + dag_name |
| 58 | + + '/dag_runs' |
| 59 | + ) |
| 60 | + # Make a POST request to IAP which then Triggers the DAG |
| 61 | + make_iap_request(webserver_url, client_id, method='POST', json=data) |
| 62 | + |
| 63 | + |
| 64 | +# This code is copied from |
| 65 | +# https://github.com/GoogleCloudPlatform/python-docs-samples/blob/master/iap/make_iap_request.py |
| 66 | +# START COPIED IAP CODE |
| 67 | +def make_iap_request(url, client_id, method='GET', **kwargs): |
| 68 | + """Makes a request to an application protected by Identity-Aware Proxy. |
| 69 | +
|
| 70 | + Args: |
| 71 | + url: The Identity-Aware Proxy-protected URL to fetch. |
| 72 | + client_id: The client ID used by Identity-Aware Proxy. |
| 73 | + method: The request method to use |
| 74 | + ('GET', 'OPTIONS', 'HEAD', 'POST', 'PUT', 'PATCH', 'DELETE') |
| 75 | + **kwargs: Any of the parameters defined for the request function: |
| 76 | + https://github.com/requests/requests/blob/master/requests/api.py |
| 77 | + If no timeout is provided, it is set to 90 by default. |
| 78 | +
|
| 79 | + Returns: |
| 80 | + The page body, or raises an exception if the page couldn't be retrieved. |
| 81 | + """ |
| 82 | + # Set the default timeout, if missing |
| 83 | + if 'timeout' not in kwargs: |
| 84 | + kwargs['timeout'] = 90 |
| 85 | + |
| 86 | + # Figure out what environment we're running in and get some preliminary |
| 87 | + # information about the service account. |
| 88 | + bootstrap_credentials, _ = google.auth.default( |
| 89 | + scopes=[IAM_SCOPE]) |
| 90 | + |
| 91 | + # For service account's using the Compute Engine metadata service, |
| 92 | + # service_account_email isn't available until refresh is called. |
| 93 | + bootstrap_credentials.refresh(Request()) |
| 94 | + |
| 95 | + signer_email = bootstrap_credentials.service_account_email |
| 96 | + if isinstance(bootstrap_credentials, |
| 97 | + google.auth.compute_engine.credentials.Credentials): |
| 98 | + # Since the Compute Engine metadata service doesn't expose the service |
| 99 | + # account key, we use the IAM signBlob API to sign instead. |
| 100 | + # In order for this to work: |
| 101 | + # 1. Your VM needs the https://www.googleapis.com/auth/iam scope. |
| 102 | + # You can specify this specific scope when creating a VM |
| 103 | + # through the API or gcloud. When using Cloud Console, |
| 104 | + # you'll need to specify the "full access to all Cloud APIs" |
| 105 | + # scope. A VM's scopes can only be specified at creation time. |
| 106 | + # 2. The VM's default service account needs the "Service Account Actor" |
| 107 | + # role. This can be found under the "Project" category in Cloud |
| 108 | + # Console, or roles/iam.serviceAccountActor in gcloud. |
| 109 | + signer = google.auth.iam.Signer( |
| 110 | + Request(), bootstrap_credentials, signer_email) |
| 111 | + else: |
| 112 | + # A Signer object can sign a JWT using the service account's key. |
| 113 | + signer = bootstrap_credentials.signer |
| 114 | + |
| 115 | + # Construct OAuth 2.0 service account credentials using the signer |
| 116 | + # and email acquired from the bootstrap credentials. |
| 117 | + service_account_credentials = google.oauth2.service_account.Credentials( |
| 118 | + signer, signer_email, token_uri=OAUTH_TOKEN_URI, additional_claims={ |
| 119 | + 'target_audience': client_id |
| 120 | + }) |
| 121 | + # service_account_credentials gives us a JWT signed by the service |
| 122 | + # account. Next, we use that to obtain an OpenID Connect token, |
| 123 | + # which is a JWT signed by Google. |
| 124 | + google_open_id_connect_token = get_google_open_id_connect_token( |
| 125 | + service_account_credentials) |
| 126 | + |
| 127 | + # Fetch the Identity-Aware Proxy-protected URL, including an |
| 128 | + # Authorization header containing "Bearer " followed by a |
| 129 | + # Google-issued OpenID Connect token for the service account. |
| 130 | + resp = requests.request( |
| 131 | + method, url, |
| 132 | + headers={'Authorization': 'Bearer {}'.format( |
| 133 | + google_open_id_connect_token)}, **kwargs) |
| 134 | + if resp.status_code == 403: |
| 135 | + raise Exception('Service account {} does not have permission to ' |
| 136 | + 'access the IAP-protected application.'.format( |
| 137 | + signer_email)) |
| 138 | + elif resp.status_code != 200: |
| 139 | + raise Exception( |
| 140 | + 'Bad response from application: {!r} / {!r} / {!r}'.format( |
| 141 | + resp.status_code, resp.headers, resp.text)) |
| 142 | + else: |
| 143 | + return resp.text |
| 144 | + |
| 145 | + |
| 146 | +def get_google_open_id_connect_token(service_account_credentials): |
| 147 | + """Get an OpenID Connect token issued by Google for the service account. |
| 148 | +
|
| 149 | + This function: |
| 150 | +
|
| 151 | + 1. Generates a JWT signed with the service account's private key |
| 152 | + containing a special "target_audience" claim. |
| 153 | +
|
| 154 | + 2. Sends it to the OAUTH_TOKEN_URI endpoint. Because the JWT in #1 |
| 155 | + has a target_audience claim, that endpoint will respond with |
| 156 | + an OpenID Connect token for the service account -- in other words, |
| 157 | + a JWT signed by *Google*. The aud claim in this JWT will be |
| 158 | + set to the value from the target_audience claim in #1. |
| 159 | +
|
| 160 | + For more information, see |
| 161 | + https://developers.google.com/identity/protocols/OAuth2ServiceAccount . |
| 162 | + The HTTP/REST example on that page describes the JWT structure and |
| 163 | + demonstrates how to call the token endpoint. (The example on that page |
| 164 | + shows how to get an OAuth2 access token; this code is using a |
| 165 | + modified version of it to get an OpenID Connect token.) |
| 166 | + """ |
| 167 | + |
| 168 | + service_account_jwt = ( |
| 169 | + service_account_credentials._make_authorization_grant_assertion()) |
| 170 | + request = google.auth.transport.requests.Request() |
| 171 | + body = { |
| 172 | + 'assertion': service_account_jwt, |
| 173 | + 'grant_type': google.oauth2._client._JWT_GRANT_TYPE, |
| 174 | + } |
| 175 | + token_response = google.oauth2._client._token_endpoint_request( |
| 176 | + request, OAUTH_TOKEN_URI, body) |
| 177 | + return token_response['id_token'] |
| 178 | +# END COPIED IAP CODE |
| 179 | + |
| 180 | +# [END composer_trigger] |
0 commit comments