Skip to content

Commit 8ea51a5

Browse files
leahecoleleahecole
authored
Small updates to samples for Composer v2 (GoogleCloudPlatform#6923)
* Add comment about pod affinity * Add link to comment * Add helpful error message to get_client_id * add note about compatiblity * add 3.10 exemption to noxfile config * add test for get client id raising error * make long line more readable * stop whining, lint
1 parent 9c62a1f commit 8ea51a5

File tree

5 files changed

+51
-19
lines changed

5 files changed

+51
-19
lines changed

‎composer/functions/composer_storage_trigger.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -37,6 +37,8 @@ def trigger_dag(data, context=None):
3737
3838
To call this function from a Python script, omit the ``context`` argument
3939
and pass in a non-null value for the ``data`` argument.
40+
41+
This function is currently only compatible with Composer v1 environments.
4042
"""
4143

4244
# Fill in with your Composer info here

‎composer/rest/get_client_id.py

Lines changed: 30 additions & 18 deletions
Original file line numberDiff line numberDiff line change
@@ -19,6 +19,11 @@
1919

2020
def get_client_id(project_id, location, composer_environment):
2121
# [START composer_get_environment_client_id]
22+
# This script is intended to be used with Composer 1 environments
23+
# In Composer 2, the Airflow Webserver is not in the tenant project
24+
# so there is no tenant client ID
25+
# See https://cloud.google.com/composer/docs/composer-2/environment-architecture
26+
# for more details
2227
import google.auth
2328
import google.auth.transport.requests
2429
import requests
@@ -27,45 +32,52 @@ def get_client_id(project_id, location, composer_environment):
2732
# Authenticate with Google Cloud.
2833
# See: https://cloud.google.com/docs/authentication/getting-started
2934
credentials, _ = google.auth.default(
30-
scopes=['https://www.googleapis.com/auth/cloud-platform'])
31-
authed_session = google.auth.transport.requests.AuthorizedSession(
32-
credentials)
35+
scopes=["https://www.googleapis.com/auth/cloud-platform"]
36+
)
37+
authed_session = google.auth.transport.requests.AuthorizedSession(credentials)
3338

3439
# project_id = 'YOUR_PROJECT_ID'
3540
# location = 'us-central1'
3641
# composer_environment = 'YOUR_COMPOSER_ENVIRONMENT_NAME'
3742

3843
environment_url = (
39-
'https://composer.googleapis.com/v1beta1/projects/{}/locations/{}'
40-
'/environments/{}').format(project_id, location, composer_environment)
41-
composer_response = authed_session.request('GET', environment_url)
44+
"https://composer.googleapis.com/v1beta1/projects/{}/locations/{}"
45+
"/environments/{}"
46+
).format(project_id, location, composer_environment)
47+
composer_response = authed_session.request("GET", environment_url)
4248
environment_data = composer_response.json()
43-
airflow_uri = environment_data['config']['airflowUri']
49+
composer_version = environment_data["config"]["softwareConfig"]["imageVersion"]
50+
if "composer-1" not in composer_version:
51+
version_error = ("This script is intended to be used with Composer 1 environments. "
52+
"In Composer 2, the Airflow Webserver is not in the tenant project, "
53+
"so there is no tenant client ID. "
54+
"See https://cloud.google.com/composer/docs/composer-2/environment-architecture for more details.")
55+
raise (RuntimeError(version_error))
56+
airflow_uri = environment_data["config"]["airflowUri"]
4457

4558
# The Composer environment response does not include the IAP client ID.
4659
# Make a second, unauthenticated HTTP request to the web server to get the
4760
# redirect URI.
4861
redirect_response = requests.get(airflow_uri, allow_redirects=False)
49-
redirect_location = redirect_response.headers['location']
62+
redirect_location = redirect_response.headers["location"]
5063

5164
# Extract the client_id query parameter from the redirect.
5265
parsed = six.moves.urllib.parse.urlparse(redirect_location)
5366
query_string = six.moves.urllib.parse.parse_qs(parsed.query)
54-
print(query_string['client_id'][0])
67+
print(query_string["client_id"][0])
5568
# [END composer_get_environment_client_id]
5669

5770

5871
# Usage: python get_client_id.py your_project_id your_region your_environment_name
59-
if __name__ == '__main__':
72+
if __name__ == "__main__":
6073
parser = argparse.ArgumentParser(
61-
description=__doc__,
62-
formatter_class=argparse.RawDescriptionHelpFormatter)
63-
parser.add_argument('project_id', help='Your Project ID.')
74+
description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
75+
)
76+
parser.add_argument("project_id", help="Your Project ID.")
77+
parser.add_argument("location", help="Region of the Cloud Composer environment.")
6478
parser.add_argument(
65-
'location', help='Region of the Cloud Composer environment.')
66-
parser.add_argument(
67-
'composer_environment', help='Name of the Cloud Composer environment.')
79+
"composer_environment", help="Name of the Cloud Composer environment."
80+
)
6881

6982
args = parser.parse_args()
70-
get_client_id(
71-
args.project_id, args.location, args.composer_environment)
83+
get_client_id(args.project_id, args.location, args.composer_environment)

‎composer/rest/get_client_id_test.py

Lines changed: 10 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -14,15 +14,25 @@
1414

1515
import os
1616

17+
import pytest
18+
1719
from .get_client_id import get_client_id
1820

1921

2022
PROJECT = os.environ['GOOGLE_CLOUD_PROJECT']
2123
COMPOSER_LOCATION = os.environ['COMPOSER_LOCATION']
2224
COMPOSER_ENVIRONMENT = os.environ['COMPOSER_ENVIRONMENT']
25+
COMPOSER2_ENVIRONMENT = os.environ['COMPOSER2_ENVIRONMENT']
2326

2427

2528
def test_get_client_id(capsys):
2629
get_client_id(PROJECT, COMPOSER_LOCATION, COMPOSER_ENVIRONMENT)
2730
out, _ = capsys.readouterr()
2831
assert '.apps.googleusercontent.com' in out
32+
33+
34+
def test_get_client_id_composer_2(capsys):
35+
with pytest.raises(RuntimeError):
36+
get_client_id(PROJECT, COMPOSER_LOCATION, COMPOSER2_ENVIRONMENT)
37+
out, _ = capsys.readouterr()
38+
assert 'This script is intended to be used with Composer 1' in out

‎composer/workflows/kubernetes_pod_operator.py

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -135,6 +135,10 @@
135135
'GOOGLE_APPLICATION_CREDENTIALS': '/var/secrets/google/service-account.json '})
136136
# [END composer_kubernetespodoperator_secretconfig]
137137
# [START composer_kubernetespodaffinity]
138+
# Pod affinity with the KubernetesPodOperator
139+
# is not supported with Composer 2
140+
# instead, create a cluster and use the GKEStartPodOperator
141+
# https://cloud.google.com/composer/docs/using-gke-operator
138142
kubernetes_affinity_ex = KubernetesPodOperator(
139143
task_id='ex-pod-affinity',
140144
name='ex-pod-affinity',
@@ -231,6 +235,10 @@
231235
# Affinity determines which nodes the Pod can run on based on the
232236
# config. For more information see:
233237
# https://kubernetes.io/docs/concepts/configuration/assign-pod-node/
238+
# Pod affinity with the KubernetesPodOperator
239+
# is not supported with Composer 2
240+
# instead, create a cluster and use the GKEStartPodOperator
241+
# https://cloud.google.com/composer/docs/using-gke-operator
234242
affinity={})
235243
# [END composer_kubernetespodoperator_fullconfig]
236244
# [END composer_kubernetespodoperator]

‎composer/workflows/noxfile_config.py

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -30,7 +30,7 @@
3030

3131
TEST_CONFIG_OVERRIDE = {
3232
# You can opt out from the test for specific Python versions.
33-
'ignored_versions': ["2.7", "3.6", "3.7", "3.9"], # Composer w/ Airflow 2 only supports Python 3.8
33+
'ignored_versions': ["2.7", "3.6", "3.7", "3.9", "3.10"], # Composer w/ Airflow 2 only supports Python 3.8
3434
# Old samples are opted out of enforcing Python type hints
3535
# All new samples should feature them
3636
"enforce_type_hints": False,

0 commit comments

Comments
 (0)