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 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()
class TestCloudSqlDatabaseHook(unittest.TestCase):
@mock.patch(
'airflow.contrib.hooks.gcp_sql_hook.CloudSqlDatabaseHook.get_connection'
)
def setUp(self, m):
super().setUp()
self.sql_connection = Connection(
conn_id='my_gcp_sql_connection',
conn_type='gcpcloudsql',
login='******',
password='******',
host='host',
schema='schema',
extra='{"database_type":"postgres", "location":"my_location", '
'"instance":"my_instance", "use_proxy": true, '
'"project_id":"my_project"}')
self.connection = Connection(
conn_id='my_gcp_connection',
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":
"your-gcp-project",
"extra__google_cloud_platform__key_path":
'/var/local/google_cloud_default.json'
}
conn_extra_json = json.dumps(conn_extra)
self.connection.set_extra(conn_extra_json)
m.side_effect = [self.sql_connection, self.connection]
self.db_hook = CloudSqlDatabaseHook(
gcp_cloudsql_conn_id='my_gcp_sql_connection',
gcp_conn_id='my_gcp_connection')
def test_get_sqlproxy_runner(self):
self.db_hook._generate_connection_uri()
sqlproxy_runner = self.db_hook.get_sqlproxy_runner()
self.assertEqual(sqlproxy_runner.gcp_conn_id, self.connection.conn_id)
project = self.sql_connection.extra_dejson['project_id']
location = self.sql_connection.extra_dejson['location']
instance = self.sql_connection.extra_dejson['instance']
instance_spec = "{project}:{location}:{instance}".format(
project=project, location=location, instance=instance)
self.assertEqual(sqlproxy_runner.instance_specification, instance_spec)
