[superset] branch master updated: feat: add SSL to new DB parameters (#14673)

[be...@apache.org]

This is an automated email from the ASF dual-hosted git repository.

beto pushed a commit to branch master
in repository https://gitbox.apache.org/repos/asf/superset.git


The following commit(s) were added to refs/heads/master by this push:
     new 971f588  feat: add SSL to new DB parameters (#14673)
971f588 is described below

commit 971f5883f0f5187c14345de95de5c65d967be44f
Author: Beto Dealmeida <ro...@dealmeida.net>
AuthorDate: Mon May 17 21:27:43 2021 -0700

    feat: add SSL to new DB parameters (#14673)
    
    * feat: add SSL to new DB parameters
    
    * Fix test
    
    * Raise if cls.encryption_parameters is empty
---
 superset/db_engine_specs/base.py        | 24 +++++++++++++++++++++---
 superset/db_engine_specs/postgres.py    |  2 ++
 tests/databases/api_tests.py            |  4 ++++
 tests/db_engine_specs/postgres_tests.py | 11 ++++++++---
 4 files changed, 35 insertions(+), 6 deletions(-)

diff --git a/superset/db_engine_specs/base.py b/superset/db_engine_specs/base.py
index 3289b4a..158162e 100644
--- a/superset/db_engine_specs/base.py
+++ b/superset/db_engine_specs/base.py
@@ -1304,6 +1304,9 @@ class BasicParametersSchema(Schema):
     query = fields.Dict(
         keys=fields.Str(), values=fields.Raw(), description=__("Additional parameters")
     )
+    encryption = fields.Boolean(
+        required=False, description=__("Use an encrypted connection to the database")
+    )
 
 
 class BasicParametersType(TypedDict, total=False):
@@ -1313,6 +1316,7 @@ class BasicParametersType(TypedDict, total=False):
     port: int
     database: str
     query: Dict[str, Any]
+    encryption: bool
 
 
 class BasicParametersMixin:
@@ -1339,8 +1343,18 @@ class BasicParametersMixin:
         "drivername://user:password@host:port/dbname[?key=value&key=value...]"
     )
 
+    # query parameter to enable encryption in the database connection
+    # for Postgres this would be `{"sslmode": "verify-ca"}`, eg.
+    encryption_parameters: Dict[str, str] = {}
+
     @classmethod
     def build_sqlalchemy_uri(cls, parameters: BasicParametersType) -> str:
+        query = parameters.get("query", {})
+        if parameters.get("encryption"):
+            if not cls.encryption_parameters:
+                raise Exception("Unable to build a URL with encryption enabled")
+            query.update(cls.encryption_parameters)
+
         return str(
             URL(
                 cls.drivername,
@@ -1349,13 +1363,16 @@ class BasicParametersMixin:
                 host=parameters["host"],
                 port=parameters["port"],
                 database=parameters["database"],
-                query=parameters.get("query", {}),
+                query=query,
             )
         )
 
-    @staticmethod
-    def get_parameters_from_uri(uri: str) -> BasicParametersType:
+    @classmethod
+    def get_parameters_from_uri(cls, uri: str) -> BasicParametersType:
         url = make_url(uri)
+        encryption = all(
+            item in url.query.items() for item in cls.encryption_parameters.items()
+        )
         return {
             "username": url.username,
             "password": url.password,
@@ -1363,6 +1380,7 @@ class BasicParametersMixin:
             "port": url.port,
             "database": url.database,
             "query": url.query,
+            "encryption": encryption,
         }
 
     @classmethod
diff --git a/superset/db_engine_specs/postgres.py b/superset/db_engine_specs/postgres.py
index b00a27a..5c8ff40 100644
--- a/superset/db_engine_specs/postgres.py
+++ b/superset/db_engine_specs/postgres.py
@@ -163,6 +163,8 @@ class PostgresEngineSpec(PostgresBaseEngineSpec, BasicParametersMixin):
     sqlalchemy_uri_placeholder = (
         "postgresql+psycopg2://user:password@host:port/dbname[?key=value&key=value...]"
     )
+    # https://www.postgresql.org/docs/9.1/libpq-ssl.html#LIBQ-SSL-CERTIFICATES
+    encryption_parameters = {"sslmode": "verify-ca"}
 
     max_column_name_length = 63
     try_remove_schema_from_table_name = False
diff --git a/tests/databases/api_tests.py b/tests/databases/api_tests.py
index 69f98cb..d7404f5 100644
--- a/tests/databases/api_tests.py
+++ b/tests/databases/api_tests.py
@@ -1392,6 +1392,10 @@ class TestDatabaseApi(SupersetTestCase):
                                 "description": "Database name",
                                 "type": "string",
                             },
+                            "encryption": {
+                                "description": "Use an encrypted connection to the database",
+                                "type": "boolean",
+                            },
                             "host": {
                                 "description": "Hostname or IP address",
                                 "type": "string",
diff --git a/tests/db_engine_specs/postgres_tests.py b/tests/db_engine_specs/postgres_tests.py
index e166726..6f6fa54 100644
--- a/tests/db_engine_specs/postgres_tests.py
+++ b/tests/db_engine_specs/postgres_tests.py
@@ -430,11 +430,12 @@ def test_base_parameters_mixin():
         "port": 5432,
         "database": "dbname",
         "query": {"foo": "bar"},
+        "encryption": True,
     }
     sqlalchemy_uri = PostgresEngineSpec.build_sqlalchemy_uri(parameters)
-    assert (
-        sqlalchemy_uri
-        == "postgresql+psycopg2://username:password@localhost:5432/dbname?foo=bar"
+    assert sqlalchemy_uri == (
+        "postgresql+psycopg2://username:password@localhost:5432/dbname?"
+        "foo=bar&sslmode=verify-ca"
     )
 
     parameters_from_uri = PostgresEngineSpec.get_parameters_from_uri(sqlalchemy_uri)
@@ -458,6 +459,10 @@ def test_base_parameters_mixin():
                 "additionalProperties": {},
             },
             "database": {"type": "string", "description": "Database name"},
+            "encryption": {
+                "type": "boolean",
+                "description": "Use an encrypted connection to the database",
+            },
         },
         "required": ["database", "host", "port", "username"],
     }