added exception handling for presto when auth=basic

Author: Matthias Ekundayo

data_diff/databases/presto.py

@@ -40,15 +40,17 @@ def __init__(self, **kw):
         if kw.get("schema"):
             self.default_schema = kw.get("schema")
 
-        if (
-            "password" in kw and "auth" in kw and kw.get("auth") == "basic"
-        ):  # if auth=basic, add basic authenticator for Presto
-            kw["auth"] = prestodb.auth.BasicAuthentication(kw.get("user"), kw.get("password"))
-            kw.pop("password")
+        try:
+            # checks if user and password are missing when auth=basic
+            kw.get("auth") == "basic" and "user" in kw and "password" in kw
+            # if auth=basic, add basic authenticator for Presto
+            kw["auth"] = prestodb.auth.BasicAuthentication(kw.pop("user"), kw.pop("password"))
+        except:
+            raise ValueError("User or password cannot be missing if auth==basic")
 
         if "cert" in kw:  # if a certificate was specified in URI, verify session with cert
-            cert = kw.get("cert")
-            kw.pop("cert")
+            cert = kw.pop("cert")
+
             self._conn = prestodb.dbapi.connect(**kw)
             self._conn._http_session.verify = cert
 

data_diff/databases/snowflake.py

@@ -38,20 +38,22 @@ def __init__(self, *, schema: str, **kw):
 
         assert '"' not in schema, "Schema name should not contain quotes!"
         if (
-            "password" not in kw 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.
+            "key" in kw
+        ):  # if private keys are used 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())
+                p_key = serialization.load_pem_private_key(
+                    key.read(),
+                    password=None if not kw.get("password") else kw.pop("password"),
+                    backend=default_backend(),
+                )
 
-            pkb = p_key.private_bytes(
+            kw["private_key"] = p_key.private_bytes(
                 encoding=serialization.Encoding.DER,
                 format=serialization.PrivateFormat.PKCS8,
                 encryption_algorithm=serialization.NoEncryption(),
             )
-            self._conn = snowflake.connector.connect(private_key=pkb, schema=f'"{schema}"', **kw)  # replaces password
 
-        else:  # otherwise use password for connection
-            self._conn = snowflake.connector.connect(schema=f'"{schema}"', **kw)
+        self._conn = snowflake.connector.connect(schema=f'"{schema}"', **kw)
 
         self.default_schema = schema