Skip to content

Commit cf12171

Browse files
leahecoletswast
andauthored
Gcf composer trigger (#2415)
* WIP: add code for GCF composer * WIP: update requirements txt, fix some lint stuff * Fix most linting things * Add region tag * Remove unused variable * Fix some nits * Update functions/composer/composer_storage_trigger.py fix nit Co-Authored-By: Tim Swast <[email protected]> * Apply suggestions from code review Co-Authored-By: Tim Swast <[email protected]> * Fix lint * fix lint * Remove app_engine Auth bits * Update functions/composer/composer_storage_trigger.py Co-Authored-By: Tim Swast <[email protected]>
1 parent a83a307 commit cf12171

File tree

3 files changed

+216
-0
lines changed

3 files changed

+216
-0
lines changed
Lines changed: 180 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,180 @@
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]
Lines changed: 34 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,34 @@
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+
import pytest
16+
import mock
17+
import composer_storage_trigger
18+
19+
# handles error in JSON body
20+
@mock.patch('composer_storage_trigger.make_iap_request',
21+
side_effect=Exception('Bad request: JSON body error'))
22+
def test_json_body_error(make_iap_request_mock):
23+
# Pass None, an input that is not valid JSON
24+
trigger_event = None
25+
with pytest.raises(Exception):
26+
composer_storage_trigger.trigger_dag(trigger_event)
27+
28+
# handles error in IAP response
29+
@mock.patch('composer_storage_trigger.make_iap_request',
30+
side_effect=Exception('Error in IAP response: unauthorized'))
31+
def test_iap_response_error(make_iap_request_mock):
32+
trigger_event = {'file': 'some-gcs-file'}
33+
with pytest.raises(Exception):
34+
composer_storage_trigger.trigger_dag(trigger_event)

functions/composer/requirements.txt

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,2 @@
1+
requests_toolbelt==0.9.1
2+
google-auth==1.6.2

0 commit comments

Comments
 (0)