enhanced support for snowflake ssh key and presto cert

Author: Matthias Ekundayo

.gitignore

@@ -141,3 +141,6 @@ benchmark_*.png
 
 # IntelliJ
 .idea
+
+# VS Code
+.vscode

data_diff/databases/presto.py

@@ -35,10 +35,28 @@ class Presto(Database):
     }
     ROUNDS_ON_PREC_LOSS = True
 
-    def __init__(self, **kw):
+    def __init__(self, host, port, user, password, *, catalog, schema=None,  **kw):
         prestodb = import_presto()
+        self.args = dict(
+            host=host, port=port, user=user, catalog=catalog, schema=schema, **kw
+        )  # include port if specified
+
+        if (
+            "cert" in self.args
+        ):  # cert used after connection to verify session, but keyword is not valid so remove from connection params
+            self.args.pop("cert")
+
+        if "auth" in kw and kw.get("auth") == "basic":  # if auth=basic, add basic authenticator for Presto
+            self.args["auth"] = prestodb.auth.BasicAuthentication(user, password)
+
+        if schema:  # if schema was specified in URI, override default
+            self.default_schema = schema
+        self._conn = prestodb.dbapi.connect(**self.args)
+        # self._conn = prestodb.dbapi.connect(**kw)
+
+        if "cert" in kw:  # if a certificate was specified in URI, verify session with cert
+            self._conn._http_session.verify = kw.get("cert")
 
-        self._conn = prestodb.dbapi.connect(**kw)
 
     def quote(self, s: str):
         return f'"{s}"'

data_diff/databases/snowflake.py

@@ -7,8 +7,10 @@
 @import_helper("snowflake")
 def import_snowflake():
     import snowflake.connector
-
-    return snowflake
+    from cryptography.hazmat.primitives import serialization
+    from cryptography.hazmat.backends import default_backend
+    
+    return snowflake, serialization, default_backend
 
 
 class Snowflake(Database):
@@ -25,8 +27,19 @@ class Snowflake(Database):
     }
     ROUNDS_ON_PREC_LOSS = False
 
-    def __init__(self, *, schema: str, **kw):
-        snowflake = import_snowflake()
+    def __init__(self,
+        account: str,
+        user: str,
+        *,
+        warehouse: str,
+        schema: str,
+        database: str,
+        role: str = None,
+        _port: int = None,
+        password: str = None,  # default to None incase ssh key is used
+        **kw,
+    ):
+        snowflake, serialization, default_backend = import_snowflake()
         logging.getLogger("snowflake.connector").setLevel(logging.WARNING)
 
         # Got an error: snowflake.connector.network.RetryRequest: could not find io module state (interpreter shutdown?)
@@ -35,10 +48,38 @@ def __init__(self, *, schema: str, **kw):
         logging.getLogger("snowflake.connector.network").disabled = True
 
         assert '"' not in schema, "Schema name should not contain quotes!"
-        self._conn = snowflake.connector.connect(
-            schema=f'"{schema}"',
-            **kw,
-        )
+        if (
+            not password and "key" in kw
+        ):  # if private keys are used instead of password for Snowflake connection, read in key from path specified and pass as "private_key" to connector.
+            with open(kw.get("key"), "rb") as key:
+                p_key = serialization.load_pem_private_key(key.read(), password=None, backend=default_backend())
+
+            pkb = p_key.private_bytes(
+                encoding=serialization.Encoding.DER,
+                format=serialization.PrivateFormat.PKCS8,
+                encryption_algorithm=serialization.NoEncryption(),
+            )
+            self._conn = snowflake.connector.connect(
+                user=user,
+                private_key=pkb,  # replaces password
+                account=account,
+                role=role,
+                database=database,
+                warehouse=warehouse,
+                schema=f'"{schema}"',
+                **kw,
+            )
+        else:  # otherwise use password for connection
+            self._conn = snowflake.connector.connect(
+                user=user,
+                password=password,
+                account=account,
+                role=role,
+                database=database,
+                warehouse=warehouse,
+                schema=f'"{schema}"',
+                **kw,
+            )
 
         self.default_schema = schema
 

docs/requirements.txt

@@ -4,6 +4,7 @@ sphinx_markdown_tables
 sphinx-copybutton
 sphinx-rtd-theme
 recommonmark
+cryptography
 
 # Requirements. TODO Use poetry instead of this redundant list
 data_diff