download server CA cert for tests (GoogleCloudPlatform#9086)

shubha-rajan · jackwotherspoon · web-flow · commit 4a288ce509a5 · 2023-02-08T08:54:28.000-06:00
* download server CA cert for tests

* add note to connect_tcp

* remove raise_for_status

* debug print response

* refresh invalid credential

* linting

* Update connection_test.py

---------

Co-authored-by: Jack Wotherspoon &lt;jackwoth@google.com&gt;
diff --git a/cloud-sql/sql-server/sqlalchemy/connect_connector.py b/cloud-sql/sql-server/sqlalchemy/connect_connector.py
@@ -41,13 +41,24 @@ def connect_with_connector() -> sqlalchemy.engine.base.Engine:
 
     connector = Connector(ip_type)
 
+    connect_args = {}
+    # If your SQL Server instance requires SSL, you need to download the CA
+    # certificate for your instance and include cafile={path to downloaded
+    # certificate} and validate_host=False. This is a workaround for a known issue.
+    if os.environ.get("DB_ROOT_CERT"):  # e.g. '/path/to/my/server-ca.pem'
+        connect_args = {
+            "cafile" : os.environ["DB_ROOT_CERT"],
+            "validate_host": False,
+        }
+
     def getconn() -> pytds.Connection:
         conn = connector.connect(
             instance_connection_name,
             "pytds",
             user=db_user,
             password=db_pass,
             db=db_name,
+            **connect_args
         )
         return conn
 
diff --git a/cloud-sql/sql-server/sqlalchemy/connect_tcp.py b/cloud-sql/sql-server/sqlalchemy/connect_tcp.py
@@ -40,6 +40,11 @@ def connect_tcp_socket() -> sqlalchemy.engine.base.Engine:
     # For deployments that connect directly to a Cloud SQL instance without
     # using the Cloud SQL Proxy, configuring SSL certificates will ensure the
     # connection is encrypted.
+
+    # If your SQL Server instance requires SSL, you need to download the CA
+    # certificate for your instance and include cafile={path to downloaded
+    # certificate} and validate_host=False, even when using the proxy.
+    # This is a workaround for a known issue.
     if os.environ.get("DB_ROOT_CERT"):  # e.g. '/path/to/my/server-ca.pem'
         connect_args = {
             "cafile" : os.environ["DB_ROOT_CERT"],
diff --git a/cloud-sql/sql-server/sqlalchemy/connection_test.py b/cloud-sql/sql-server/sqlalchemy/connection_test.py
@@ -16,14 +16,21 @@
 import os
 
 from flask.testing import FlaskClient
-
+from google.auth import default
+import google.auth.transport.requests
 import pytest
+import requests
 
 import app
 
 logger = logging.getLogger()
 
 
+CA_FILENAME = "certs/ca.pem"
+SQLADMIN_API_ENDPOINT = "https://sqladmin.googleapis.com"
+SQLADMIN_API_VERSION = "v1beta4"
+
+
 # load proper environment variables
 def setup_test_env():
     os.environ["DB_USER"] = os.environ["SQLSERVER_USER"]
@@ -33,6 +40,32 @@ def setup_test_env():
     os.environ["INSTANCE_HOST"] = os.environ["SQLSERVER_INSTANCE_HOST"]
     os.environ["INSTANCE_CONNECTION_NAME"] = os.environ["SQLSERVER_INSTANCE"]
 
+    project, _, instance = os.environ["INSTANCE_CONNECTION_NAME"].split(":")
+    download_ca_cert(project, instance)
+    os.environ["DB_ROOT_CERT"] = CA_FILENAME
+
+
+def download_ca_cert(project, instance):
+    """ Download server CA cert"""
+    scopes = ["https://www.googleapis.com/auth/sqlservice.admin"]
+    credentials, _ = default(scopes=scopes)
+
+    if not credentials.valid:
+        request = google.auth.transport.requests.Request()
+        credentials.refresh(request)
+
+    headers = {
+        "Authorization": f"Bearer {credentials.token}",
+    }
+    url = (f"{SQLADMIN_API_ENDPOINT}/sql/{SQLADMIN_API_VERSION}"
+           f"/projects/{project}/instances/{instance}/connectSettings")
+
+    resp = requests.get(url, headers=headers)
+    server_ca_cert = resp.json()["serverCaCert"]["cert"]
+
+    with open(CA_FILENAME, "w+") as ca_out:
+        ca_out.write(server_ca_cert)
+
 
 @pytest.fixture(scope="module")
 def client() -> FlaskClient:
diff --git a/cloud-sql/sql-server/sqlalchemy/requirements-test.txt b/cloud-sql/sql-server/sqlalchemy/requirements-test.txt
@@ -1 +1,3 @@
-pytest==7.0.1
+pytest==7.0.1
+google-auth==2.16.0
+Requests==2.28.2
