Repository: dbt-labs/dbt-snowflake (python)

Issue:
[Feature] Improve key-pair auth performance 
### Is this your first time submitting a feature request?

- [X] I have read the [expectations for open source contributors](https://docs.getdbt.com/docs/contributing/oss-expectations)
- [X] I have searched the existing issues, and I could not find an existing issue for this feature
- [X] I am requesting a straightforward extension of existing dbt-snowflake functionality, rather than a Big Idea better suited to a discussion

### Describe the feature

@VersusFacit's analysis:
```
Params:
Unix time command -- acceptable imprecision for the orders of magnitude we're dealing with here
5000 node project
average 2 runs each for each authentication method
dbt snowflake - user pass: 427.83s user 42.43s system 16% cpu 46:49.12 total
dbt snowflake - key pair: 1011.76s user 44.96s system 32% cpu 53:42.99 total
```
400 vs 1000 is quite a dramatic difference!

Avenues for investigation:
- We want to first try caching the keyfile contents
- Redo timing analysis to see if there's any benefit
- Look at any extra calls being made to snowflake
- If that fails, create a ticket for further investigation


Pull request:
Improve run times when using key pair auth by caching the private key
resolves #1082

### Problem

We evaluate and validate the private key on every connection when using key pair auth despite the fact that the key does not change during a dbt run.

### Solution

- create a function that uses caching outside of the `Credentials` class and pass the parameters to this function
- use `functools.lru_cache` if we're using Python 3.8
- use `functools.cache` if we're using Python 3.9
- remove the check on Python version when we retire Python 3.8 support

### Checklist

- [x] I have read [the contributing guide](https://github.com/dbt-labs/dbt-core/blob/main/CONTRIBUTING.md) and understand what's expected of me
- [x] I have run this code in development and it appears to resolve the stated issue
- [x] This PR includes tests, or tests are not required/relevant for this PR
- [x] This PR has no interface changes (e.g. macros, cli, logs, json artifacts, config files, adapter interface, etc) or this PR has already received feedback and approval from Product or DX


Patch:
diff --git a/.changes/unreleased/Features-20240710-172345.yaml b/.changes/unreleased/Features-20240710-172345.yaml
new file mode 100644
index 0000000..e68f638
--- /dev/null
+++ b/.changes/unreleased/Features-20240710-172345.yaml
@@ -0,0 +1,6 @@
+kind: Features
+body: Improve run times when using key pair auth by caching the private key
+time: 2024-07-10T17:23:45.046905-04:00
+custom:
+  Author: mikealfare aranke
+  Issue: "1082"
diff --git a/dbt/adapters/snowflake/auth.py b/dbt/adapters/snowflake/auth.py
new file mode 100644
index 0000000..e914b6f
--- /dev/null
+++ b/dbt/adapters/snowflake/auth.py
@@ -0,0 +1,57 @@
+import base64
+import sys
+from typing import Optional
+
+if sys.version_info < (3, 9):
+    from functools import lru_cache
+
+    cache = lru_cache(maxsize=None)
+else:
+    from functools import cache
+
+from cryptography.hazmat.backends import default_backend
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
+
+
+@cache
+def private_key_from_string(
+    private_key_string: str, passphrase: Optional[str] = None
+) -> RSAPrivateKey:
+
+    if passphrase:
+        encoded_passphrase = passphrase.encode()
+    else:
+        encoded_passphrase = None
+
+    if private_key_string.startswith("-"):
+        return serialization.load_pem_private_key(
+            data=bytes(private_key_string, "utf-8"),
+            password=encoded_passphrase,
+            backend=default_backend(),
+        )
+    return serialization.load_der_private_key(
+        data=base64.b64decode(private_key_string),
+        password=encoded_passphrase,
+        backend=default_backend(),
+    )
+
+
+@cache
+def private_key_from_file(
+    private_key_path: str, passphrase: Optional[str] = None
+) -> RSAPrivateKey:
+
+    if passphrase:
+        encoded_passphrase = passphrase.encode()
+    else:
+        encoded_passphrase = None
+
+    with open(private_key_path, "rb") as file:
+        private_key_bytes = file.read()
+
+    return serialization.load_pem_private_key(
+        data=private_key_bytes,
+        password=encoded_passphrase,
+        backend=default_backend(),
+    )
diff --git a/dbt/adapters/snowflake/connections.py b/dbt/adapters/snowflake/connections.py
index 6e9a5aa..a9e83cd 100644
--- a/dbt/adapters/snowflake/connections.py
+++ b/dbt/adapters/snowflake/connections.py
@@ -1,6 +1,14 @@
 import base64
 import datetime
 import os
+import sys
+
+if sys.version_info < (3, 9):
+    from functools import lru_cache
+
+    cache = lru_cache(maxsize=None)
+else:
+    from functools import cache
 
 import pytz
 import re
@@ -11,8 +19,8 @@ from time import sleep
 
 from typing import Optional, Tuple, Union, Any, List, Iterable, TYPE_CHECKING
 
-from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 import requests
 import snowflake.connector
 import snowflake.connector.constants
@@ -44,6 +52,8 @@ from dbt_common.events.functions import warn_or_error
 from dbt.adapters.events.types import AdapterEventWarning, AdapterEventError
 from dbt_common.ui import line_wrap_message, warning_tag
 
+from dbt.adapters.snowflake.auth import private_key_from_file, private_key_from_string
+
 if TYPE_CHECKING:
     import agate
 
@@ -63,6 +73,15 @@ ERROR_REDACTION_PATTERNS = {
 }
 
 
+@cache
+def snowflake_private_key(private_key: RSAPrivateKey) -> bytes:
+    return private_key.private_bytes(
+        encoding=serialization.Encoding.DER,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.NoEncryption(),
+    )
+
+
 @dataclass
 class SnowflakeAdapterResponse(AdapterResponse):
     query_id: str = ""
@@ -273,44 +292,17 @@ class SnowflakeCredentials(Credentials):
             )
         return result_json["access_token"]
 
-    def _get_private_key(self):
+    def _get_private_key(self) -> Optional[bytes]:
         """Get Snowflake private key by path, from a Base64 encoded DER bytestring or None."""
         if self.private_key and self.private_key_path:
             raise DbtConfigError("Cannot specify both `private_key`  and `private_key_path`")
-
-        if self.private_key_passphrase:
-            encoded_passphrase = self.private_key_passphrase.encode()
-        else:
-            encoded_passphrase = None
-
-        if self.private_key:
-            if self.private_key.startswith("-"):
-                p_key = serialization.load_pem_private_key(
-                    data=bytes(self.private_key, "utf-8"),
-                    password=encoded_passphrase,
-                    backend=default_backend(),
-                )
-
-            else:
-                p_key = serialization.load_der_private_key(
-                    data=base64.b64decode(self.private_key),
-                    password=encoded_passphrase,
-                    backend=default_backend(),
-                )
-
+        elif self.private_key:
+            private_key = private_key_from_string(self.private_key, self.private_key_passphrase)
         elif self.private_key_path:
-            with open(self.private_key_path, "rb") as key:
-                p_key = serialization.load_pem_private_key(
-                    key.read(), password=encoded_passphrase, backend=default_backend()
-                )
+            private_key = private_key_from_file(self.private_key_path, self.private_key_passphrase)
         else:
             return None
-
-        return p_key.private_bytes(
-            encoding=serialization.Encoding.DER,
-            format=serialization.PrivateFormat.PKCS8,
-            encryption_algorithm=serialization.NoEncryption(),
-        )
+        return snowflake_private_key(private_key)
 
 
 class SnowflakeConnectionManager(SQLConnectionManager):


Test patch:
diff --git a/tests/functional/oauth/test_jwt.py b/tests/functional/auth_tests/test_jwt.py
similarity index 100%
rename from tests/functional/oauth/test_jwt.py
rename to tests/functional/auth_tests/test_jwt.py
diff --git a/tests/functional/auth_tests/test_key_pair.py b/tests/functional/auth_tests/test_key_pair.py
new file mode 100644
index 0000000..6d3254f
--- /dev/null
+++ b/tests/functional/auth_tests/test_key_pair.py
@@ -0,0 +1,26 @@
+import os
+
+from dbt.tests.util import run_dbt
+import pytest
+
+
+class TestKeyPairAuth:
+    @pytest.fixture(scope="class", autouse=True)
+    def dbt_profile_target(self):
+        return {
+            "type": "snowflake",
+            "threads": 4,
+            "account": os.getenv("SNOWFLAKE_TEST_ACCOUNT"),
+            "user": os.getenv("SNOWFLAKE_TEST_USER"),
+            "private_key": os.getenv("SNOWFLAKE_TEST_PRIVATE_KEY"),
+            "private_key_passphrase": os.getenv("SNOWFLAKE_TEST_PRIVATE_KEY_PASSPHRASE"),
+            "database": os.getenv("SNOWFLAKE_TEST_DATABASE"),
+            "warehouse": os.getenv("SNOWFLAKE_TEST_WAREHOUSE"),
+        }
+
+    @pytest.fixture(scope="class")
+    def models(self):
+        return {"my_model.sql": "select 1 as id"}
+
+    def test_connection(self, project):
+        run_dbt()
diff --git a/tests/functional/oauth/test_oauth.py b/tests/functional/auth_tests/test_oauth.py
similarity index 100%
rename from tests/functional/oauth/test_oauth.py
rename to tests/functional/auth_tests/test_oauth.py
diff --git a/tests/unit/test_private_keys.py b/tests/unit/test_private_keys.py
new file mode 100644
index 0000000..59b8522
--- /dev/null
+++ b/tests/unit/test_private_keys.py
@@ -0,0 +1,61 @@
+import os
+import tempfile
+from typing import Generator
+
+from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric import rsa
+import pytest
+
+from dbt.adapters.snowflake.auth import private_key_from_file, private_key_from_string
+
+
+PASSPHRASE = "password1234"
+
+
+def serialize(private_key: rsa.RSAPrivateKey) -> bytes:
+    return private_key.private_bytes(
+        serialization.Encoding.DER,
+        serialization.PrivateFormat.PKCS8,
+        serialization.NoEncryption(),
+    )
+
+
+@pytest.fixture(scope="session")
+def private_key() -> rsa.RSAPrivateKey:
+    return rsa.generate_private_key(public_exponent=65537, key_size=2048)
+
+
+@pytest.fixture(scope="session")
+def private_key_string(private_key) -> str:
+    private_key_bytes = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.BestAvailableEncryption(PASSPHRASE.encode()),
+    )
+    return private_key_bytes.decode("utf-8")
+
+
+@pytest.fixture(scope="session")
+def private_key_file(private_key) -> Generator[str, None, None]:
+    private_key_bytes = private_key.private_bytes(
+        encoding=serialization.Encoding.PEM,
+        format=serialization.PrivateFormat.PKCS8,
+        encryption_algorithm=serialization.BestAvailableEncryption(PASSPHRASE.encode()),
+    )
+    file = tempfile.NamedTemporaryFile()
+    file.write(private_key_bytes)
+    file.seek(0)
+    yield file.name
+    file.close()
+
+
+def test_private_key_from_string_pem(private_key_string, private_key):
+    assert isinstance(private_key_string, str)
+    calculated_private_key = private_key_from_string(private_key_string, PASSPHRASE)
+    assert serialize(calculated_private_key) == serialize(private_key)
+
+
+def test_private_key_from_file(private_key_file, private_key):
+    assert os.path.exists(private_key_file)
+    calculated_private_key = private_key_from_file(private_key_file, PASSPHRASE)
+    assert serialize(calculated_private_key) == serialize(private_key)