Skip to content

Commit cb28bb6

Browse files
juskeeratanandurnfdogci.datadog-api-spec
authored
[DELA-208] Adding delegated token authentication in python client (#2860)
* changed template files + generate * rename file * rename files to match go client * fix aws tests * fix conftest * Restore docs/datadog_api_client.rst file * regen * print header * fix headers * fix config propogation * feedback updates * rest changes * gen * update * del * fix tests * static config * update config * updates tests * update imports * config * Update src/datadog_api_client/api_client.py Co-authored-by: Kevin L <96131879+urnfdog@users.noreply.github.com> * Update src/datadog_api_client/api_client.py Co-authored-by: Kevin L <96131879+urnfdog@users.noreply.github.com> * pre-commit fixes --------- Co-authored-by: Kevin L <96131879+urnfdog@users.noreply.github.com> Co-authored-by: ci.datadog-api-spec <packages@datadoghq.com>
1 parent dbc4581 commit cb28bb6

File tree

16 files changed

+1884
-25
lines changed

16 files changed

+1884
-25
lines changed

.generator/conftest.py

Lines changed: 17 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -82,9 +82,12 @@ def encode(self, obj):
8282
JINJA_ENV.filters["safe_snake_case"] = safe_snake_case
8383
JINJA_ENV.globals["format_data_with_schema"] = format_data_with_schema
8484
JINJA_ENV.globals["format_parameters"] = format_parameters
85+
JINJA_ENV.globals["package"] = "datadog_api_client"
8586

8687
PYTHON_EXAMPLE_J2 = JINJA_ENV.get_template("example.j2")
87-
88+
DATADOG_EXAMPLES_J2 = {
89+
"aws.py": JINJA_ENV.get_template("example_aws.j2")
90+
}
8891

8992
def pytest_bdd_after_scenario(request, feature, scenario):
9093
try:
@@ -137,6 +140,19 @@ def pytest_bdd_after_scenario(request, feature, scenario):
137140
with output.open("w") as f:
138141
f.write(data)
139142

143+
for file_name, template in DATADOG_EXAMPLES_J2.items():
144+
output = ROOT_PATH / "examples" / "datadog" / file_name
145+
output.parent.mkdir(parents=True, exist_ok=True)
146+
147+
data = template.render(
148+
context=context,
149+
version=version,
150+
scenario=scenario,
151+
operation_spec=operation_spec.spec,
152+
)
153+
with output.open("w") as f:
154+
f.write(data)
155+
140156

141157
def pytest_bdd_apply_tag(tag, function):
142158
"""Register tags as custom markers and skip test for '@skip' ones."""

.generator/src/generator/cli.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -72,6 +72,8 @@ def cli(specs, output):
7272
"exceptions.py": env.get_template("exceptions.j2"),
7373
"model_utils.py": env.get_template("model_utils.j2"),
7474
"rest.py": env.get_template("rest.j2"),
75+
"delegated_auth.py": env.get_template("delegated_auth.j2"),
76+
"aws.py": env.get_template("aws.j2"),
7577
}
7678

7779
top_package = output / PACKAGE_NAME

.generator/src/generator/templates/api_client.j2

Lines changed: 65 additions & 12 deletions
Original file line numberDiff line numberDiff line change
@@ -54,6 +54,17 @@ class ApiClient:
5454
self.default_headers["Accept-Encoding"] = "gzip"
5555
# Set default User-Agent.
5656
self.user_agent = user_agent()
57+
58+
# Initialize delegated token config if delegated auth is configured
59+
self._delegated_token_config = None
60+
if (self.configuration.delegated_auth_provider is not None and
61+
self.configuration.delegated_auth_org_uuid is not None):
62+
from {{ package }}.delegated_auth import DelegatedTokenConfig
63+
self._delegated_token_config = DelegatedTokenConfig(
64+
org_uuid=self.configuration.delegated_auth_org_uuid,
65+
provider="aws",
66+
provider_auth=self.configuration.delegated_auth_provider,
67+
)
5768

5869
def __enter__(self) -> Self:
5970
return self
@@ -454,6 +465,32 @@ class ApiClient:
454465
return "application/json"
455466
return content_types[0]
456467

468+
def use_delegated_token_auth(self, headers: Dict[str, Any]) -> None:
469+
"""Use delegated token authentication if configured.
470+
471+
:param headers: Header parameters dict to be updated.
472+
:raises: ApiValueError if delegated token authentication fails
473+
"""
474+
# Skip if no delegated token config
475+
if self._delegated_token_config is None:
476+
return
477+
478+
# Check if we need to get or refresh the token
479+
if (self.configuration._delegated_token_credentials is None or
480+
self.configuration._delegated_token_credentials.is_expired()):
481+
482+
# Get new token from provider, passing the API configuration
483+
try:
484+
self.configuration._delegated_token_credentials = self.configuration.delegated_auth_provider.authenticate(
485+
self._delegated_token_config, self.configuration
486+
)
487+
except Exception as e:
488+
raise ApiValueError(f"Failed to get delegated token: {str(e)}")
489+
490+
# Set the Authorization header with the delegated token
491+
token = self.configuration._delegated_token_credentials.delegated_token
492+
headers["Authorization"] = f"Bearer {token}"
493+
457494

458495
class ThreadedApiClient(ApiClient):
459496

@@ -824,18 +861,34 @@ class Endpoint:
824861
if not self.settings["auth"]:
825862
return
826863

827-
for auth in self.settings["auth"]:
828-
auth_setting = self.api_client.configuration.auth_settings().get(auth)
829-
if auth_setting:
830-
if auth_setting["in"] == "header":
831-
if auth_setting["type"] != "http-signature":
832-
if auth_setting["value"] is None:
833-
raise ApiValueError("Invalid authentication token for {}".format(auth_setting["key"]))
834-
headers[auth_setting["key"]] = auth_setting["value"]
835-
elif auth_setting["in"] == "query":
836-
queries.append((auth_setting["key"], auth_setting["value"]))
837-
else:
838-
raise ApiValueError("Authentication token must be in `query` or `header`")
864+
# check if endpoint uses appKeyAuth and if delegated token config is available
865+
has_app_key_auth = "appKeyAuth" in self.settings["auth"]
866+
867+
# Check if delegated auth is configured (using our actual attributes)
868+
has_delegated_auth = (
869+
hasattr(self.api_client.configuration, 'delegated_auth_provider') and
870+
self.api_client.configuration.delegated_auth_provider is not None and
871+
hasattr(self.api_client.configuration, 'delegated_auth_org_uuid') and
872+
self.api_client.configuration.delegated_auth_org_uuid is not None
873+
)
874+
875+
if has_app_key_auth and has_delegated_auth:
876+
# Use delegated token authentication
877+
self.api_client.use_delegated_token_auth(headers)
878+
else:
879+
# Use regular authentication
880+
for auth in self.settings["auth"]:
881+
auth_setting = self.api_client.configuration.auth_settings().get(auth)
882+
if auth_setting:
883+
if auth_setting["in"] == "header":
884+
if auth_setting["type"] != "http-signature":
885+
if auth_setting["value"] is None:
886+
raise ApiValueError("Invalid authentication token for {}".format(auth_setting["key"]))
887+
headers[auth_setting["key"]] = auth_setting["value"]
888+
elif auth_setting["in"] == "query":
889+
queries.append((auth_setting["key"], auth_setting["value"]))
890+
else:
891+
raise ApiValueError("Authentication token must be in `query` or `header`")
839892

840893

841894
def user_agent() -> str:

0 commit comments

Comments
 (0)