Skip to content
Merged
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
34 changes: 30 additions & 4 deletions examples/agent_pool.py
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,14 @@
This example demonstrates:
1. Agent Pool CRUD operations (Create, Read, Update, Delete)
2. Agent token creation and management
3. Using the organization SDK client
3. Workspace assignment using assign_to_workspaces and remove_from_workspaces
4. Proper error handling

Make sure to set the following environment variables:
- TFE_TOKEN: Your Terraform Cloud/Enterprise API token
- TFE_ADDRESS: Your Terraform Cloud/Enterprise URL (optional, defaults to https://app.terraform.io)
- TFE_ORG: Your organization name
- TFE_WORKSPACE_ID: A workspace ID for testing workspace assignment (optional)

Usage:
export TFE_TOKEN="your-token-here"
Expand All @@ -24,8 +25,10 @@
from pytfe.errors import NotFound
from pytfe.models import (
AgentPoolAllowedWorkspacePolicy,
AgentPoolAssignToWorkspacesOptions,
AgentPoolCreateOptions,
AgentPoolListOptions,
AgentPoolRemoveFromWorkspacesOptions,
AgentPoolUpdateOptions,
AgentTokenCreateOptions,
)
Expand All @@ -37,6 +40,9 @@ def main():
token = os.environ.get("TFE_TOKEN")
org = os.environ.get("TFE_ORG")
address = os.environ.get("TFE_ADDRESS", "https://app.terraform.io")
workspace_id = os.environ.get(
"TFE_WORKSPACE_ID"
) # optional, for workspace assignment

if not token:
print("TFE_TOKEN environment variable is required")
Expand Down Expand Up @@ -96,7 +102,27 @@ def main():
updated_pool = client.agent_pools.update(new_pool.id, update_options)
print(f"Updated agent pool name to: {updated_pool.name}")

# Example 5: Create an agent token
# Example 5: Workspace assignment
# assign_to_workspaces sends PATCH /agent-pools/:id with relationships.allowed-workspaces
# remove_from_workspaces sends PATCH /agent-pools/:id with relationships.excluded-workspaces
if workspace_id:
print("\n Assigning workspace to agent pool...")
updated_pool = client.agent_pools.assign_to_workspaces(
new_pool.id,
AgentPoolAssignToWorkspacesOptions(workspace_ids=[workspace_id]),
)
print(f" Assigned workspace {workspace_id} to pool {updated_pool.name}")

print("\n Removing workspace from agent pool...")
updated_pool = client.agent_pools.remove_from_workspaces(
new_pool.id,
AgentPoolRemoveFromWorkspacesOptions(workspace_ids=[workspace_id]),
)
print(f" Removed workspace {workspace_id} from pool {updated_pool.name}")
else:
print("\n Skipping workspace assignment (set TFE_WORKSPACE_ID to test)")

# Example 6: Create an agent token
print("\n Creating agent token...")
token_options = AgentTokenCreateOptions(
description="SDK example token" # Optional description
Expand All @@ -107,7 +133,7 @@ def main():
if agent_token.token:
print(f" Token (first 10 chars): {agent_token.token[:10]}...")

# Example 6: List agent tokens
# Example 7: List agent tokens
print("\n Listing agent tokens...")
tokens = client.agent_tokens.list(new_pool.id)

Expand All @@ -117,7 +143,7 @@ def main():
for token in token_list:
print(f" - {token.description or 'No description'} (ID: {token.id})")

# Example 7: Clean up - delete the token and pool
# Example 8: Clean up - delete the token and pool
print("\n Cleaning up...")
client.agent_tokens.delete(agent_token.id)
print("Deleted agent token")
Expand Down
8 changes: 8 additions & 0 deletions src/pytfe/models/agent.py
Original file line number Diff line number Diff line change
Expand Up @@ -82,6 +82,10 @@ class AgentPoolCreateOptions(BaseModel):
organization_scoped: bool | None = None
# Optional: Allowed workspace policy
allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None
# Optional: IDs of workspaces allowed to use this pool (sent as relationships.allowed-workspaces)
allowed_workspace_ids: list[str] = Field(default_factory=list)
# Optional: IDs of workspaces excluded from this pool (sent as relationships.excluded-workspaces)
excluded_workspace_ids: list[str] = Field(default_factory=list)


class AgentPoolUpdateOptions(BaseModel):
Expand All @@ -93,6 +97,10 @@ class AgentPoolUpdateOptions(BaseModel):
organization_scoped: bool | None = None
# Optional: Allowed workspace policy
allowed_workspace_policy: AgentPoolAllowedWorkspacePolicy | None = None
# Optional: Full replacement list of workspace IDs allowed to use this pool
allowed_workspace_ids: list[str] = Field(default_factory=list)
# Optional: Full replacement list of workspace IDs excluded from this pool
excluded_workspace_ids: list[str] = Field(default_factory=list)


class AgentPoolReadOptions(BaseModel):
Expand Down
161 changes: 140 additions & 21 deletions src/pytfe/resources/agent_pools.py
Original file line number Diff line number Diff line change
Expand Up @@ -203,7 +203,27 @@ def create(self, organization: str, options: AgentPoolCreateOptions) -> AgentPoo
options.allowed_workspace_policy.value
)

payload = {"data": {"type": "agent-pools", "attributes": attributes}}
relationships: dict[str, Any] = {}
if options.allowed_workspace_ids:
relationships["allowed-workspaces"] = {
"data": [
{"type": "workspaces", "id": ws_id}
for ws_id in options.allowed_workspace_ids
]
}
if options.excluded_workspace_ids:
relationships["excluded-workspaces"] = {
"data": [
{"type": "workspaces", "id": ws_id}
for ws_id in options.excluded_workspace_ids
]
}

payload: dict[str, Any] = {
"data": {"type": "agent-pools", "attributes": attributes}
}
if relationships:
payload["data"]["relationships"] = relationships

response = self.t.request("POST", path, json_body=payload)
data = response.json()["data"]
Expand Down Expand Up @@ -320,13 +340,31 @@ def update(self, agent_pool_id: str, options: AgentPoolUpdateOptions) -> AgentPo
options.allowed_workspace_policy.value
)

