feat: add pytds support (#57)

* feat: add pytds support

* move proxy port number to constant

* scope dependencies in setup.py

* Update requirements.txt

* add type hints

* update setup.py

Author: shubha-rajan

.mypy.ini

@@ -7,4 +7,10 @@ plugins = sqlmypy
 ignore_missing_imports = True
 
 [mypy-pg8000]
+ignore_missing_imports = True
+
+[mypy-pytds]
+ignore_missing_imports = True
+
+[mypy-pytest]
 ignore_missing_imports = True

google/cloud/sql/connector/instance_connection_manager.py

@@ -44,9 +44,11 @@
 if TYPE_CHECKING:
     import pymysql
     import pg8000
+    import pytds
 logger = logging.getLogger(name=__name__)
 
 APPLICATION_NAME = "cloud-sql-python-connector"
+SERVER_PROXY_PORT = 3307
 
 # The default delay is set to 55 minutes since each ephemeral certificate is only
 # valid for an hour. This gives five minutes of buffer time.
@@ -351,6 +353,7 @@ async def _connect(self, driver: str, **kwargs: Any) -> Any:
         connect_func = {
             "pymysql": self._connect_with_pymysql,
             "pg8000": self._connect_with_pg8000,
+            "pytds": self._connect_with_pytds,
         }
 
         instance_data: InstanceMetadata = await self._current
@@ -391,7 +394,8 @@ def _connect_with_pymysql(
 
         # Create socket and wrap with context.
         sock = ctx.wrap_socket(
-            socket.create_connection((ip_address, 3307)), server_hostname=ip_address
+            socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+            server_hostname=ip_address,
         )
 
         # Create pymysql connection object and hand in pre-made connection
@@ -431,7 +435,43 @@ def _connect_with_pg8000(
             database=db,
             password=passwd,
             host=ip_address,
-            port=3307,
+            port=SERVER_PROXY_PORT,
             ssl_context=ctx,
             **kwargs,
         )
+
+    def _connect_with_pytds(
+        self, ip_address: str, ctx: ssl.SSLContext, **kwargs: Any
+    ) -> "pytds.Connection":
+        """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: pytds.Connection
+        :returns: A pytds Connection object for the Cloud SQL instance.
+        """
+        try:
+            import pytds
+        except ImportError:
+            raise ImportError(
+                'Unable to import module "pytds." Please install and try again.'
+            )
+        user = kwargs.pop("user")
+        db = kwargs.pop("db")
+        passwd = kwargs.pop("password")
+
+        # Create socket and wrap with context.
+        sock = ctx.wrap_socket(
+            socket.create_connection((ip_address, SERVER_PROXY_PORT)),
+            server_hostname=ip_address,
+        )
+        return pytds.connect(
+            ip_address, database=db, user=user, password=passwd, sock=sock, **kwargs
+        )

requirements-test.txt

@@ -2,4 +2,5 @@ pytest==6.2.2
 mock==4.0.3
 pytest-cov==2.11.1
 pytest-asyncio==0.14.0
-SQLAlchemy==1.3.23
+SQLAlchemy==1.3.23
+sqlalchemy-pytds

requirements.txt

@@ -2,7 +2,8 @@ aiohttp==3.7.4.post0
 cryptography==3.4.7
 PyMySQL==1.0.2
 pg8000==1.19.2
+# python-tds==1.10.0
+git+https://github.com/denisenkom/pytds.git
 pyopenssl==20.0.1
-pytest==6.2.3
 Requests==2.25.1
 google-api-python-client==2.2.0

setup.py

@@ -37,11 +37,10 @@
 version = version["__version__"]
 
 release_status = "Development Status :: 3 - Alpha"
-dependencies = [
+core_dependencies = [
     "aiohttp",
     "cryptography",
-    "PyMySQL",
-    "pytest",
+    "pyopenssl",
     "Requests",
     "google-api-python-client",
 ]
@@ -67,7 +66,12 @@
     platforms="Posix; MacOS X",
     packages=packages,
     namespace_packages=namespaces,
-    install_requires=dependencies,
+    install_requires=core_dependencies,
+    extras_require={
+        "pymysql": ["PyMySQL==1.0.2"],
+        "pg8000": ["pg8000==1.19.2"],
+        "pytds": ["python-tds @ git+https://github.com/denisenkom/pytds.git"]
+    },
     python_requires=">=3.6",
     include_package_data=True,
     zip_safe=False,

tests/system/test_pytds_connection.py

@@ -0,0 +1,76 @@
+""""
+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
+from typing import Generator
+
+import pytds
+import pytest
+import sqlalchemy
+from google.cloud.sql.connector import connector
+
+table_name = f"books_{uuid.uuid4().hex}"
+
+
+def init_connection_engine() -> sqlalchemy.engine.Engine:
+    def getconn() -> pytds.Connection:
+        conn = connector.connect(
+            os.environ["SQLSERVER_CONNECTION_NAME"],
+            "pytds",
+            user=os.environ["SQLSERVER_USER"],
+            password=os.environ["SQLSERVER_PASS"],
+            db=os.environ["SQLSERVER_DB"],
+        )
+        return conn
+
+    engine = sqlalchemy.create_engine(
+        "mssql+pytds://localhost",
+        creator=getconn,
+    )
+    engine.dialect.description_encoding = None
+    return engine
+
+
+@pytest.fixture(name="pool")
+def setup() -> Generator:
+    pool = init_connection_engine()
+
+    with pool.connect() as conn:
+        conn.execute(
+            f"CREATE TABLE {table_name}"
+            " ( id CHAR(20) NOT NULL, title TEXT NOT NULL );"
+        )
+
+    yield pool
+
+    with pool.connect() as conn:
+        conn.execute(f"DROP TABLE {table_name}")
+
+
+def test_pooled_connection_with_pytds(pool: sqlalchemy.engine.Engine) -> None:
+    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"]