Skip to content

Commit dea2cc9

Browse files
lazebnyioctavia-squidington-iii
andauthored
feat(low-code): added json.loads to jwt authenticator (#301)
Co-authored-by: octavia-squidington-iii <[email protected]>
1 parent a6d55be commit dea2cc9

File tree

2 files changed

+31
-11
lines changed
  • airbyte_cdk/sources/declarative/auth
  • unit_tests/sources/declarative/auth

2 files changed

+31
-11
lines changed

airbyte_cdk/sources/declarative/auth/jwt.py

Lines changed: 17 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -3,6 +3,7 @@
33
#
44

55
import base64
6+
import json
67
from dataclasses import InitVar, dataclass
78
from datetime import datetime
89
from typing import Any, Mapping, Optional, Union
@@ -104,21 +105,21 @@ def __post_init__(self, parameters: Mapping[str, Any]) -> None:
104105
)
105106

106107
def _get_jwt_headers(self) -> dict[str, Any]:
107-
""" "
108+
"""
108109
Builds and returns the headers used when signing the JWT.
109110
"""
110-
headers = self._additional_jwt_headers.eval(self.config)
111+
headers = self._additional_jwt_headers.eval(self.config, json_loads=json.loads)
111112
if any(prop in headers for prop in ["kid", "alg", "typ", "cty"]):
112113
raise ValueError(
113114
"'kid', 'alg', 'typ', 'cty' are reserved headers and should not be set as part of 'additional_jwt_headers'"
114115
)
115116

116117
if self._kid:
117-
headers["kid"] = self._kid.eval(self.config)
118+
headers["kid"] = self._kid.eval(self.config, json_loads=json.loads)
118119
if self._typ:
119-
headers["typ"] = self._typ.eval(self.config)
120+
headers["typ"] = self._typ.eval(self.config, json_loads=json.loads)
120121
if self._cty:
121-
headers["cty"] = self._cty.eval(self.config)
122+
headers["cty"] = self._cty.eval(self.config, json_loads=json.loads)
122123
headers["alg"] = self._algorithm
123124
return headers
124125

@@ -130,18 +131,19 @@ def _get_jwt_payload(self) -> dict[str, Any]:
130131
exp = now + self._token_duration if isinstance(self._token_duration, int) else now
131132
nbf = now
132133

133-
payload = self._additional_jwt_payload.eval(self.config)
134+
payload = self._additional_jwt_payload.eval(self.config, json_loads=json.loads)
134135
if any(prop in payload for prop in ["iss", "sub", "aud", "iat", "exp", "nbf"]):
135136
raise ValueError(
136137
"'iss', 'sub', 'aud', 'iat', 'exp', 'nbf' are reserved properties and should not be set as part of 'additional_jwt_payload'"
137138
)
138139

139140
if self._iss:
140-
payload["iss"] = self._iss.eval(self.config)
141+
payload["iss"] = self._iss.eval(self.config, json_loads=json.loads)
141142
if self._sub:
142-
payload["sub"] = self._sub.eval(self.config)
143+
payload["sub"] = self._sub.eval(self.config, json_loads=json.loads)
143144
if self._aud:
144-
payload["aud"] = self._aud.eval(self.config)
145+
payload["aud"] = self._aud.eval(self.config, json_loads=json.loads)
146+
145147
payload["iat"] = now
146148
payload["exp"] = exp
147149
payload["nbf"] = nbf
@@ -151,7 +153,7 @@ def _get_secret_key(self) -> str:
151153
"""
152154
Returns the secret key used to sign the JWT.
153155
"""
154-
secret_key: str = self._secret_key.eval(self.config)
156+
secret_key: str = self._secret_key.eval(self.config, json_loads=json.loads)
155157
return (
156158
base64.b64encode(secret_key.encode()).decode()
157159
if self._base64_encode_secret_key
@@ -176,7 +178,11 @@ def _get_header_prefix(self) -> Union[str, None]:
176178
"""
177179
Returns the header prefix to be used when attaching the token to the request.
178180
"""
179-
return self._header_prefix.eval(self.config) if self._header_prefix else None
181+
return (
182+
self._header_prefix.eval(self.config, json_loads=json.loads)
183+
if self._header_prefix
184+
else None
185+
)
180186

181187
@property
182188
def auth_header(self) -> str:

unit_tests/sources/declarative/auth/test_jwt.py

Lines changed: 14 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -126,6 +126,20 @@ def test_get_secret_key(self, base64_encode_secret_key, secret_key, expected):
126126
)
127127
assert authenticator._get_secret_key() == expected
128128

129+
def test_get_secret_key_from_config(
130+
self,
131+
):
132+
authenticator = JwtAuthenticator(
133+
config={"secrets": '{"secret_key": "test"}'},
134+
parameters={},
135+
secret_key="{{ json_loads(config['secrets'])['secret_key'] }}",
136+
algorithm="test_algo",
137+
token_duration=1200,
138+
base64_encode_secret_key=False,
139+
)
140+
expected = "test"
141+
assert authenticator._get_secret_key() == expected
142+
129143
def test_get_signed_token(self):
130144
authenticator = JwtAuthenticator(
131145
config={},

0 commit comments

Comments
 (0)