Skip to content
Open
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
18 changes: 17 additions & 1 deletion .generator/conftest.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,9 +82,12 @@ def encode(self, obj):
JINJA_ENV.filters["safe_snake_case"] = safe_snake_case
JINJA_ENV.globals["format_data_with_schema"] = format_data_with_schema
JINJA_ENV.globals["format_parameters"] = format_parameters
JINJA_ENV.globals["package"] = "datadog_api_client"

PYTHON_EXAMPLE_J2 = JINJA_ENV.get_template("example.j2")

DATADOG_EXAMPLES_J2 = {
"aws.py": JINJA_ENV.get_template("example_aws.j2")
}

def pytest_bdd_after_scenario(request, feature, scenario):
try:
Expand Down Expand Up @@ -137,6 +140,19 @@ def pytest_bdd_after_scenario(request, feature, scenario):
with output.open("w") as f:
f.write(data)

for file_name, template in DATADOG_EXAMPLES_J2.items():
output = ROOT_PATH / "examples" / "datadog" / file_name
output.parent.mkdir(parents=True, exist_ok=True)

data = template.render(
context=context,
version=version,
scenario=scenario,
operation_spec=operation_spec.spec,
)
with output.open("w") as f:
f.write(data)


def pytest_bdd_apply_tag(tag, function):
"""Register tags as custom markers and skip test for '@skip' ones."""
Expand Down
2 changes: 2 additions & 0 deletions .generator/src/generator/cli.py
Original file line number Diff line number Diff line change
Expand Up @@ -72,6 +72,8 @@ def cli(specs, output):
"exceptions.py": env.get_template("exceptions.j2"),
"model_utils.py": env.get_template("model_utils.j2"),
"rest.py": env.get_template("rest.j2"),
"delegated_auth.py": env.get_template("delegated_auth.j2"),
"aws.py": env.get_template("aws.j2"),
}

top_package = output / PACKAGE_NAME
Expand Down
77 changes: 65 additions & 12 deletions .generator/src/generator/templates/api_client.j2
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,17 @@ class ApiClient:
self.default_headers["Accept-Encoding"] = "gzip"
# Set default User-Agent.
self.user_agent = user_agent()

# Initialize delegated token config if delegated auth is configured
self._delegated_token_config = None
if (self.configuration.delegated_auth_provider is not None and
self.configuration.delegated_auth_org_uuid is not None):
from {{ package }}.delegated_auth import DelegatedTokenConfig
self._delegated_token_config = DelegatedTokenConfig(
org_uuid=self.configuration.delegated_auth_org_uuid,
provider="aws",
provider_auth=self.configuration.delegated_auth_provider,
)

def __enter__(self) -> Self:
return self
Expand Down Expand Up @@ -454,6 +465,32 @@ class ApiClient:
return "application/json"
return content_types[0]

def use_delegated_token_auth(self, headers: Dict[str, Any]) -> None:
"""Use delegated token authentication if configured.

:param headers: Header parameters dict to be updated.
:raises: ApiValueError if delegated token authentication fails
"""
# Skip if no delegated token config
if self._delegated_token_config is None:
return

# Check if we need to get or refresh the token
if (self.configuration._delegated_token_credentials is None or
self.configuration._delegated_token_credentials.is_expired()):

# Get new token from provider, passing the API configuration
try:
self.configuration._delegated_token_credentials = self.configuration.delegated_auth_provider.authenticate(
self._delegated_token_config, self.configuration
)
except Exception as e:
raise ApiValueError(f"Failed to get delegated token: {str(e)}")

# Set the Authorization header with the delegated token
token = self.configuration._delegated_token_credentials.delegated_token
headers["Authorization"] = f"Bearer {token}"


class ThreadedApiClient(ApiClient):

Expand Down Expand Up @@ -824,18 +861,34 @@ class Endpoint:
if not self.settings["auth"]:
return

for auth in self.settings["auth"]:
auth_setting = self.api_client.configuration.auth_settings().get(auth)
if auth_setting:
if auth_setting["in"] == "header":
if auth_setting["type"] != "http-signature":
if auth_setting["value"] is None:
raise ApiValueError("Invalid authentication token for {}".format(auth_setting["key"]))
headers[auth_setting["key"]] = auth_setting["value"]
elif auth_setting["in"] == "query":
queries.append((auth_setting["key"], auth_setting["value"]))
else:
raise ApiValueError("Authentication token must be in `query` or `header`")
# check if endpoint uses appKeyAuth and if delegated token config is available
has_app_key_auth = "appKeyAuth" in self.settings["auth"]

# Check if delegated auth is configured (using our actual attributes)
has_delegated_auth = (
hasattr(self.api_client.configuration, 'delegated_auth_provider') and
self.api_client.configuration.delegated_auth_provider is not None and
hasattr(self.api_client.configuration, 'delegated_auth_org_uuid') and
self.api_client.configuration.delegated_auth_org_uuid is not None
)

if has_app_key_auth and has_delegated_auth:
# Use delegated token authentication
self.api_client.use_delegated_token_auth(headers)
else:
# Use regular authentication
for auth in self.settings["auth"]:
auth_setting = self.api_client.configuration.auth_settings().get(auth)
if auth_setting:
if auth_setting["in"] == "header":
if auth_setting["type"] != "http-signature":
if auth_setting["value"] is None:
raise ApiValueError("Invalid authentication token for {}".format(auth_setting["key"]))
headers[auth_setting["key"]] = auth_setting["value"]
elif auth_setting["in"] == "query":
queries.append((auth_setting["key"], auth_setting["value"]))
else:
raise ApiValueError("Authentication token must be in `query` or `header`")


def user_agent() -> str:
Expand Down
Loading