def test_cloudsql_hook_delete_connection_on_exception(
self, get_connections, run, get_connection, delete_connection):
connection = Connection()
connection.parse_from_uri(
"gcpcloudsql://*****:*****@127.0.0.1:3200/testdb?database_type=mysql&"
"project_id=example-project&location=europe-west1&instance=testdb&"
"use_proxy=False")
get_connection.return_value = connection
db_connection = Connection()
db_connection.host = "127.0.0.1"
db_connection.set_extra(
json.dumps({
"project_id": "example-project",
"location": "europe-west1",
"instance": "testdb",
"database_type": "mysql"
}))
get_connections.return_value = [db_connection]
run.side_effect = Exception("Exception when running a query")
operator = CloudSqlQueryOperator(sql=['SELECT * FROM TABLE'],
task_id='task_id')
with self.assertRaises(Exception) as cm:
operator.execute(None)
err = cm.exception
self.assertEqual("Exception when running a query", str(err))
delete_connection.assert_called_once_with()
def add_gcp_connection(ds, **kwargs):
""""Add a airflow connection for GCP"""
new_conn = Connection(
conn_id='bigquery',
conn_type='google_cloud_platform',
)
scopes = [
"https://www.googleapis.com/auth/bigquery",
]
conn_extra = {
"extra__google_cloud_platform__scope":
",".join(scopes),
"extra__google_cloud_platform__project":
os.environ['GOOGLE_CLOUD_PROJECT'],
"extra__google_cloud_platform__key_path":
os.environ['AIRFLOW_CONN_GOOGLE_CLOUD_PLATFORM']
}
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()
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)
def test_cloudsql_hook_delete_connection_on_exception(
self, get_connections, run, get_connection, delete_connection):
connection = Connection()
connection.parse_from_uri(
"gcpcloudsql://*****:*****@127.0.0.1:3200/testdb?database_type=mysql&"
"project_id=example-project&location=europe-west1&instance=testdb&"
"use_proxy=False")
get_connection.return_value = connection
db_connection = Connection()
db_connection.host = "127.0.0.1"
db_connection.set_extra(json.dumps({"project_id": "example-project",
"location": "europe-west1",
"instance": "testdb",
"database_type": "mysql"}))
get_connections.return_value = [db_connection]
run.side_effect = Exception("Exception when running a query")
operator = CloudSqlQueryOperator(
sql=['SELECT * FROM TABLE'],
task_id='task_id'
)
with self.assertRaises(Exception) as cm:
operator.execute(None)
err = cm.exception
self.assertEqual("Exception when running a query", str(err))
delete_connection.assert_called_once_with()
def add_gcp_connection(ds, **kwargs):
""""Add a airflow connection for GCP"""
new_conn = Connection(
conn_id="my_gcp_connection", # TODO: parameterize
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": "wam-bam-258119", # TODO: parameterize
"extra__google_cloud_platform__key_path": "service_account.json", # TODO: parameterize
}
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()
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 set_connection(
self,
conn_id: str,
conn_type: str,
host: Optional[str] = None,
schema: Optional[str] = None,
login: Optional[str] = None,
password: Optional[str] = None,
port: Optional[int] = None,
extra: Optional[Union[str, dict]] = None,
):
assert str(settings.engine.url) == self.sql_alchemy_conn
session = settings.Session()
try:
new_conn = Connection(conn_id=conn_id,
conn_type=conn_type,
host=host,
login=login,
password=password,
schema=schema,
port=port)
if extra is not None:
new_conn.set_extra(
extra if isinstance(extra, str) else json.dumps(extra))
session.add(new_conn)
session.commit()
finally:
session.close()
def add_gcp_connection(ds, **kwargs):
""""Add a airflow connection for GCP"""
json_data = json.loads(open('/usr/local/airflow/gcp.json').read())
new_conn = Connection(
conn_id='AirflowGCPKey',
conn_type='google_cloud_platform',
)
scopes = [
"https://www.googleapis.com/auth/pubsub",
"https://www.googleapis.com/auth/datastore",
"https://www.googleapis.com/auth/bigquery",
"https://www.googleapis.com/auth/devstorage.read_write",
"https://www.googleapis.com/auth/logging.write",
"https://www.googleapis.com/auth/cloud-platform",
]
conn_extra = {
"extra__google_cloud_platform__scope": ",".join(scopes),
"extra__google_cloud_platform__project":
"<insert your project platform>",
"extra__google_cloud_platform__key_path":
'/usr/local/airflow/gcp.json',
"extra__google_cloud_platform__key_json": json_data
}
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()
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(**kwargs):
new_conn = Connection(
conn_id="google_cloud_default",
#conn_id="airflowgcp_connection",
conn_type='google_cloud_platform',
)
extra_field = {
"extra__google_cloud_platform__scope":
"https://www.googleapis.com/auth/cloud-platform",
"extra__google_cloud_platform__project":
"airflowgcp",
"extra__google_cloud_platform__key_path":
'/usr/local/airflow/dags/gcpLicense.json'
}
session = settings.Session()
#checking if connection exist
if session.query(Connection).filter(
Connection.conn_id == new_conn.conn_id).first():
my_connection = session.query(Connection).filter(
Connection.conn_id == new_conn.conn_id).one()
my_connection.set_extra(json.dumps(extra_field))
session.add(my_connection)
session.commit()
else: #if it doesn't exit create one
new_conn.set_extra(json.dumps(extra_field))
session.add(new_conn)
session.commit()
def add_gcp_connection(**kwargs):
new_conn = Connection(
conn_id="google_cloud_default",
conn_type='google_cloud_platform',
)
extra_field = {
"extra__google_cloud_platform__scope":
"https://www.googleapis.com/auth/cloud-platform",
"extra__google_cloud_platform__project": "task-ts",
"extra__google_cloud_platform__key_path": os.environ["GCP_KEY_PATH"]
}
session = settings.Session()
#checking if connection exist
if session.query(Connection).filter(
Connection.conn_id == new_conn.conn_id).first():
my_connection = session.query(Connection).filter(
Connection.conn_id == new_conn.conn_id).one()
my_connection.set_extra(json.dumps(extra_field))
session.add(my_connection)
session.commit()
else: #if it doesn't exit create one
new_conn.set_extra(json.dumps(extra_field))
session.add(new_conn)
session.commit()
def add_connection(
conn_id,
uri=None,
conn_type=None,
host=None,
login=None,
password=None,
port=None,
schema=None,
extra=None,
):
if uri:
new_conn = Connection(conn_id=conn_id, uri=uri)
else:
new_conn = Connection(
conn_id=conn_id,
conn_type=conn_type,
host=host,
login=login,
password=password,
schema=schema,
port=port,
)
if extra is not None:
new_conn.set_extra(extra)
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\tSuccessfully added `conn_id`={conn_id} : {uri}\n"
msg = msg.format(
conn_id=new_conn.conn_id,
uri=uri or urlunparse((
conn_type,
"{login}:{password}@{host}:{port}".format(
login=login or "",
password=password or "",
host=host or "",
port=port or "",
),
schema or "",
"",
"",
"",
)),
)
logger.info(msg)
logger.info("extra=%s" % str(extra))
else:
msg = "\n\tA connection with `conn_id`={conn_id} already exists\n"
raise DatabandRuntimeError(msg.format(conn_id=new_conn.conn_id))
def connections_add(args):
"""Adds new connection"""
# Check that the conn_id and conn_uri args were passed to the command:
missing_args = []
invalid_args = []
if args.conn_uri:
for arg in alternative_conn_specs:
if getattr(args, arg) is not None:
invalid_args.append(arg)
elif not args.conn_type:
missing_args.append('conn-uri or conn-type')
if missing_args:
msg = ('The following args are required to add a connection:' +
' {missing!r}'.format(missing=missing_args))
raise SystemExit(msg)
if invalid_args:
msg = ('The following args are not compatible with the ' +
'add flag and --conn-uri flag: {invalid!r}')
msg = msg.format(invalid=invalid_args)
raise SystemExit(msg)
if args.conn_uri:
new_conn = Connection(conn_id=args.conn_id, uri=args.conn_uri)
else:
new_conn = Connection(conn_id=args.conn_id,
conn_type=args.conn_type,
host=args.conn_host,
login=args.conn_login,
password=args.conn_password,
schema=args.conn_schema,
port=args.conn_port)
if args.conn_extra is not None:
new_conn.set_extra(args.conn_extra)
with create_session() as session:
if not (session.query(Connection)
.filter(Connection.conn_id == new_conn.conn_id).first()):
session.add(new_conn)
msg = '\n\tSuccessfully added `conn_id`={conn_id} : {uri}\n'
msg = msg.format(conn_id=new_conn.conn_id,
uri=args.conn_uri or
urlunparse((args.conn_type,
'{login}:{password}@{host}:{port}'
.format(login=args.conn_login or '',
password='******' if args.conn_password else '',
host=args.conn_host or '',
port=args.conn_port or ''),
args.conn_schema or '', '', '', '')))
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 connections_add(args):
"""Adds new connection"""
# Check that the conn_id and conn_uri args were passed to the command:
missing_args = []
invalid_args = []
if args.conn_uri:
if not _valid_uri(args.conn_uri):
raise SystemExit(f'The URI provided to --conn-uri is invalid: {args.conn_uri}')
for arg in alternative_conn_specs:
if getattr(args, arg) is not None:
invalid_args.append(arg)
elif not args.conn_type:
missing_args.append('conn-uri or conn-type')
if missing_args:
raise SystemExit(f'The following args are required to add a connection: {missing_args!r}')
if invalid_args:
raise SystemExit(
f'The following args are not compatible with the '
f'add flag and --conn-uri flag: {invalid_args!r}'
)
if args.conn_uri:
new_conn = Connection(conn_id=args.conn_id, description=args.conn_description, uri=args.conn_uri)
else:
new_conn = Connection(
conn_id=args.conn_id,
conn_type=args.conn_type,
description=args.conn_description,
host=args.conn_host,
login=args.conn_login,
password=args.conn_password,
schema=args.conn_schema,
port=args.conn_port,
)
if args.conn_extra is not None:
new_conn.set_extra(args.conn_extra)
with create_session() as session:
if not session.query(Connection).filter(Connection.conn_id == new_conn.conn_id).first():
session.add(new_conn)
msg = 'Successfully added `conn_id`={conn_id} : {uri}'
msg = msg.format(
conn_id=new_conn.conn_id,
uri=args.conn_uri
or urlunparse(
(
args.conn_type,
'{login}:{password}@{host}:{port}'.format(
login=args.conn_login or '',
password='******' if args.conn_password else '',
host=args.conn_host or '',
port=args.conn_port or '',
),
args.conn_schema or '',
'',
'',
'',
)
),
)
print(msg)
else:
msg = f'A connection with `conn_id`={new_conn.conn_id} already exists.'
raise SystemExit(msg)
def connections_add(args):
"""Adds new connection"""
has_uri = bool(args.conn_uri)
has_json = bool(args.conn_json)
has_type = bool(args.conn_type)
if not has_type and not (has_json or has_uri):
raise SystemExit(
'Must supply either conn-uri or conn-json if not supplying conn-type'
)
if has_json and has_uri:
raise SystemExit('Cannot supply both conn-uri and conn-json')
if has_type and not (args.conn_type in _get_connection_types()):
warnings.warn(
f'The type provided to --conn-type is invalid: {args.conn_type}')
warnings.warn(f'Supported --conn-types are:{_get_connection_types()}.'
'Hence overriding the conn-type with generic')
args.conn_type = 'generic'
if has_uri or has_json:
invalid_args = []
if has_uri and not _valid_uri(args.conn_uri):
raise SystemExit(
f'The URI provided to --conn-uri is invalid: {args.conn_uri}')
for arg in alternative_conn_specs:
if getattr(args, arg) is not None:
invalid_args.append(arg)
if has_json and args.conn_extra:
invalid_args.append("--conn-extra")
if invalid_args:
raise SystemExit(
"The following args are not compatible with "
f"the --conn-{'uri' if has_uri else 'json'} flag: {invalid_args!r}"
)
if args.conn_uri:
new_conn = Connection(conn_id=args.conn_id,
description=args.conn_description,
uri=args.conn_uri)
if args.conn_extra is not None:
new_conn.set_extra(args.conn_extra)
elif args.conn_json:
new_conn = Connection.from_json(conn_id=args.conn_id,
value=args.conn_json)
if not new_conn.conn_type:
raise SystemExit('conn-json is invalid; must supply conn-type')
else:
new_conn = Connection(
conn_id=args.conn_id,
conn_type=args.conn_type,
description=args.conn_description,
host=args.conn_host,
login=args.conn_login,
password=args.conn_password,
schema=args.conn_schema,
port=args.conn_port,
)
if args.conn_extra is not None:
new_conn.set_extra(args.conn_extra)
with create_session() as session:
if not session.query(Connection).filter(
Connection.conn_id == new_conn.conn_id).first():
session.add(new_conn)
msg = 'Successfully added `conn_id`={conn_id} : {uri}'
msg = msg.format(
conn_id=new_conn.conn_id,
uri=args.conn_uri or urlunparse((
new_conn.conn_type,
f"{new_conn.login or ''}:{'******' if new_conn.password else ''}"
f"@{new_conn.host or ''}:{new_conn.port or ''}",
new_conn.schema or '',
'',
'',
'',
)),
)
print(msg)
else:
msg = f'A connection with `conn_id`={new_conn.conn_id} already exists.'
raise SystemExit(msg)
