Post-merge fixes for asymmetric encryption

GitOrigin-RevId: eb9c4c29e3627380aaf2c4e9c7867e7e77de3061
gretelai · May 14, 2024 · 279a99d · 279a99d
1 parent 1895674
commit 279a99d
Show file tree

Hide file tree

Showing 4 changed files with 13 additions and 5 deletions.
diff --git a/requirements/base.txt b/requirements/base.txt
@@ -5,6 +5,7 @@ docker==6.1.2
 kubernetes==28.1.0
 opentelemetry-distro==0.44b0
 opentelemetry-exporter-prometheus==0.44b0
+pycryptodome>=3.19,<4
 pydantic==1.10.13
 python_dateutil>=2.8.0
 pyyaml==6.0.1

diff --git a/src/gretel_client/_hybrid/asymmetric.py b/src/gretel_client/_hybrid/asymmetric.py
@@ -19,12 +19,13 @@
 
 class AsymmetricCredentialsEncryption(BaseCredentialsEncryption):
 
-    _projects_api: ProjectsApi
-    _asymmetric_key_metadata: ConfigAsymmetricKeyMetadata
+    _projects_api: Optional[ProjectsApi]
+    _asymmetric_key_metadata: Optional[ConfigAsymmetricKeyMetadata]
 
     def __init__(
         self,
-        projects_api: ProjectsApi,
+        *,
+        projects_api: Optional[ProjectsApi] = None,
         asymmetric_key_metadata: Optional[ConfigAsymmetricKeyMetadata] = None,
     ):
         self._projects_api = projects_api
@@ -37,6 +38,10 @@ def apply(self, credentials: dict, *, project_guid: Optional[str] = None):
                 raise ValueError(
                     "can not apply asymmetric encryption for connections not specifying a project ID"
                 )
+            if not self._projects_api:
+                raise ValueError(
+                    "encryption mechanism is not configured for dynamic retrieval of asymmetric key"
+                )
 
             project = self._projects_api.get_project(
                 project_guid, expand=["cluster"]

diff --git a/src/gretel_client/_hybrid/config.py b/src/gretel_client/_hybrid/config.py
@@ -131,7 +131,9 @@ def __init__(
         self._creds_encryption = (
             creds_encryption
             if creds_encryption is not None
-            else AsymmetricCredentialsEncryption(session.get_v1_api(ProjectsV1Api))
+            else AsymmetricCredentialsEncryption(
+                projects_api=session.get_v1_api(ProjectsV1Api)
+            )
         )
         self._deployment_user = deployment_user
         self._default_cluster_guid = default_cluster_guid

diff --git a/src/gretel_client/cli/connections.py b/src/gretel_client/cli/connections.py
@@ -121,7 +121,7 @@ def create(
                 )
 
             encryption_provider = AsymmetricCredentialsEncryption(
-                projects_v1_api, asymmetric_key_metadata=cluster.config.asymmetric_key
+                asymmetric_key_metadata=cluster.config.asymmetric_key
             )
 
     if conn.get("encrypted_credentials") is not None: