SSL: Use SSLMode enum class from asyncpg

Author: amotl

src/sqlalchemy_cratedb/dialect.py

@@ -25,6 +25,7 @@
 
 from sqlalchemy import types as sqltypes
 from sqlalchemy.engine import default, reflection
+from sqlalchemy.exc import SQLAlchemyError
 from sqlalchemy.sql import functions
 from sqlalchemy.util import asbool, to_list
 
@@ -35,6 +36,7 @@
 )
 from .sa_version import SA_1_4, SA_2_0, SA_VERSION
 from .type import FloatVector, ObjectArray, ObjectType
+from .util import SSLMode
 
 TYPES_MAP = {
     "boolean": sqltypes.Boolean,
@@ -228,8 +230,7 @@ def connect(self, host=None, port=None, *args, **kwargs):
             server = kwargs.pop("servers")
         servers = to_list(server)
 
-        # Process SSL options, old and new.
-        # TODO: Switch to the canonical default `sslmode=prefer` later.
+        # Process legacy SSL option `ssl`.
         if "ssl" in kwargs:
             warnings.warn(
                 "The `ssl=true` option will be deprecated, "
@@ -238,11 +239,25 @@ def connect(self, host=None, port=None, *args, **kwargs):
                 stacklevel=2,
             )
         use_ssl = asbool(kwargs.pop("ssl", False))
-        sslmode = kwargs.pop("sslmode", "disable")
-        if sslmode in ["allow", "prefer", "require", "verify-ca", "verify-full"]:
-            use_ssl = True
-        if sslmode in ["allow", "prefer", "require"]:
-            kwargs["verify_ssl_cert"] = False
+
+        # Process new SSL option `sslmode`.
+        # Please consult https://www.postgresql.org/docs/18/libpq-connect.html.
+        if "sslmode" in kwargs:
+            try:
+                sslmode = SSLMode.parse(kwargs.pop("sslmode"))
+            except AttributeError as exc:
+                modes = ", ".join(SSLMode.modes)
+                raise SQLAlchemyError(
+                    "`sslmode` parameter must be one of: {}".format(modes)
+                ) from exc
+            if sslmode < SSLMode.allow:
+                use_ssl = False
+            else:
+                use_ssl = True
+                if sslmode >= SSLMode.verify_ca:
+                    kwargs["verify_ssl_cert"] = True
+                else:
+                    kwargs["verify_ssl_cert"] = False
 
         if not servers:
             servers = [self.dbapi.http.Client.default_server.replace("http://", "")]

src/sqlalchemy_cratedb/util.py

@@ -0,0 +1,27 @@
+import enum
+
+from sqlalchemy.util import classproperty
+
+
+class SSLMode(enum.IntEnum):
+    """
+    SSLMode class from asyncpg, with a little improvement.
+    https://github.com/MagicStack/asyncpg/blob/v0.31.0/asyncpg/connect_utils.py#L36-L48
+    """
+
+    disable = 0
+    allow = 1
+    prefer = 2
+    require = 3
+    verify_ca = 4
+    verify_full = 5
+
+    @classmethod
+    def parse(cls, sslmode):
+        if isinstance(sslmode, cls):
+            return sslmode
+        return getattr(cls, sslmode.replace("-", "_"))
+
+    @classproperty
+    def modes(cls):
+        return [m.name.replace("_", "-") for m in cls]

tests/connection_test.py

@@ -24,7 +24,7 @@
 
 import pytest
 import sqlalchemy as sa
-from sqlalchemy.exc import NoSuchModuleError
+from sqlalchemy.exc import NoSuchModuleError, SQLAlchemyError
 
 from sqlalchemy_cratedb import SA_1_4, SA_VERSION
 from tests.util import ExtraAssertions
@@ -104,6 +104,15 @@ def test_connection_server_uri_https_sslmode_disabled(self):
         self.assertEqual(["http://otherhost:19201"], servers)
         engine.dispose()
 
+    def test_connection_server_uri_https_sslmode_invalid(self):
+        with pytest.raises(SQLAlchemyError) as exc_info:
+            engine = sa.create_engine("crate://otherhost:19201/?sslmode=foo")
+            engine.raw_connection()
+        exc_info.match(
+            "`sslmode` parameter must be one of: "
+            "disable, allow, prefer, require, verify-ca, verify-full"
+        )
+
     def test_connection_server_uri_invalid_port(self):
         with self.assertRaises(ValueError) as context:
             sa.create_engine("crate://foo:bar")