chore: typing

Author: pnilan

airbyte_cdk/sources/declarative/auth/jwt.py

@@ -6,18 +6,27 @@
 import json
 from dataclasses import InitVar, dataclass
 from datetime import datetime
-from typing import Any, Mapping, Optional, Union
+from typing import Any, Mapping, Optional, Union, cast
 
 import jwt
 from cryptography.hazmat.backends import default_backend
 from cryptography.hazmat.primitives import serialization
+from cryptography.hazmat.primitives.asymmetric.ec import EllipticCurvePrivateKey
+from cryptography.hazmat.primitives.asymmetric.ed448 import Ed448PrivateKey
+from cryptography.hazmat.primitives.asymmetric.ed25519 import Ed25519PrivateKey
+from cryptography.hazmat.primitives.asymmetric.rsa import RSAPrivateKey
 from cryptography.hazmat.primitives.asymmetric.types import PrivateKeyTypes
 
 from airbyte_cdk.sources.declarative.auth.declarative_authenticator import DeclarativeAuthenticator
 from airbyte_cdk.sources.declarative.interpolation.interpolated_boolean import InterpolatedBoolean
 from airbyte_cdk.sources.declarative.interpolation.interpolated_mapping import InterpolatedMapping
 from airbyte_cdk.sources.declarative.interpolation.interpolated_string import InterpolatedString
 
+# Type alias for keys that JWT library accepts
+JwtKeyTypes = Union[
+    RSAPrivateKey, EllipticCurvePrivateKey, Ed25519PrivateKey, Ed448PrivateKey, str, bytes
+]
+
 
 class JwtAlgorithm(str):
     """
@@ -158,18 +167,21 @@ def _get_jwt_payload(self) -> dict[str, Any]:
         payload["nbf"] = nbf
         return payload
 
-    def _get_secret_key(self) -> PrivateKeyTypes | str | bytes:
+    def _get_secret_key(self) -> JwtKeyTypes:
         """
         Returns the secret key used to sign the JWT.
         """
         secret_key: str = self._secret_key.eval(self.config, json_loads=json.loads)
 
         if self._passphrase:
-            return serialization.load_pem_private_key(
+            # Load encrypted private key and cast to JWT-compatible type
+            # The JWT algorithms we support (RSA, ECDSA, EdDSA) use compatible key types
+            private_key = serialization.load_pem_private_key(
                 secret_key.encode(),
                 password=self._passphrase.eval(self.config, json_loads=json.loads).encode(),
                 backend=default_backend(),
             )
+            return cast(JwtKeyTypes, private_key)
         else:
             return (
                 base64.b64encode(secret_key.encode()).decode()

airbyte_cdk/sources/declarative/interpolation/macros.py

@@ -208,9 +208,12 @@ def camel_case_to_snake_case(value: str) -> str:
     return re.sub(r"(?<!^)(?=[A-Z])", "_", value).lower()
 
 
-def random_uuid() -> str:
+def generate_uuid() -> str:
     """
     Generates a UUID4
+
+    Usage:
+    `"{{ generate_uuid() }}"`
     """
     return str(uuid.uuid4())
 
@@ -228,5 +231,6 @@ def random_uuid() -> str:
     str_to_datetime,
     sanitize_url,
     camel_case_to_snake_case,
+    generate_uuid,
 ]
 macros = {f.__name__: f for f in _macros_list}

unit_tests/sources/declarative/interpolation/test_macros.py

@@ -3,6 +3,7 @@
 #
 
 import datetime
+import uuid
 
 import pytest
 
@@ -20,6 +21,7 @@
         ("test_format_datetime", "format_datetime", True),
         ("test_duration", "duration", True),
         ("test_camel_case_to_snake_case", "camel_case_to_snake_case", True),
+        ("test_generate_uuid", "generate_uuid", True),
         ("test_not_a_macro", "thisisnotavalidmacro", False),
     ],
 )
@@ -275,3 +277,26 @@ def test_sanitize_url(test_name, input_value, expected_output):
 )
 def test_camel_case_to_snake_case(value, expected_value):
     assert macros["camel_case_to_snake_case"](value) == expected_value
+
+
+def test_generate_uuid():
+    """Test uuid macro generates valid UUID4 strings."""
+    uuid_fn = macros["generate_uuid"]
+
+    # Test that uuid function returns a string
+    result = uuid_fn()
+    assert isinstance(result, str)
+
+    # Test that the result is a valid UUID format
+    # This will raise ValueError if not a valid UUID
+    parsed_uuid = uuid.UUID(result)
+
+    # Test that it's specifically a UUID4 (version 4)
+    assert parsed_uuid.version == 4
+
+    # Test that multiple calls return different UUIDs
+    result2 = uuid_fn()
+    assert result != result2
+
+    # Test that both results are valid UUIDs
+    uuid.UUID(result2)  # Will raise ValueError if invalid