Adding conditions on claiming connections

astronomer · Nov 4, 2023 · d66a8a3 · d66a8a3
1 parent c00813a
commit d66a8a3
Show file tree

Hide file tree

Showing 2 changed files with 20 additions and 0 deletions.
diff --git a/cosmos/profiles/snowflake/user_encrypted_privatekey_env_variable.py b/cosmos/profiles/snowflake/user_encrypted_privatekey_env_variable.py
@@ -45,6 +45,16 @@ class SnowflakeEncryptedPrivateKeyPemProfileMapping(BaseProfileMapping):
         "private_key_passphrase": "password",
     }
 
+    def can_claim_connection(self) -> bool:
+        # Make sure this isn't a private key path credential
+        result = super().can_claim_connection()
+        if (
+            result
+            and self.conn.extra_dejson.get("private_key_file") is not None
+        ):
+            return False
+        return result
+
     @property
     def conn(self) -> Connection:
         """

diff --git a/cosmos/profiles/snowflake/user_encrypted_privatekey_file.py b/cosmos/profiles/snowflake/user_encrypted_privatekey_file.py
@@ -44,6 +44,16 @@ class SnowflakeEncryptedPrivateKeyFilePemProfileMapping(BaseProfileMapping):
         "private_key_path": "extra.private_key_file",
     }
 
+    def can_claim_connection(self) -> bool:
+        # Make sure this isn't a private key enviroment variable
+        result = super().can_claim_connection()
+        if (
+            result
+            and self.conn.extra_dejson.get("private_key_content") is not None
+        ):
+            return False
+        return result
+
     @property
     def conn(self) -> Connection:
         """