payload = {
relationships: dict[str, Any] = {}
if options.allowed_workspace_ids:
relationships["allowed-workspaces"] = {
"data": [
{"type": "workspaces", "id": ws_id}
for ws_id in options.allowed_workspace_ids
]
}
if options.excluded_workspace_ids:
relationships["excluded-workspaces"] = {
"data": [
{"type": "workspaces", "id": ws_id}
for ws_id in options.excluded_workspace_ids
]
}

payload: dict[str, Any] = {
"data": {
"type": "agent-pools",
"id": agent_pool_id,
"attributes": attributes,
}
}
if relationships:
payload["data"]["relationships"] = relationships

response = self.t.request("PATCH", path, json_body=payload)
data = response.json()["data"]
Expand Down Expand Up @@ -371,13 +409,20 @@ def delete(self, agent_pool_id: str) -> None:

def assign_to_workspaces(
self, agent_pool_id: str, options: AgentPoolAssignToWorkspacesOptions
) -> None:
"""Assign an agent pool to workspaces.
) -> AgentPool:
"""Assign an agent pool to workspaces by updating the allowed-workspaces
relationship via PATCH /agent-pools/:id.

The provided workspace IDs become the new complete list of allowed
workspaces for this pool (full replacement, not append).

Args:
agent_pool_id: Agent pool ID
options: Assignment options containing workspace IDs

Returns:
Updated AgentPool object

Raises:
ValueError: If parameters are invalid
TFEError: If API request fails
Expand All @@ -388,26 +433,67 @@ def assign_to_workspaces(
if not options.workspace_ids:
raise ValueError("At least one workspace ID is required")

path = f"/api/v2/agent-pools/{agent_pool_id}/relationships/workspaces"

# Create data payload with workspace references
workspace_data = []
for workspace_id in options.workspace_ids:
if not valid_string_id(workspace_id):
raise ValueError(f"Invalid workspace ID: {workspace_id}")
workspace_data.append({"type": "workspaces", "id": workspace_id})

payload = {"data": workspace_data}
self.t.request("POST", path, json_body=payload)
path = f"/api/v2/agent-pools/{agent_pool_id}"
payload: dict[str, Any] = {
"data": {
"type": "agent-pools",
"id": agent_pool_id,
"attributes": {},
"relationships": {
"allowed-workspaces": {
"data": [
{"type": "workspaces", "id": ws_id}
for ws_id in options.workspace_ids
]
}
},
}
}
response = self.t.request("PATCH", path, json_body=payload)
data = response.json()["data"]

# Extract agent pool data from response
attr = data.get("attributes", {}) or {}
agent_pool_data = {
"id": _safe_str(data.get("id")),
"name": _safe_str(attr.get("name")),
"created_at": attr.get("created-at"),
"organization_scoped": attr.get("organization-scoped"),
"allowed_workspace_policy": attr.get("allowed-workspace-policy"),
"agent_count": attr.get("agent-count", 0),
}

return AgentPool(
id=_safe_str(agent_pool_data["id"]) or "",
name=_safe_str(agent_pool_data["name"]),
created_at=cast(Any, agent_pool_data["created_at"]),
organization_scoped=_safe_bool(agent_pool_data["organization_scoped"]),
allowed_workspace_policy=_safe_workspace_policy(
agent_pool_data["allowed_workspace_policy"]
),
agent_count=_safe_int(agent_pool_data["agent_count"]),
)

def remove_from_workspaces(
self, agent_pool_id: str, options: AgentPoolRemoveFromWorkspacesOptions
) -> None:
"""Remove an agent pool from workspaces.
) -> AgentPool:
"""Exclude workspaces from an agent pool by updating the excluded-workspaces
relationship via PATCH /agent-pools/:id.

Use this for organization-scoped pools where most workspaces are allowed
but you want to block specific ones. The provided list becomes the new
complete excluded-workspaces list (full replacement, not append).

Args:
agent_pool_id: Agent pool ID
options: Removal options containing workspace IDs
options: Removal options containing workspace IDs to exclude

Returns:
Updated AgentPool object

Raises:
ValueError: If parameters are invalid
Expand All @@ -419,14 +505,47 @@ def remove_from_workspaces(
if not options.workspace_ids:
raise ValueError("At least one workspace ID is required")

path = f"/api/v2/agent-pools/{agent_pool_id}/relationships/workspaces"

# Create data payload with workspace references
workspace_data = []
for workspace_id in options.workspace_ids:
if not valid_string_id(workspace_id):
raise ValueError(f"Invalid workspace ID: {workspace_id}")
workspace_data.append({"type": "workspaces", "id": workspace_id})

payload = {"data": workspace_data}
self.t.request("DELETE", path, json_body=payload)
path = f"/api/v2/agent-pools/{agent_pool_id}"
payload: dict[str, Any] = {
"data": {
"type": "agent-pools",
"id": agent_pool_id,
"attributes": {},
"relationships": {
"excluded-workspaces": {
"data": [
{"type": "workspaces", "id": ws_id}
for ws_id in options.workspace_ids
]
}
},
}
}
response = self.t.request("PATCH", path, json_body=payload)
data = response.json()["data"]

# Extract agent pool data from response
attr = data.get("attributes", {}) or {}
agent_pool_data = {
"id": _safe_str(data.get("id")),
"name": _safe_str(attr.get("name")),
"created_at": attr.get("created-at"),
"organization_scoped": attr.get("organization-scoped"),
"allowed_workspace_policy": attr.get("allowed-workspace-policy"),
"agent_count": attr.get("agent-count", 0),
}

return AgentPool(
id=_safe_str(agent_pool_data["id"]) or "",
name=_safe_str(agent_pool_data["name"]),
created_at=cast(Any, agent_pool_data["created_at"]),
organization_scoped=_safe_bool(agent_pool_data["organization_scoped"]),
allowed_workspace_policy=_safe_workspace_policy(
agent_pool_data["allowed_workspace_policy"]
),
agent_count=_safe_int(agent_pool_data["agent_count"]),
)
Loading
Loading