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)
def _execute_query(self, hook: CloudSqlDatabaseHook, database_hook: Union[PostgresHook, MySqlHook]):
cloud_sql_proxy_runner = None
try:
if hook.use_proxy:
cloud_sql_proxy_runner = hook.get_sqlproxy_runner()
hook.free_reserved_port()
# There is very, very slim chance that the socket will
# be taken over here by another bind(0).
# It's quite unlikely to happen though!
cloud_sql_proxy_runner.start_proxy()
self.log.info('Executing: "%s"', self.sql)
database_hook.run(self.sql, self.autocommit, parameters=self.parameters)
finally:
if cloud_sql_proxy_runner:
cloud_sql_proxy_runner.stop_proxy()