def add_docker_connection(ds, **kwargs):
""""Add a airflow connection for google container registry"""
new_conn = Connection(
conn_id=CONN_PARAMS_DICT.get("gcr_conn_id"),
conn_type="docker",
host="gcr.io/" + CONN_PARAMS_DICT.get("gcp_project"),
login="******",
)
# save contents of service account key into encrypted password field
json_key = get_secret(
project_name=CONN_PARAMS_DICT.get("gcp_project"),
secret_name=CONN_PARAMS_DICT.get("secret_name"),
)
data = str(json.loads(json_key))
new_conn.set_password(data)
session = settings.Session()
if not (session.query(Connection).filter(Connection.conn_id == new_conn.conn_id).first()):
session.add(new_conn)
session.commit()
msg = "\n\tA connection with `conn_id`={conn_id} is newly created\n"
msg = msg.format(conn_id=new_conn.conn_id)
print(msg)
else:
msg = "\n\tA connection with `conn_id`={conn_id} already exists\n"
msg = msg.format(conn_id=new_conn.conn_id)
print(msg)
def add_gcp_connection(ds, **kwargs):
""""Add a airflow connection for GCP"""
new_conn = Connection(
conn_id=CONN_PARAMS_DICT["gcp_conn_id"], conn_type="google_cloud_platform",
)
scopes = [
"https://www.googleapis.com/auth/cloud-platform",
]
conn_extra = {
"extra__google_cloud_platform__scope": ",".join(scopes),
"extra__google_cloud_platform__project": CONN_PARAMS_DICT.get("gcp_project"),
"extra__google_cloud_platform__keyfile_dict": get_secret(
project_name=CONN_PARAMS_DICT.get("gcp_project"),
secret_name=CONN_PARAMS_DICT.get("secret_name"),
),
}
conn_extra_json = json.dumps(conn_extra)
new_conn.set_extra(conn_extra_json)
session = settings.Session()
if not (session.query(Connection).filter(Connection.conn_id == new_conn.conn_id).first()):
session.add(new_conn)
session.commit()
msg = "\n\tA connection with `conn_id`={conn_id} is newly created\n"
msg = msg.format(conn_id=new_conn.conn_id)
print(msg)
else:
msg = "\n\tA connection with `conn_id`={conn_id} already exists\n"
msg = msg.format(conn_id=new_conn.conn_id)
print(msg)
