From 32a16b059d7eed70eaa1e42930477d51be2bbe26 Mon Sep 17 00:00:00 2001
From: Ayush Shah <ayush@getcollate.io>
Date: Fri, 22 Apr 2022 20:07:06 +0530
Subject: [PATCH] Modify Snowflake and Databricks schema (#4367)

---
 .../database/databricksConnection.json        | 10 ++--
 .../database/snowflakeConnection.json         | 11 +++--
 ingestion/src/metadata/utils/connections.py   | 47 ++++++++++++++++++-
 3 files changed, 57 insertions(+), 11 deletions(-)

diff --git a/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/databricksConnection.json b/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/databricksConnection.json
index e7173056c44..6b55fc1cd39 100644
--- a/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/databricksConnection.json
+++ b/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/databricksConnection.json
@@ -59,6 +59,11 @@
       "type": "string",
       "format": "password"
     },
+    "httpPath": {
+      "title": "Http Path",
+      "description": "Databricks compute resources URL",
+      "type": "string"
+    },
     "connectionOptions": {
       "title": "Connection Options",
       "$ref": "../connectionBasicType.json#/definitions/connectionOptions"
@@ -77,8 +82,5 @@
     }
   },
   "additionalProperties": false,
-  "required": [
-    "hostPort",
-    "token"
-  ]
+  "required": ["hostPort", "token"]
 }
diff --git a/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/snowflakeConnection.json b/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/snowflakeConnection.json
index 37574bc44ac..099c2f8d79f 100644
--- a/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/snowflakeConnection.json
+++ b/catalog-rest-service/src/main/resources/json/schema/entity/services/connections/database/snowflakeConnection.json
@@ -68,6 +68,11 @@
       "description": "Snowflake warehouse.",
       "type": "string"
     },
+    "privateKey": {
+      "title": "Private Key",
+      "description": "Connection to Snowflake instance via Private Key",
+      "type": "string"
+    },
     "connectionOptions": {
       "title": "Connection Options",
       "$ref": "../connectionBasicType.json#/definitions/connectionOptions"
@@ -89,9 +94,5 @@
     }
   },
   "additionalProperties": false,
-  "required": [
-    "hostPort",
-    "username",
-    "account"
-  ]
+  "required": ["hostPort", "username", "account"]
 }
diff --git a/ingestion/src/metadata/utils/connections.py b/ingestion/src/metadata/utils/connections.py
index d717b7ea1cd..79d9b31bfef 100644
--- a/ingestion/src/metadata/utils/connections.py
+++ b/ingestion/src/metadata/utils/connections.py
@@ -29,14 +29,17 @@ from metadata.generated.schema.entity.services.connections.connectionBasicType i
 from metadata.generated.schema.entity.services.connections.database.bigQueryConnection import (
     BigQueryConnection,
 )
+from metadata.generated.schema.entity.services.connections.database.databricksConnection import (
+    DatabricksConnection,
+)
 from metadata.generated.schema.entity.services.connections.database.dynamoDBConnection import (
     DynamoDBConnection,
 )
 from metadata.generated.schema.entity.services.connections.database.glueConnection import (
     GlueConnection,
 )
-from metadata.generated.schema.entity.services.connections.database.sampleDataConnection import (
-    SampleDataConnection,
+from metadata.generated.schema.entity.services.connections.database.snowflakeConnection import (
+    SnowflakeConnection,
 )
 from metadata.utils.aws_client import AWSClient, DynamoClient, GlueClient
 from metadata.utils.credentials import set_google_credentials
@@ -83,6 +86,46 @@ def get_connection(
     return create_generic_connection(connection, verbose)
 
 
+@get_connection.register
+def _(connection: DatabricksConnection, verbose: bool = False):
+    args = connection.connectionArguments
+    if not args:
+        connection.connectionArguments = dict()
+        connection.connectionArguments["http_path"] = connection.httpPath
+    return create_generic_connection(connection, verbose)
+
+
+@get_connection.register
+def _(connection: SnowflakeConnection, verbose: bool = False):
+    if connection.privateKey:
+        import os
+
+        from cryptography.hazmat.backends import default_backend
+        from cryptography.hazmat.primitives import serialization
+
+        snowflake_private_key_passphrase = os.environ.get(
+            "SNOWFLAKE_PRIVATE_KEY_PASSPHRASE", ""
+        )
+        if not snowflake_private_key_passphrase:
+            logger.warning(
+                "Snowflake Private Key Passphrase not found, replacing it with empty string"
+            )
+        p_key = serialization.load_pem_private_key(
+            bytes(connection.privateKey, "utf-8"),
+            password=snowflake_private_key_passphrase.encode(),
+            backend=default_backend(),
+        )
+        pkb = p_key.private_bytes(
+            encoding=serialization.Encoding.DER,
+            format=serialization.PrivateFormat.PKCS8,
+            encryption_algorithm=serialization.NoEncryption(),
+        )
+        if not connection.connectionArguments:
+            connection.connectionArguments = dict()
+            connection.connectionArguments["private_key"] = pkb
+    return create_generic_connection(connection, verbose)
+
+
 @get_connection.register
 def _(connection: BigQueryConnection, verbose: bool = False):
     """