Adjust connect arguments to enable SSL using SQLAlchemy DB URI

Author: amotl

CHANGES.txt

@@ -13,6 +13,8 @@ Unreleased
 
 - Adjusted connect arguments to accept credentials within the HTTP URI.
 
+- Adjusted connect arguments to enable SSL using SQLAlchemy DB URI.
+
 2020/09/28 0.26.0
 =================
 

src/crate/client/sqlalchemy/dialect.py

@@ -25,6 +25,7 @@
 from sqlalchemy import types as sqltypes
 from sqlalchemy.engine import default, reflection
 from sqlalchemy.sql import functions
+from sqlalchemy.util import asbool, to_list
 
 from .compiler import (
     CrateCompiler,
@@ -194,8 +195,16 @@ def connect(self, host=None, port=None, *args, **kwargs):
             server = '{0}:{1}'.format(host, port or '4200')
         if 'servers' in kwargs:
             server = kwargs.pop('servers')
-        if server:
-            return self.dbapi.connect(servers=server, **kwargs)
+        servers = to_list(server)
+        if servers:
+            # Evaluate "ssl" connection URI query parameter.
+            # TODO: Evaluate more parameters like `ssl_ca`, `ssl_key`, `ssl_cert`,
+            #       `ssl_capath` and `ssl_cipher`.
+            if "ssl" in kwargs:
+                use_ssl = asbool(kwargs.pop("ssl"))
+                if use_ssl:
+                    servers = ["https://" + server for server in servers]
+            return self.dbapi.connect(servers=servers, **kwargs)
         return self.dbapi.connect(**kwargs)
 
     def _get_default_schema_name(self, connection):

src/crate/client/sqlalchemy/tests/connection_test.py

@@ -21,28 +21,59 @@
 
 from unittest import TestCase
 import sqlalchemy as sa
+from sqlalchemy.exc import NoSuchModuleError
 
 
 class SqlAlchemyConnectionTest(TestCase):
 
-    def setUp(self):
-        self.engine = sa.create_engine('crate://')
-        self.connection = self.engine.connect()
+    def test_connection_server_uri_unknown_sa_plugin(self):
+        with self.assertRaises(NoSuchModuleError):
+            sa.create_engine("foobar://otherhost:19201")
 
     def test_default_connection(self):
         engine = sa.create_engine('crate://')
         conn = engine.raw_connection()
         self.assertEqual("<Connection <Client ['http://127.0.0.1:4200']>>",
                          repr(conn.connection))
 
-    def test_connection_server(self):
+    def test_connection_server_uri_http(self):
         engine = sa.create_engine(
             "crate://otherhost:19201")
         conn = engine.raw_connection()
         self.assertEqual("<Connection <Client ['http://otherhost:19201']>>",
                          repr(conn.connection))
 
-    def test_connection_multiple_server(self):
+    def test_connection_server_uri_https(self):
+        engine = sa.create_engine(
+            "crate://otherhost:19201/?ssl=true")
+        conn = engine.raw_connection()
+        self.assertEqual("<Connection <Client ['https://otherhost:19201']>>",
+                         repr(conn.connection))
+
+    def test_connection_server_uri_invalid_port(self):
+        with self.assertRaises(ValueError) as context:
+            sa.create_engine("crate://foo:bar")
+        self.assertTrue("invalid literal for int() with base 10: 'bar'" in str(context.exception))
+
+    def test_connection_server_uri_https_with_trusted_user(self):
+        engine = sa.create_engine(
+            "crate://foo@otherhost:19201/?ssl=true")
+        conn = engine.raw_connection()
+        self.assertEqual("<Connection <Client ['https://otherhost:19201']>>",
+                         repr(conn.connection))
+        self.assertEqual(conn.connection.client.username, "foo")
+        self.assertEqual(conn.connection.client.password, None)
+
+    def test_connection_server_uri_https_with_credentials(self):
+        engine = sa.create_engine(
+            "crate://foo:bar@otherhost:19201/?ssl=true")
+        conn = engine.raw_connection()
+        self.assertEqual("<Connection <Client ['https://otherhost:19201']>>",
+                         repr(conn.connection))
+        self.assertEqual(conn.connection.client.username, "foo")
+        self.assertEqual(conn.connection.client.password, "bar")
+
+    def test_connection_multiple_server_http(self):
         engine = sa.create_engine(
             "crate://", connect_args={
                 'servers': ['localhost:4201', 'localhost:4202']
@@ -53,3 +84,16 @@ def test_connection_multiple_server(self):
             "<Connection <Client ['http://localhost:4201', " +
             "'http://localhost:4202']>>",
             repr(conn.connection))
+
+    def test_connection_multiple_server_https(self):
+        engine = sa.create_engine(
+            "crate://", connect_args={
+                'servers': ['localhost:4201', 'localhost:4202'],
+                'ssl': True,
+            }
+        )
+        conn = engine.raw_connection()
+        self.assertEqual(
+            "<Connection <Client ['https://localhost:4201', " +
+            "'https://localhost:4202']>>",
+            repr(conn.connection))