Google Cloud Composer Python Samples

.. image:: https://gstatic.com/cloudssh/images/open-btn.png

:target: https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/GoogleCloudPlatform/python-docs-samples&page=editor&open_in_editor=composer/tools/README.rst

This directory contains samples for Google Cloud Composer. `Google Cloud Composer`_ is a managed Apache Airflow service that helps you create, schedule, monitor and manage workflows. Cloud Composer automation helps you create Airflow environments quickly and use Airflow-native tools, such as the powerful Airflow web interface and command line tools, so you can focus on your workflows and not your infrastructure.

.. _Google Cloud Composer: https://cloud.google.com/composer/docs

Setup

Authentication

This sample requires you to have authentication setup. Refer to the

`Authentication Getting Started Guide`_ for instructions on setting up

credentials for applications.

.. _Authentication Getting Started Guide:

https://cloud.google.com/docs/authentication/getting-started

Install Dependencies

#. Clone python-docs-samples and change directory to the sample directory you want to use.

.. code-block:: bash

#. Install `pip`_ and `virtualenv`_ if you do not already have them. You may want to refer to the `Python Development Environment Setup Guide`_ for Google Cloud Platform for instructions.

.. _Python Development Environment Setup Guide:

https://cloud.google.com/python/setup

#. Create a virtualenv. Samples are compatible with Python 2.7 and 3.4+.

.. code-block:: bash

#. Install the dependencies needed to run the samples.

.. code-block:: bash

.. _pip: https://pip.pypa.io/

.. _virtualenv: https://virtualenv.pypa.io/

Samples

Create a new Composer environment based on an existing environment

.. image:: https://gstatic.com/cloudssh/images/open-btn.png

:target: https://console.cloud.google.com/cloudshell/open?git_repo=https://github.com/GoogleCloudPlatform/python-docs-samples&page=editor&open_in_editor=composer/tools/copy_environment.py,composer/tools/README.rst

To run this sample:

.. code-block:: bash

# This file is used to generate README.rst

product:

name: Google Cloud Composer

short_name: Cloud Composer

url: https://cloud.google.com/composer/docs

description: >

`Google Cloud Composer`_ is a managed Apache Airflow service that helps

you create, schedule, monitor and manage workflows. Cloud Composer

automation helps you create Airflow environments quickly and use

Airflow-native tools, such as the powerful Airflow web interface and

command line tools, so you can focus on your workflows and not your

infrastructure.

setup:

- auth

- install_deps

samples:

- name: Create a new Composer environment based on an existing environment

file: copy_environment.py

show_help: True

cloud_client_library: false

folder: composer/tools

import contextlib

import os

from cryptography import fernet

from googleapiclient import discovery, errors

from kubernetes import client, config

from mysql import connector

import six

from six.moves import configparser

def get_fernet_key(composer_env):

print("Retrieving fernet key for Composer Environment {}.".format(

composer_env.get('name', '')))

gke_cluster_resource = composer_env.get("config", {}).get("gkeCluster")

project_zone_cluster = re.match(

"projects/([^/]*)/zones/([^/]*)/clusters/([^/]*)", gke_cluster_resource

).groups()

tmp_dir_name = None

try:

print("Getting cluster credentials {} to retrieve fernet key.".format(

gke_cluster_resource))

tmp_dir_name = tempfile.mkdtemp()

kubeconfig_file = tmp_dir_name + "/config"

os.environ["KUBECONFIG"] = kubeconfig_file

if subprocess.call(

[

"gcloud",

"container",

"clusters",

"get-credentials",

project_zone_cluster[2],

"--zone",

project_zone_cluster[1],

"--project",

project_zone_cluster[0]

]

):

print("Failed to retrieve cluster credentials: {}.".format(

gke_cluster_resource))

sys.exit(1)

kubernetes_client = client.CoreV1Api(

api_client=config.new_client_from_config(

config_file=kubeconfig_file))

airflow_configmap = kubernetes_client.read_namespaced_config_map(

"airflow-configmap", "default")

config_str = airflow_configmap.data['airflow.cfg']

with contextlib.closing(six.StringIO(config_str)) as config_buffer:

config_parser = configparser.ConfigParser()

config_parser.readfp(config_buffer)

return config_parser.get("core", "fernet_key")

except Exception as exc:

print("Failed to get fernet key for cluster: {}.".format(str(exc)))

sys.exit(1)

finally:

if tmp_dir_name:

shutil.rmtree(tmp_dir_name)

def reencrypt_variables_connections(old_fernet_key_str, new_fernet_key_str):

old_fernet_key = fernet.Fernet(old_fernet_key_str.encode("utf-8"))

new_fernet_key = fernet.Fernet(new_fernet_key_str.encode("utf-8"))

db = connector.connect(

host="127.0.0.1",

user="root",

database="airflow-db",

)

variable_cursor = db.cursor()

variable_cursor.execute("SELECT id, val, is_encrypted FROM variable")

rows = variable_cursor.fetchall()

for row in rows:

id = row[0]

val = row[1]

is_encrypted = row[2]

if is_encrypted:

updated_val = new_fernet_key.encrypt(

old_fernet_key.decrypt(bytes(val))).decode()

variable_cursor.execute(

"UPDATE variable SET val=%s WHERE id=%s", (updated_val, id))

db.commit()

conn_cursor = db.cursor()

conn_cursor.execute(

"SELECT id, password, extra, is_encrypted, is_extra_encrypted FROM "

"connection")

rows = conn_cursor.fetchall()

for row in rows:

id = row[0]

password = row[1]

extra = row[2]

is_encrypted = row[3]

is_extra_encrypted = row[4]

if is_encrypted:

updated_password = new_fernet_key.encrypt(

old_fernet_key.decrypt(bytes(password))).decode()

conn_cursor.execute(

"UPDATE connection SET password=%s WHERE id=%s",

(updated_password, id))

if is_extra_encrypted:

updated_extra = new_fernet_key.encrypt(

old_fernet_key.decrypt(bytes(extra))).decode()

conn_cursor.execute(

"UPDATE connection SET extra=%s WHERE id=%s",

(updated_extra, id))

db.commit()

sql_client,

service_account_key,

project,

instance,

gcs_bucket,

filename,

old_fernet_key,

new_fernet_key

print("Reencrypting variables and connections.")

reencrypt_variables_connections(old_fernet_key, new_fernet_key)

except Exception as exc:

print("Failed to copy database: {}".format(str(exc)))

print("Obtaining fernet keys for Composer Environments.")

old_fernet_key = get_fernet_key(existing_env)

new_fernet_key = get_fernet_key(new_env)

old_fernet_key,

new_fernet_key,

delete_bucket(gcs_db_dump_bucket)