feat: add pg8000 support (#40)

shubha-rajan · web-flow · commit d810d7da9a50 · 2021-02-26T11:31:21.000-08:00
* feat: add support for connecting with pg8000

* add docstring

* linting

* update _connect_with_pg8000 to work with updated version of pg8000

* update env var names

* subclass ssl.SSLContext to get python3.6 tests passing

* blacken/lint

* fix typo in docstring
diff --git a/.gitignore b/.gitignore
@@ -2,3 +2,4 @@
 env
 venv
 *.pyc
+.python-version
diff --git a/google/cloud/sql/connector/InstanceConnectionManager.py b/google/cloud/sql/connector/InstanceConnectionManager.py
@@ -42,6 +42,16 @@
 _sql_api_version: str = "v1beta4"
 
 
+class ConnectionSSLContext(ssl.SSLContext):
+    """Subclass of ssl.SSLContext with added request_ssl attribute. This is
+    required for compatibility with pg8000 driver.
+    """
+
+    def __init__(self, *args, **kwargs):
+        self.request_ssl = False
+        super(ConnectionSSLContext, self).__init__(*args, **kwargs)
+
+
 class InstanceMetadata:
     ip_address: str
     _ca_fileobject: NamedTemporaryFile
@@ -73,7 +83,7 @@ def __init__(
         self._cert_fileobject.seek(0)
         self._key_fileobject.seek(0)
 
-        self.context = ssl.SSLContext()
+        self.context = ConnectionSSLContext()
         self.context.load_cert_chain(
             self._cert_fileobject.name, keyfile=self._key_fileobject.name
         )
@@ -455,8 +465,13 @@ def connect(self, driver: str, **kwargs) -> Any:
         with self._lock:
             instance_data: InstanceMetadata = self._current.result()
 
+        connect_func = {
+            "pymysql": self._connect_with_pymysql,
+            "pg8000": self._connect_with_pg8000,
+        }
+
         try:
-            connector = {"pymysql": self._connect_with_pymysql}[driver]
+            connector = connect_func[driver]
         except KeyError:
             raise KeyError("Driver {} is not supported.".format(driver))
 
@@ -499,3 +514,38 @@ def _connect_with_pymysql(self, ip_address: str, ctx: ssl.SSLContext, **kwargs):
         conn = pymysql.Connection(host=ip_address, defer_connect=True, **kwargs)
         conn.connect(sock)
         return conn
+
+    def _connect_with_pg8000(self, ip_address: str, ctx: ssl.SSLContext, **kwargs):
+        """Helper function to create a pg8000 DB-API connection object.
+
+        :type ip_address: str
+        :param ip_address: A string containing an IP address for the Cloud SQL
+            instance.
+
+        :type ctx: ssl.SSLContext
+        :param ctx: An SSLContext object created from the Cloud SQL server CA
+            cert and ephemeral cert.
+
+
+        :rtype: pg8000.dbapi.Connection
+        :returns: A pg8000 Connection object for the Cloud SQL instance.
+        """
+        try:
+            import pg8000
+        except ImportError:
+            raise ImportError(
+                'Unable to import module "pg8000." Please install and try again.'
+            )
+        user = kwargs.pop("user")
+        db = kwargs.pop("db")
+        passwd = kwargs.pop("password")
+        ctx.request_ssl = False
+        return pg8000.dbapi.connect(
+            user,
+            database=db,
+            password=passwd,
+            host=ip_address,
+            port=3307,
+            ssl_context=ctx,
+            **kwargs
+        )
diff --git a/requirements.txt b/requirements.txt
@@ -1,6 +1,8 @@
 aiohttp==3.7.4
 cryptography==3.4.6
 PyMySQL==1.0.2
+#pg8000==1.17.0
+git+https://github.com/tlocke/pg8000.git@37dcbe3#egg=pg8000
 pyopenssl==20.0.1
 pytest==6.2.2
 Requests==2.25.1
diff --git a/tests/system/test_pg8000_connection.py b/tests/system/test_pg8000_connection.py
@@ -0,0 +1,74 @@
+""""
+Copyright 2021 Google LLC
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+  https://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+"""
+import os
+import uuid
+
+import pytest
+import sqlalchemy
+from google.cloud.sql.connector import connector
+
+table_name = f"books_{uuid.uuid4().hex}"
+
+
+def init_connection_engine():
+    def getconn():
+        conn = connector.connect(
+            os.environ["POSTGRES_CONNECTION_NAME"],
+            "pg8000",
+            user=os.environ["POSTGRES_USER"],
+            password=os.environ["POSTGRES_PASS"],
+            db=os.environ["POSTGRES_DB"],
+        )
+        return conn
+
+    engine = sqlalchemy.create_engine(
+        "postgresql+pg8000://",
+        creator=getconn,
+    )
+    engine.dialect.description_encoding = None
+    return engine
+
+
+@pytest.fixture(name="pool")
+def setup():
+    pool = init_connection_engine()
+
+    with pool.connect() as conn:
+        conn.execute(
+            f"CREATE TABLE IF NOT EXISTS {table_name}"
+            " ( id CHAR(20) NOT NULL, title TEXT NOT NULL );"
+        )
+
+    yield pool
+
+    with pool.connect() as conn:
+        conn.execute(f"DROP TABLE IF EXISTS {table_name}")
+
+
+def test_pooled_connection_with_pymysql(pool):
+    insert_stmt = sqlalchemy.text(
+        f"INSERT INTO {table_name} (id, title) VALUES (:id, :title)",
+    )
+    with pool.connect() as conn:
+        conn.execute(insert_stmt, id="book1", title="Book One")
+        conn.execute(insert_stmt, id="book2", title="Book Two")
+
+    select_stmt = sqlalchemy.text(f"SELECT title FROM {table_name} ORDER BY ID;")
+    with pool.connect() as conn:
+        rows = conn.execute(select_stmt).fetchall()
+        titles = [row[0] for row in rows]
+
+    assert titles == ["Book One", "Book Two"]

-Original file line number
+Diff line change
 env
 venv
 *.pyc
 +.python-version