Skip to content

Commit

Permalink
Add scape_org_role_arn to account information (#635)
Browse files Browse the repository at this point in the history
  • Loading branch information
meln1k authored Oct 11, 2024
1 parent 38dd9f7 commit e5f33ac
Show file tree
Hide file tree
Showing 5 changed files with 26 additions and 4 deletions.
2 changes: 1 addition & 1 deletion .devcontainer/devcontainer.json
Original file line number Diff line number Diff line change
Expand Up @@ -8,7 +8,7 @@
// Features to add to the dev container. More info: https://containers.dev/features.
"features": {
"ghcr.io/devcontainers-contrib/features/poetry:2": {},
"ghcr.io/devcontainers/features/node:1": {},
"ghcr.io/devcontainers/features/node:1.5": {}
},
// Use 'forwardPorts' to make a list of ports inside the container available locally.
// This can be used to network with other containers or the host.
Expand Down
1 change: 1 addition & 0 deletions fixbackend/collect/collect_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -54,6 +54,7 @@ class AwsAccountInformation(AccountInformation):
aws_account_id: CloudAccountId
aws_account_name: Optional[str]
aws_role_arn: AwsARN
scrape_org_role_arn: Optional[AwsARN]
external_id: ExternalId


Expand Down
16 changes: 14 additions & 2 deletions fixbackend/dispatcher/dispatcher_service.py
Original file line number Diff line number Diff line change
Expand Up @@ -488,6 +488,7 @@ async def trigger_collect(
reason: str,
defer_by: Optional[timedelta] = None,
retry_failed_for: Optional[timedelta] = None,
privileged_account_id: Optional[CloudAccountId] = None,
**kwargs: Any,
) -> None:
set_cloud_account_id(account.account_id)
Expand All @@ -507,6 +508,11 @@ async def account_information() -> Optional[AccountInformation]:
aws_account_id=account.account_id,
aws_account_name=account.final_name(),
aws_role_arn=AwsARN(f"arn:aws:iam::{account.account_id}:role/{role_name}"),
scrape_org_role_arn=(
AwsARN(f"arn:aws:iam::{privileged_account_id}:role/{role_name}")
if privileged_account_id
else None
),
external_id=external_id,
)
case GcpCloudAccess(service_account_key_id):
Expand Down Expand Up @@ -571,6 +577,7 @@ async def schedule_next_runs(self) -> None:
accounts = healthy_accounts + degraded_accounts
product_tier = await self.workspace_repository.get_product_tier(workspace_id)
log.info(f"scheduling next run for workspace {workspace_id}, {len(accounts)} accounts")
priveleged_account_id = next((acc.account_id for acc in accounts if acc.privileged), None)
for account in accounts:
reason = "regular_collect"

Expand All @@ -579,9 +586,14 @@ async def schedule_next_runs(self) -> None:

if account.cloud == CloudNames.Azure and not azure_graph_scheduled:
azure_graph_scheduled = True
await self.trigger_collect(account, reason=reason, collect_microsoft_graph=True)
await self.trigger_collect(
account,
reason=reason,
collect_microsoft_graph=True,
privileged_account_id=priveleged_account_id,
)
else:
await self.trigger_collect(account, reason=reason)
await self.trigger_collect(account, reason=reason, privileged_account_id=priveleged_account_id)

for account in degraded_accounts:
await self.cloud_account_repo.update(
Expand Down
8 changes: 7 additions & 1 deletion tests/fixbackend/collect/collect_queue_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -49,7 +49,11 @@ async def test_redis_collect_queue(
assert set(await arq_redis.keys()) == set()
# enqueue new job
aws_account = AwsAccountInformation(
CloudAccountId("123"), CloudAccountName("test"), AwsARN("arn"), ExternalId(uuid4())
CloudAccountId("123"),
CloudAccountName("test"),
AwsARN("arn"),
AwsARN("scarpe_arn"),
ExternalId(uuid4()),
)
await collect_queue.enqueue(graph_db_access, aws_account, job_id="test")
assert set(await arq_redis.keys()) == {b"arq:queue", b"arq:job:test"}
Expand Down Expand Up @@ -80,13 +84,15 @@ def test_aws_account_info_json() -> None:
aws_account_id=CloudAccountId("123456789012"),
aws_account_name=CloudAccountName("test"),
aws_role_arn=AwsARN("arn:aws:iam::123456789012:role/test"),
scrape_org_role_arn=AwsARN("arn:aws:iam::123456789012:role/scrape-test"),
external_id=external_id,
)
assert aws_account_info.to_json() == {
"kind": "aws_account_information",
"aws_account_id": "123456789012",
"aws_account_name": "test",
"aws_role_arn": "arn:aws:iam::123456789012:role/test",
"scrape_org_role_arn": "arn:aws:iam::123456789012:role/scrape-test",
"external_id": str(external_id),
}

Expand Down
3 changes: 3 additions & 0 deletions tests/fixbackend/dispatcher/dispatcher_service_test.py
Original file line number Diff line number Diff line change
Expand Up @@ -241,6 +241,7 @@ async def jobs_mapping_hash_len() -> int:
aws_account_id=aws_account_id,
aws_account_name=CloudAccountName("test"),
aws_role_arn=AwsARN("arn"),
scrape_org_role_arn=AwsARN("scrape_arn"),
external_id=external_id,
),
job_id,
Expand Down Expand Up @@ -391,6 +392,7 @@ async def jobs_mapping_hash_len() -> int:
aws_account_id=aws_account_id,
aws_account_name=CloudAccountName("test"),
aws_role_arn=AwsARN("arn"),
scrape_org_role_arn=AwsARN("scrape_arn"),
external_id=external_id,
),
job_id,
Expand Down Expand Up @@ -478,6 +480,7 @@ async def jobs_mapping_hash_len() -> int:
aws_account_id=aws_account_id,
aws_account_name=CloudAccountName("test"),
aws_role_arn=AwsARN("arn"),
scrape_org_role_arn=AwsARN("scrape_arn"),
external_id=external_id,
),
job_id,
Expand Down

0 comments on commit e5f33ac

Please sign in to comment.