get_dag_prefix.py

# 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 a Cloud Composer environment via the REST API.

This code sample gets a Cloud Composer environment resource and prints the
Cloud Storage path used to store Apache Airflow DAGs.
"""

import argparse


def get_dag_prefix(project_id, location, composer_environment):
    # [START composer_get_environment_dag_prefix]
    import google.auth
    import google.auth.transport.requests

    # 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/%7B%7D/locations/%7B%7D'
        '/environments/{}').format(project_id, location, composer_environment)
    response = authed_session.request('GET', environment_url)
    environment_data = response.json()

    # Print the bucket name from the response body.
    print(environment_data['config']['dagGcsPrefix'])
    # [END composer_get_environment_dag_prefix]


# Usage: python get_dag_prefix.py your_project_id your_region your_environment_name
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 environment.')
    parser.add_argument(
        'composer_environment', help='Name of the Cloud Composer environment.')

    args = parser.parse_args()
    get_dag_prefix(
        args.project_id, args.location, args.composer_environment)