This is an automated email from the ASF dual-hosted git repository.
shahar1 pushed a commit to branch main
in repository https://gitbox.apache.org/repos/asf/airflow.git
The following commit(s) were added to refs/heads/main by this push:
new f396a3b2271 Add Cloud SQL Auth Proxy v2 support to CloudSqlProxyRunner
(#73646)
f396a3b2271 is described below
commit f396a3b227196d9d90a4c8ffee6eb574ed316ccf
Author: Arno Fukuda <[email protected]>
AuthorDate: Sat Nov 7 00:19:47 2026 +0900
Add Cloud SQL Auth Proxy v2 support to CloudSqlProxyRunner (#73646)
---
providers/google/docs/connections/gcp_sql.rst | 55 ++++-
.../providers/google/cloud/hooks/cloud_sql.py | 116 ++++++++--
.../unit/google/cloud/hooks/test_cloud_sql.py | 234 +++++++++++++++++++++
3 files changed, 385 insertions(+), 20 deletions(-)
diff --git a/providers/google/docs/connections/gcp_sql.rst
b/providers/google/docs/connections/gcp_sql.rst
index c7280b33ea3..743a3c36cb6 100644
--- a/providers/google/docs/connections/gcp_sql.rst
+++ b/providers/google/docs/connections/gcp_sql.rst
@@ -78,6 +78,50 @@ Extra (optional)
export
AIRFLOW_CONN_GOOGLE_CLOUD_SQL_DEFAULT='gcpcloudsql://user:[email protected]:3306/mydb?database_type=mysql&project_id=example-project&location=europe-west1&instance=testinstance&use_proxy=True&sql_proxy_use_tcp=False'
+Using Cloud SQL Auth Proxy v2
+-----------------------------
+
+By default Airflow runs Cloud SQL Auth Proxy v1. To run Cloud SQL Auth Proxy
v2 instead, set
+``"sql_proxy_major_version": 2`` in the ``extra`` field. ``sql_proxy_version``
must also be set to the v2
+release to download, as there is no "latest" download for v2, unless the
``sql_proxy_binary_path`` argument
+of the hook or operator points to an existing v2 binary. Both must be v2.2.0
or newer.
+
+Cloud SQL for MySQL 8.4 users authenticating with ``caching_sha2_password``
need v2.19.0 or newer, and
+must keep the default Unix socket mode (``sql_proxy_use_tcp`` disabled). v1
and older v2 releases cannot
+complete that plugin's full authentication, so connections fail with ``1045
Access denied`` after the
+server's authentication cache is emptied, for example by an instance restart.
Over TCP, full authentication
+needs the client to retrieve the server's public key, which the MySQL hook
cannot be configured to do.
+
+Cloud SQL Auth Proxy v2 connects only to the instance's public IP by default,
while v1 falls back to the
+private IP when the instance has no public IP. Airflow passes ``--auto-ip`` to
v2 to keep the v1 behavior,
+so instances with only a private IP work with both versions. ``--auto-ip`` was
added in v2.2.0, which is
+why older v2 releases exit at startup.
+
+The two versions also differ in authentication:
+
+* When the Google Cloud connection has neither ``key_path`` nor
``keyfile_dict``, v1 tries the
+ ``GOOGLE_APPLICATION_CREDENTIALS`` file, then the active ``gcloud`` account,
then Application Default
+ Credentials. v2 uses Application Default Credentials only. A worker
authenticated only with
+ ``gcloud auth login`` or ``gcloud auth activate-service-account`` can
therefore work with v1 but fail,
+ or connect as a different identity, with v2.
+* v2 attributes its Cloud SQL Admin API calls to the quota project of its
credentials. Credentials whose
+ quota project is unrelated to the instance, which is common for local
Application Default Credentials,
+ make v2 fail with ``403 SERVICE_DISABLED`` naming that other project.
+
+Example "extras" field:
+
+.. code-block:: json
+
+ {
+ "database_type": "mysql",
+ "project_id": "example-project",
+ "location": "europe-west1",
+ "instance": "testinstance",
+ "use_proxy": true,
+ "sql_proxy_major_version": 2,
+ "sql_proxy_version": "v2.27.0"
+ }
+
Configuring and using IAM authentication
----------------------------------------
@@ -124,9 +168,10 @@ Configure ``gcpcloudsql`` connection with Cloud SQL Auth
Proxy IAM authenticatio
"""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""""
For using Cloud SQL Auth Proxy IAM authentication, enable ``"use_proxy":
"True"`` and
-``"sql_proxy_enable_iam_login": "True"`` in the ``extra`` field. With the
current Cloud SQL Auth Proxy
-v1 integration this option is supported for both Postgres and MySQL. Airflow
passes
-``-enable_iam_login`` to the proxy, so the ``password`` field can be empty.
+``"sql_proxy_enable_iam_login": "True"`` in the ``extra`` field. Airflow
passes ``-enable_iam_login`` to
+Cloud SQL Auth Proxy v1 or ``--auto-iam-authn`` to v2, so the ``password``
field can be empty.
+Cloud SQL Auth Proxy v1 supports IAM database authentication for Postgres
only; for MySQL, use
+Cloud SQL Auth Proxy v2 (see `Using Cloud SQL Auth Proxy v2`_).
Example "extras" field for Postgres:
@@ -153,7 +198,9 @@ Example "extras" field for MySQL:
"instance": "testinstance",
"use_proxy": true,
"sql_proxy_use_tcp": true,
- "sql_proxy_enable_iam_login": true
+ "sql_proxy_enable_iam_login": true,
+ "sql_proxy_major_version": 2,
+ "sql_proxy_version": "v2.27.0"
}
.. note::
diff --git
a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
index 80ef340027a..9fbb6f30245 100644
--- a/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
+++ b/providers/google/src/airflow/providers/google/cloud/hooks/cloud_sql.py
@@ -35,7 +35,7 @@ import time
import uuid
from collections.abc import Sequence
from pathlib import Path
-from subprocess import PIPE, Popen
+from subprocess import PIPE, STDOUT, Popen
from tempfile import NamedTemporaryFile, _TemporaryFileWrapper, gettempdir
from typing import TYPE_CHECKING, Any, cast
from urllib.parse import quote_plus
@@ -81,6 +81,7 @@ UNIX_PATH_MAX = 108
TIME_TO_SLEEP_IN_SECONDS = 20
CLOUD_SQL_PROXY_VERSION_REGEX = re.compile(r"^v?(\d+\.\d+\.\d+)(-\w*.?\d?)?$")
+CLOUD_SQL_PROXY_V2_VERSION_REGEX = re.compile(r"v2\.\d+\.\d+")
class CloudSQLImportError(AirflowException):
@@ -528,6 +529,9 @@ CLOUD_SQL_PROXY_DOWNLOAD_URL =
"https://dl.google.com/cloudsql/cloud_sql_proxy.{
CLOUD_SQL_PROXY_VERSION_DOWNLOAD_URL = (
"https://storage.googleapis.com/cloudsql-proxy/{}/cloud_sql_proxy.{}.{}";
)
+CLOUD_SQL_PROXY_V2_DOWNLOAD_URL = (
+
"https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/{}/cloud-sql-proxy.{}.{}";
+)
class CloudSQLAsyncHook(GoogleBaseAsyncHook):
@@ -589,15 +593,22 @@ class CloudSqlProxyRunner(LoggingMixin):
-instances parameter (typically in the form of
``<project>:<region>:<instance>``
for UNIX socket connections and in the form of
``<project>:<region>:<instance>=tcp:<port>`` for TCP connections.
+ With ``sql_proxy_major_version=2`` use
``<project>:<region>:<instance>?port=<port>``
+ for TCP connections instead.
:param gcp_conn_id: Id of Google Cloud connection to use for
authentication
:param project_id: Optional id of the Google Cloud project to connect to -
it overwrites
default project id taken from the Google Cloud connection.
:param sql_proxy_version: Specific version of SQL proxy to download
(for example 'v1.13'). By default, latest version is downloaded.
+ With ``sql_proxy_major_version=2`` the version (for example 'v2.27.0')
is required
+ to download the proxy, as there is no "latest" download for v2.
:param sql_proxy_binary_path: If specified, then proxy will be
used from the path specified rather than dynamically generated. This
means
that if the binary is not present in that path it will also be
downloaded.
+ :param sql_proxy_enable_iam_login: Whether the proxy should use IAM
database authentication.
+ :param sql_proxy_major_version: Major version of Cloud SQL Auth Proxy to
run: ``1`` (default)
+ or ``2``. The two versions have different download locations and
command line flags.
"""
def __init__(
@@ -610,11 +621,15 @@ class CloudSqlProxyRunner(LoggingMixin):
sql_proxy_binary_path: str | None = None,
*,
sql_proxy_enable_iam_login: bool = False,
+ sql_proxy_major_version: int = 1,
) -> None:
super().__init__()
self.path_prefix = path_prefix
if not self.path_prefix:
raise AirflowException("The path_prefix must not be empty!")
+ if sql_proxy_major_version not in (1, 2):
+ raise ValueError(f"The sql_proxy_major_version must be 1 or 2, got
{sql_proxy_major_version!r}!")
+ self.sql_proxy_major_version = sql_proxy_major_version
self.sql_proxy_was_downloaded = False
self.sql_proxy_version = sql_proxy_version
self.download_sql_proxy_dir = None
@@ -627,14 +642,44 @@ class CloudSqlProxyRunner(LoggingMixin):
self.cloud_sql_proxy_socket_directory = self.path_prefix
self.sql_proxy_path = sql_proxy_binary_path or
f"{self.path_prefix}_cloud_sql_proxy"
self.credentials_path = self.path_prefix + "_credentials.json"
+ self._validate_sql_proxy_configuration()
self._build_command_line_parameters()
+ def _validate_sql_proxy_configuration(self) -> None:
+ # Validated here rather than in start_proxy(): callers stop the proxy
on failure, and
+ # stop_proxy() on a proxy that never started raises and hides the
original error.
+ if self.sql_proxy_major_version == 1:
+ if self.sql_proxy_version and
self.sql_proxy_version.lstrip("v").startswith("2."):
+ raise ValueError(
+ f"The sql_proxy_version {self.sql_proxy_version!r} is a
Cloud SQL Auth Proxy v2 "
+ "release. Set sql_proxy_major_version to 2 to use it!"
+ )
+ return
+ if not self.instance_specification:
+ raise ValueError(
+ "Cloud SQL Auth Proxy v2 does not support forwarding all
instances of a project. "
+ "The instance_specification must be provided!"
+ )
+ if not os.path.isfile(self.sql_proxy_path):
+ self._get_sql_proxy_download_url()
+
def _build_command_line_parameters(self) -> None:
+ if self.sql_proxy_major_version == 2:
+ self._build_v2_command_line_parameters()
+ return
self.command_line_parameters.extend(["-dir",
self.cloud_sql_proxy_socket_directory])
self.command_line_parameters.extend(["-instances",
self.instance_specification])
if self.sql_proxy_enable_iam_login:
self.command_line_parameters.append("-enable_iam_login")
+ def _build_v2_command_line_parameters(self) -> None:
+ self.command_line_parameters.extend(["--unix-socket",
self.cloud_sql_proxy_socket_directory])
+ # v2 connects to the public IP only by default; --auto-ip keeps v1's
fallback to the private IP.
+ self.command_line_parameters.append("--auto-ip")
+ if self.sql_proxy_enable_iam_login:
+ self.command_line_parameters.append("--auto-iam-authn")
+ self.command_line_parameters.append(self.instance_specification)
+
def _download_sql_proxy_if_needed(self) -> None:
if os.path.isfile(self.sql_proxy_path):
self.log.info("cloud-sql-proxy is already present")
@@ -665,6 +710,18 @@ class CloudSqlProxyRunner(LoggingMixin):
processor = "amd64"
elif processor == "aarch64":
processor = "arm64"
+ if self.sql_proxy_major_version == 2:
+ if not self.sql_proxy_version:
+ raise ValueError(
+ "The sql_proxy_version must be specified to download Cloud
SQL Auth Proxy v2 "
+ "(for example 'v2.27.0')!"
+ )
+ if not
CLOUD_SQL_PROXY_V2_VERSION_REGEX.fullmatch(self.sql_proxy_version):
+ raise ValueError(
+ "The sql_proxy_version should match the regular expression
"
+ f"{CLOUD_SQL_PROXY_V2_VERSION_REGEX.pattern}"
+ )
+ return
CLOUD_SQL_PROXY_V2_DOWNLOAD_URL.format(self.sql_proxy_version, system,
processor)
if not self.sql_proxy_version:
download_url = CLOUD_SQL_PROXY_DOWNLOAD_URL.format(system,
processor)
else:
@@ -682,8 +739,12 @@ class CloudSqlProxyRunner(LoggingMixin):
extras =
GoogleBaseHook.get_connection(conn_id=self.gcp_conn_id).extra_dejson
key_path = get_field(extras, "key_path")
keyfile_dict = get_field(extras, "keyfile_dict")
+ if self.sql_proxy_major_version == 2:
+ credential_file_flag = "--credentials-file"
+ else:
+ credential_file_flag = "-credential_file"
if key_path:
- credential_params = ["-credential_file", key_path]
+ credential_params = [credential_file_flag, key_path]
elif keyfile_dict:
keyfile_content = keyfile_dict if isinstance(keyfile_dict, dict)
else json.loads(keyfile_dict)
self.log.info("Saving credentials to %s", self.credentials_path)
@@ -693,12 +754,12 @@ class CloudSqlProxyRunner(LoggingMixin):
fd = os.open(self.credentials_path, os.O_WRONLY | os.O_CREAT |
os.O_TRUNC, 0o600)
with os.fdopen(fd, "w") as file:
json.dump(keyfile_content, file)
- credential_params = ["-credential_file", self.credentials_path]
+ credential_params = [credential_file_flag, self.credentials_path]
else:
self.log.info(
"The credentials are not supplied by neither key_path nor "
"keyfile_dict of the gcp connection %s. Falling back to "
- "default activated account",
+ "the default credentials of Cloud SQL Auth Proxy",
self.gcp_conn_id,
)
credential_params = []
@@ -733,14 +794,19 @@ class CloudSqlProxyRunner(LoggingMixin):
command_to_run.extend(self._get_credential_parameters())
self.log.info("Running the command: `%s`", " ".join(command_to_run))
- self.sql_proxy_process = Popen(command_to_run, stdin=PIPE,
stdout=PIPE, stderr=PIPE)
+ if self.sql_proxy_major_version == 2:
+ # v2 logs informational messages, including the ready line, to
stdout
+ self.sql_proxy_process = Popen(command_to_run, stdin=PIPE,
stdout=PIPE, stderr=STDOUT)
+ output = self.sql_proxy_process.stdout
+ else:
+ self.sql_proxy_process = Popen(command_to_run, stdin=PIPE,
stdout=PIPE, stderr=PIPE)
+ output = self.sql_proxy_process.stderr
self.log.info("The pid of cloud_sql_proxy: %s",
self.sql_proxy_process.pid)
while True:
- line = (
- self.sql_proxy_process.stderr.readline().decode("utf-8")
- if self.sql_proxy_process.stderr
- else ""
- )
+ if output:
+ line = output.readline().decode("utf-8")
+ else:
+ line = ""
return_code = self.sql_proxy_process.poll()
if line == "" and return_code is not None:
self.sql_proxy_process = None
@@ -750,7 +816,8 @@ class CloudSqlProxyRunner(LoggingMixin):
if "googleapi: Error" in line or "invalid instance name:" in line:
self.stop_proxy()
raise AirflowException(f"Error when starting the
cloud_sql_proxy {line}!")
- if "Ready for new connections" in line:
+ # v1: "Ready for new connections", v2: "... is ready for new
connections!"
+ if "ready for new connections" in line.lower():
return
def stop_proxy(self) -> None:
@@ -789,7 +856,8 @@ class CloudSqlProxyRunner(LoggingMixin):
command_to_run.extend(["--version"])
command_to_run.extend(self._get_credential_parameters())
result = subprocess.check_output(command_to_run).decode("utf-8")
- matched = re.search("[Vv]ersion (.*?);", result)
+ # v1: "... version 1.13; ...", v2: "cloud-sql-proxy version
2.14.0+linux.amd64"
+ matched = re.search(r"[Vv]ersion ([^;\s]+)", result)
if matched:
return matched.group(1)
return None
@@ -869,12 +937,15 @@ class CloudSQLDatabaseHook(BaseHook):
* **use_iam** - (default False) Whether IAM should be used to connect to
Cloud SQL DB.
With using IAM password field should be empty string.
* **sql_proxy_enable_iam_login** - (default False) Whether Cloud SQL Auth
Proxy should use
- IAM database authentication. This requires ``use_proxy`` and is
supported with the current
- Cloud SQL Auth Proxy v1 integration for both Postgres and MySQL.
+ IAM database authentication. This requires ``use_proxy``. With Cloud SQL
Auth Proxy v1 it is
+ supported for Postgres only, with v2 for both Postgres and MySQL.
* **sql_proxy_use_tcp** - (default False) If set to true, TCP is used to
connect via
proxy, otherwise UNIX sockets are used.
* **sql_proxy_version** - Specific version of the proxy to download (for
example
- v1.13). If not specified, the latest version is downloaded.
+ v1.13). If not specified, the latest version is downloaded. Required when
+ ``sql_proxy_major_version`` is 2 and the binary has to be downloaded
(for example v2.27.0).
+ * **sql_proxy_major_version** - (default 1) Major version of Cloud SQL
Auth Proxy to use,
+ either 1 or 2.
* **sslcert** - Path to client certificate to authenticate when SSL is
used.
* **sslkey** - Path to client private key to authenticate when SSL is used.
* **sslrootcert** - Path to server's certificate to authenticate when SSL
is used.
@@ -940,6 +1011,9 @@ class CloudSQLDatabaseHook(BaseHook):
)
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_major_version = self._get_sql_proxy_major_version(
+ self.extras.get("sql_proxy_major_version", 1)
+ )
self.sql_proxy_binary_path = sql_proxy_binary_path
self.public_ip = self.cloudsql_connection.host
self.public_port = self.cloudsql_connection.port
@@ -1055,6 +1129,12 @@ class CloudSQLDatabaseHook(BaseHook):
return False
return True
+ @staticmethod
+ def _get_sql_proxy_major_version(val: Any) -> int:
+ if str(val) not in ("1", "2"):
+ raise ValueError(f"The extra 'sql_proxy_major_version' must be 1
or 2, got {val!r}")
+ return int(val)
+
@staticmethod
def _check_ssl_file(file_to_check, name) -> None:
if not file_to_check:
@@ -1222,7 +1302,10 @@ class CloudSQLDatabaseHook(BaseHook):
def _get_sqlproxy_instance_specification(self) -> str:
instance_specification = self._get_instance_socket_name()
if self.sql_proxy_use_tcp:
- instance_specification += f"=tcp:{self.sql_proxy_tcp_port}"
+ if self.sql_proxy_major_version == 2:
+ instance_specification += f"?port={self.sql_proxy_tcp_port}"
+ else:
+ instance_specification += f"=tcp:{self.sql_proxy_tcp_port}"
return instance_specification
def _generate_connection_parameters(self) -> dict:
@@ -1325,6 +1408,7 @@ class CloudSQLDatabaseHook(BaseHook):
sql_proxy_binary_path=self.sql_proxy_binary_path,
gcp_conn_id=self.gcp_conn_id,
sql_proxy_enable_iam_login=self.sql_proxy_enable_iam_login,
+ sql_proxy_major_version=self.sql_proxy_major_version,
)
def get_database_hook(self, connection: Connection) -> DbApiHook:
diff --git a/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
b/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
index 07e2f9be6e8..6bf93f9f2fa 100644
--- a/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
+++ b/providers/google/tests/unit/google/cloud/hooks/test_cloud_sql.py
@@ -24,6 +24,7 @@ import platform
import stat
import tempfile
from pathlib import Path
+from subprocess import PIPE, STDOUT
from unittest import mock
from unittest.mock import PropertyMock, call, mock_open
from urllib.parse import parse_qsl, unquote, urlsplit
@@ -1216,6 +1217,75 @@ class TestCloudSqlDatabaseHook:
proxy_runner = hook.get_sqlproxy_runner()
assert proxy_runner is not None
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.BaseHook.get_connection",
autospec=True)
+ def test_cloudsql_database_hook_get_sqlproxy_runner_v2_tcp(self,
get_connection):
+ if AIRFLOW_V_3_1_PLUS:
+ connection = Connection(
+ conn_id="test_conn_id",
**_parse_from_uri("http://user:password@host:80/database";)
+ )
+ else:
+ connection =
Connection(uri="http://user:password@host:80/database";)
+ extras = json.dumps(
+ {
+ "location": "test",
+ "instance": "instance",
+ "database_type": "postgres",
+ "use_proxy": "True",
+ "sql_proxy_use_tcp": "True",
+ "sql_proxy_major_version": "2",
+ "sql_proxy_version": "v2.14.0",
+ }
+ )
+ if AIRFLOW_V_3_1_PLUS:
+ connection.extra = extras
+ else:
+ connection.set_extra(extras)
+ get_connection.return_value = connection
+ hook = CloudSQLDatabaseHook(
+ gcp_cloudsql_conn_id="cloudsql_connection",
default_gcp_project_id="google_connection"
+ )
+ hook.create_connection()
+ proxy_runner = hook.get_sqlproxy_runner()
+
+ assert proxy_runner.sql_proxy_major_version == 2
+ assert proxy_runner.command_line_parameters == [
+ "--unix-socket",
+ hook.sql_proxy_unique_path,
+ "--auto-ip",
+ f"google_connection:test:instance?port={hook.sql_proxy_tcp_port}",
+ ]
+
+ @pytest.mark.parametrize("major_version", ["", "3", "v2", None])
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.BaseHook.get_connection",
autospec=True)
+ def test_cloudsql_database_hook_rejects_invalid_sql_proxy_major_version(
+ self, get_connection, major_version
+ ):
+ if AIRFLOW_V_3_1_PLUS:
+ connection = Connection(
+ conn_id="test_conn_id",
**_parse_from_uri("http://user:password@host:80/database";)
+ )
+ else:
+ connection =
Connection(uri="http://user:password@host:80/database";)
+ extras = json.dumps(
+ {
+ "location": "test",
+ "instance": "instance",
+ "database_type": "postgres",
+ "use_proxy": "True",
+ "sql_proxy_major_version": major_version,
+ }
+ )
+ if AIRFLOW_V_3_1_PLUS:
+ connection.extra = extras
+ else:
+ connection.set_extra(extras)
+ get_connection.return_value = connection
+
+ with pytest.raises(ValueError, match="The extra
'sql_proxy_major_version' must be 1 or 2"):
+ CloudSQLDatabaseHook(
+ gcp_cloudsql_conn_id="cloudsql_connection",
default_gcp_project_id="google_connection"
+ )
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.CloudSQLDatabaseHook.get_connection")
def test_cloudsql_database_hook_get_database_hook(self, get_connection):
if AIRFLOW_V_3_1_PLUS:
@@ -2051,6 +2121,170 @@ class TestCloudSqlProxyRunner:
assert runner._get_credential_parameters() == ["-credential_file",
"/tmp/key.json"]
assert "-enable_iam_login" in runner.command_line_parameters
+ def test_cloud_sql_proxy_runner_rejects_unknown_major_version(self):
+ with pytest.raises(ValueError, match="The sql_proxy_major_version must
be 1 or 2"):
+ CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="project:us-east-1:instance",
+ sql_proxy_major_version=3,
+ )
+
+ def test_cloud_sql_proxy_runner_v1_rejects_v2_version(self):
+ with pytest.raises(ValueError, match="Set sql_proxy_major_version to
2"):
+ CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="project:us-east-1:instance",
+ sql_proxy_version="v2.14.0",
+ )
+
+ @pytest.mark.parametrize(
+ "instance_specification",
+ ["project:us-east-1:instance", "project:us-east-1:instance?port=5432"],
+ ids=["unix_socket", "tcp"],
+ )
+ def test_cloud_sql_proxy_runner_v2_command_line(self,
instance_specification):
+ runner = CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification=instance_specification,
+ sql_proxy_version="v2.14.0",
+ sql_proxy_major_version=2,
+ )
+
+ assert runner.command_line_parameters == [
+ "--unix-socket",
+ "12345678",
+ "--auto-ip",
+ instance_specification,
+ ]
+
+ def test_cloud_sql_proxy_runner_v2_adds_auto_iam_authn_flag(self):
+ runner = CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="project:us-east-1:instance",
+ sql_proxy_version="v2.14.0",
+ sql_proxy_enable_iam_login=True,
+ sql_proxy_major_version=2,
+ )
+
+ assert "--auto-iam-authn" in runner.command_line_parameters
+ assert "-enable_iam_login" not in runner.command_line_parameters
+
+ def test_cloud_sql_proxy_runner_v2_version_ok(self):
+ runner = CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="project:us-east-1:instance",
+ sql_proxy_version="v2.14.0",
+ sql_proxy_major_version=2,
+ )
+
+ assert runner._get_sql_proxy_download_url() == (
+
"https://storage.googleapis.com/cloud-sql-connectors/cloud-sql-proxy/v2.14.0/cloud-sql-proxy.";
+ f"{platform.system().lower()}.{get_processor()}"
+ )
+
+ @pytest.mark.parametrize("version", ["v1.23.0", "2.14.0", "v2.14",
"v2.14.0\n"])
+ def test_cloud_sql_proxy_runner_v2_version_nok(self, version):
+ with pytest.raises(ValueError, match="The sql_proxy_version should
match the regular expression"):
+ CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="project:us-east-1:instance",
+ sql_proxy_version=version,
+ sql_proxy_major_version=2,
+ )
+
+ def test_cloud_sql_proxy_runner_v2_requires_version_to_download(self):
+ with pytest.raises(ValueError, match="The sql_proxy_version must be
specified"):
+ CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="project:us-east-1:instance",
+ sql_proxy_major_version=2,
+ )
+
+ def test_cloud_sql_proxy_runner_v2_requires_instance_specification(self):
+ with pytest.raises(ValueError, match="does not support forwarding all
instances"):
+ CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="",
+ sql_proxy_version="v2.14.0",
+ sql_proxy_major_version=2,
+ )
+
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.GoogleBaseHook.get_connection",
autospec=True)
+ def test_cloud_sql_proxy_runner_v2_credentials_file_flag(self,
get_connection):
+ connection = Connection(conn_id="google_conn",
conn_type="google_cloud_platform")
+ if AIRFLOW_V_3_1_PLUS:
+ connection.extra = json.dumps({"key_path": "/tmp/key.json"})
+ else:
+ connection.set_extra(json.dumps({"key_path": "/tmp/key.json"}))
+ get_connection.return_value = connection
+ runner = CloudSqlProxyRunner(
+ path_prefix="12345678",
+ instance_specification="project:us-east-1:instance",
+ gcp_conn_id="google_conn",
+ sql_proxy_version="v2.14.0",
+ sql_proxy_major_version=2,
+ )
+
+ assert runner._get_credential_parameters() == ["--credentials-file",
"/tmp/key.json"]
+
+ @mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.Popen",
autospec=True)
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.GoogleBaseHook.get_connection",
autospec=True)
+ def
test_cloud_sql_proxy_runner_v2_start_proxy_reads_ready_line_from_stdout(
+ self, get_connection, mock_popen, tmp_path
+ ):
+ get_connection.return_value = Connection(conn_id="google_conn",
conn_type="google_cloud_platform")
+ # A pre-existing binary means no download, so no sql_proxy_version is
needed.
+ proxy_binary = tmp_path / "cloud-sql-proxy"
+ proxy_binary.touch()
+ process = mock_popen.return_value
+ process.pid = 1234
+ process.poll.return_value = None
+ process.stdout = mock.Mock(spec=["readline"])
+ process.stdout.readline.side_effect = [
+ b"Authorizing with Application Default Credentials\n",
+ b"The proxy has started successfully and is ready for new
connections!\n",
+ ]
+ runner = CloudSqlProxyRunner(
+ path_prefix=str(tmp_path / "12345678"),
+ instance_specification="project:us-east-1:instance",
+ gcp_conn_id="google_conn",
+ sql_proxy_binary_path=str(proxy_binary),
+ sql_proxy_major_version=2,
+ )
+
+ runner.start_proxy()
+
+ mock_popen.assert_called_once_with(
+ [
+ str(proxy_binary),
+ "--unix-socket",
+ str(tmp_path / "12345678"),
+ "--auto-ip",
+ "project:us-east-1:instance",
+ ],
+ stdin=PIPE,
+ stdout=PIPE,
+ stderr=STDOUT,
+ )
+ assert process.stdout.readline.call_count == 2
+
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.subprocess.check_output",
autospec=True)
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.GoogleBaseHook.get_connection",
autospec=True)
+ def test_cloud_sql_proxy_runner_v2_get_proxy_version(self, get_connection,
check_output, tmp_path):
+ get_connection.return_value = Connection(conn_id="google_conn",
conn_type="google_cloud_platform")
+ check_output.return_value = b"cloud-sql-proxy version
2.14.0+linux.amd64\n"
+ proxy_binary = tmp_path / "cloud-sql-proxy"
+ proxy_binary.touch()
+ runner = CloudSqlProxyRunner(
+ path_prefix=str(tmp_path / "12345678"),
+ instance_specification="project:us-east-1:instance",
+ gcp_conn_id="google_conn",
+ sql_proxy_binary_path=str(proxy_binary),
+ sql_proxy_major_version=2,
+ )
+
+ assert runner.get_proxy_version() == "2.14.0+linux.amd64"
+
@mock.patch("airflow.providers.google.cloud.hooks.cloud_sql.GoogleBaseHook.get_connection")
def test_credentials_file_from_keyfile_dict_is_chmod_0600(self,
get_connection, tmp_path):
"""The keyfile_dict credentials file must be written with explicit
0600 permissions.
