Move cloud_sql_binary_path from connection to Hook (#29499)

Specifying cloud_sql_binary_path in connection should never be needed, This is at most property of the Hook (and Operator by transition) if you want to override it, rather than extra in the connection.
apache · Feb 13, 2023 · 32c571e · 32c571e
1 parent 8e24387
commit 32c571e
Show file tree

Hide file tree

Showing 3 changed files with 14 additions and 6 deletions.
diff --git a/airflow/providers/google/cloud/example_dags/example_cloud_sql_query.py b/airflow/providers/google/cloud/example_dags/example_cloud_sql_query.py
@@ -205,8 +205,7 @@ def get_absolute_path(path):
  "location={location}&"
  "instance={instance}&"
  "use_proxy=True&"
- "sql_proxy_binary_path={sql_proxy_binary_path}&"
- "sql_proxy_use_tcp=False".format(sql_proxy_binary_path=quote_plus(sql_proxy_binary_path), **mysql_kwargs)
+ "sql_proxy_use_tcp=False".format(**mysql_kwargs)
 )
 
 # MySQL: connect directly via TCP (non-SSL)
@@ -279,7 +278,10 @@ def get_absolute_path(path):
 
  for connection_name in connection_names:
  task = CloudSQLExecuteQueryOperator(
- gcp_cloudsql_conn_id=connection_name, task_id="example_gcp_sql_task_" + connection_name, sql=SQL
+ gcp_cloudsql_conn_id=connection_name,
+ task_id="example_gcp_sql_task_" + connection_name,
+ sql=SQL,
+ sql_proxy_binary_path=sql_proxy_binary_path,
  )
  tasks.append(task)
  if prev_task:

diff --git a/airflow/providers/google/cloud/hooks/cloud_sql.py b/airflow/providers/google/cloud/hooks/cloud_sql.py
@@ -668,6 +668,8 @@ class CloudSQLDatabaseHook(BaseHook):
  * **public_ip** - IP to connect to for public connection (from host of the URI).
  * **public_port** - Port to connect to for public connection (from port of the URI).
  * **database** - Database to connect to (from schema of the URI).
+ * **sql_proxy_binary_path** - Optional path to Cloud SQL Proxy binary. If the binary
+ is not specified or the binary is not present, it is automatically downloaded.
 
  Remaining parameters are retrieved from the extras (URI query parameters):
 
@@ -682,8 +684,6 @@ class CloudSQLDatabaseHook(BaseHook):
  You cannot use proxy and SSL together.
  * **sql_proxy_use_tcp** - (default False) If set to true, TCP is used to connect via
  proxy, otherwise UNIX sockets are used.
- * **sql_proxy_binary_path** - Optional path to Cloud SQL Proxy binary. If the binary
- is not specified or the binary is not present, it is automatically downloaded.
  * **sql_proxy_version** - Specific version of the proxy to download (for example
  v1.13). If not specified, the latest version is downloaded.
  * **sslcert** - Path to client certificate to authenticate when SSL is used.
@@ -707,6 +707,7 @@ def __init__(
  gcp_cloudsql_conn_id: str = "google_cloud_sql_default",
  gcp_conn_id: str = "google_cloud_default",
  default_gcp_project_id: str | None = None,
+ sql_proxy_binary_path: str | None = None,
  ) -> None:
  super().__init__()
  self.gcp_conn_id = gcp_conn_id
@@ -722,7 +723,7 @@ def __init__(
  self.use_ssl = self._get_bool(self.extras.get("use_ssl", "False"))
  self.sql_proxy_use_tcp = self._get_bool(self.extras.get("sql_proxy_use_tcp", "False"))
  self.sql_proxy_version = self.extras.get("sql_proxy_version")
- self.sql_proxy_binary_path = self.extras.get("sql_proxy_binary_path")
+ self.sql_proxy_binary_path = sql_proxy_binary_path
  self.user = self.cloudsql_connection.login
  self.password = self.cloudsql_connection.password
  self.public_ip = self.cloudsql_connection.host

diff --git a/airflow/providers/google/cloud/operators/cloud_sql.py b/airflow/providers/google/cloud/operators/cloud_sql.py
@@ -1045,6 +1045,8 @@ class CloudSQLExecuteQueryOperator(BaseOperator):
  its schema should be gcpcloudsql://.
  See :class:`~airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLDatabaseHook` for
  details on how to define ``gcpcloudsql://`` connection.
+ :param sql_proxy_binary_path: (optional) Path to the cloud-sql-proxy binary.
+ is not specified or the binary is not present, it is automatically downloaded.
  """
 
  # [START gcp_sql_query_template_fields]
@@ -1062,6 +1064,7 @@ def __init__(
  parameters: Iterable | Mapping | None = None,
  gcp_conn_id: str = "google_cloud_default",
  gcp_cloudsql_conn_id: str = "google_cloud_sql_default",
+ sql_proxy_binary_path: str | None = None,
  **kwargs,
  ) -> None:
  super().__init__(**kwargs)
@@ -1071,6 +1074,7 @@ def __init__(
  self.autocommit = autocommit
  self.parameters = parameters
  self.gcp_connection: Connection | None = None
+ self.sql_proxy_binary_path = sql_proxy_binary_path
 
  def _execute_query(self, hook: CloudSQLDatabaseHook, database_hook: PostgresHook | MySqlHook) -> None:
  cloud_sql_proxy_runner = None
@@ -1094,6 +1098,7 @@ def execute(self, context: Context):
  gcp_cloudsql_conn_id=self.gcp_cloudsql_conn_id,
  gcp_conn_id=self.gcp_conn_id,
  default_gcp_project_id=get_field(self.gcp_connection.extra_dejson, "project"),
+ sql_proxy_binary_path=self.sql_proxy_binary_path,
  )
  hook.validate_ssl_certs()
  connection = hook.create_connection()