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()
class CloudSqlQueryOperator(BaseOperator):
"""
Performs DML or DDL query on an existing Cloud Sql instance. It optionally uses
cloud-sql-proxy to establish secure connection with the database.
.. seealso::
For more information on how to use this operator, take a look at the guide:
:ref:`howto/operator:CloudSqlQueryOperator`
:param sql: SQL query or list of queries to run (should be DML or DDL query -
this operator does not return any data from the database,
so it is useless to pass it DQL queries. Note that it is responsibility of the
author of the queries to make sure that the queries are idempotent. For example
you can use CREATE TABLE IF NOT EXISTS to create a table.
:type sql: str or list[str]
:param parameters: (optional) the parameters to render the SQL query with.
:type parameters: dict or iterable
:param autocommit: if True, each command is automatically committed.
(default value: False)
:type autocommit: bool
:param gcp_conn_id: The connection ID used to connect to Google Cloud Platform for
cloud-sql-proxy authentication.
:type gcp_conn_id: str
:param gcp_cloudsql_conn_id: The connection ID used to connect to Google Cloud SQL
its schema should be gcpcloudsql://.
See :class:`~airflow.contrib.hooks.gcp_sql_hook.CloudSqlDatabaseHook` for
details on how to define gcpcloudsql:// connection.
:type gcp_cloudsql_conn_id: str
"""
# [START gcp_sql_query_template_fields]
template_fields = ('sql', 'gcp_cloudsql_conn_id', 'gcp_conn_id')
template_ext = ('.sql', )
# [END gcp_sql_query_template_fields]
@apply_defaults
def __init__(self,
sql: Union[List[str], str],
autocommit: bool = False,
parameters: Optional[Union[Dict, Iterable]] = None,
gcp_conn_id: str = 'google_cloud_default',
gcp_cloudsql_conn_id: str = 'google_cloud_sql_default',
*args,
**kwargs) -> None:
super().__init__(*args, **kwargs)
self.sql = sql
self.gcp_conn_id = gcp_conn_id
self.gcp_cloudsql_conn_id = gcp_cloudsql_conn_id
self.autocommit = autocommit
self.parameters = parameters
self.gcp_connection = BaseHook.get_connection(self.gcp_conn_id)
self.cloudsql_db_hook = CloudSqlDatabaseHook(
gcp_cloudsql_conn_id=gcp_cloudsql_conn_id,
gcp_conn_id=gcp_conn_id,
default_gcp_project_id=self.gcp_connection.extra_dejson.get(
'extra__google_cloud_platform__project'))
self.cloud_sql_proxy_runner = None
self.database_hook = None
def _execute_query(self):
try:
if self.cloudsql_db_hook.use_proxy:
self.cloud_sql_proxy_runner = self.cloudsql_db_hook. \
get_sqlproxy_runner()
self.cloudsql_db_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!
self.cloud_sql_proxy_runner.start_proxy()
self.log.info('Executing: "%s"', self.sql)
self.database_hook.run(self.sql,
self.autocommit,
parameters=self.parameters)
finally:
if self.cloud_sql_proxy_runner:
self.cloud_sql_proxy_runner.stop_proxy()
self.cloud_sql_proxy_runner = None
def execute(self, context):
self.cloudsql_db_hook.validate_ssl_certs()
self.cloudsql_db_hook.create_connection()
try:
self.cloudsql_db_hook.validate_socket_path_length()
self.database_hook = self.cloudsql_db_hook.get_database_hook()
try:
self._execute_query()
finally:
self.cloudsql_db_hook.cleanup_database_hook()
finally:
self.cloudsql_db_hook.delete_connection()
self.cloudsql_db_hook = None