diff --git a/component/lambda/functions/si_auth_api.py b/component/lambda/functions/si_auth_api.py index 1515ce160f..bc95a976ce 100644 --- a/component/lambda/functions/si_auth_api.py +++ b/component/lambda/functions/si_auth_api.py @@ -1,5 +1,6 @@ +from typing import Literal, Optional, TypeVar, TypedDict, cast from pip._vendor import requests -from si_types import WorkspaceId, OwnerPk +from si_types import WorkspaceId, OwnerPk, IsoTimestamp, Ulid import urllib.parse import logging @@ -51,10 +52,40 @@ def post(self, path: str, json): # Function to query owner workspaces def owner_workspaces(self, workspace_id: WorkspaceId): - return self.get(f"/workspaces/{workspace_id}/ownerWorkspaces").json() + result = self.get(f"/workspaces/{workspace_id}/ownerWorkspaces").json() + return cast(WorkspaceOwnerWorkspaces, result) class HTTPError(Exception): def __init__(self, *args, json, **kwargs): super().__init__(*args, **kwargs) self.json = json +InstanceEnvType = Literal['LOCAL', 'PRIVATE', 'SI'] +Role = Literal['OWNER', 'APPROVER', 'EDITOR'] + +class WorkspaceOwnerWorkspaces(TypedDict): + workspaceId: WorkspaceId + workspaceOwnerId: OwnerPk + workspaces: list['OwnedWorkspace'] + +class Workspace(TypedDict): + id: WorkspaceId + instanceEnvType: InstanceEnvType + instanceUrl: str + displayName: str + description: Optional[str] + isDefault: bool + isFavourite: bool + creatorUserId: OwnerPk + creatorUser: 'WorkspaceCreatorUser' + token: Ulid + deletedAt: Optional[IsoTimestamp] + quarantinedAt: Optional[IsoTimestamp] + +class OwnedWorkspace(Workspace): + role: Role + invitedAt: Optional[IsoTimestamp] + +class WorkspaceCreatorUser(TypedDict): + firstName: str + lastName: str diff --git a/component/lambda/functions/si_lambda.py b/component/lambda/functions/si_lambda.py index f98851410d..a8c8483d15 100644 --- a/component/lambda/functions/si_lambda.py +++ b/component/lambda/functions/si_lambda.py @@ -30,6 +30,7 @@ class SiLambdaEnv(TypedDict): """ARN to an AWS secret containing the Lago API token""" LAGO_API_TOKEN_ARN: NotRequired[str] + """Auth API URL. Defaults to https://auth-api.systeminit.com""" AUTH_API_URL: NotRequired[str] BILLING_USER_EMAIL: NotRequired[str] BIlLING_USER_PASSWORD_ARN: NotRequired[str] @@ -44,6 +45,7 @@ def __init__(self, event: SiLambdaEnv, session = boto3.Session()): self.dry_run = event.get("SI_DRY_RUN", False) self._lago = None self._redshift = None + self._auth_api = None logging.getLogger().setLevel(self.getenv("SI_LOG_LEVEL", self.getenv("LOG_LEVEL", "INFO"))) @property @@ -78,12 +80,12 @@ def redshift(self): return self._redshift + @property def auth_api(self): """Get the Auth API client, configured from the lambda environment """ if self._auth_api is None: - auth_api_url = self.getenv("AUTH_API_URL") - if auth_api_url is None: - return None + auth_api_url = self.getenv("AUTH_API_URL", "https://auth-api.systeminit.com") + assert auth_api_url is not None, "AUTH_API_URL must be set" billing_user_email = self.getenv("BILLING_USER_EMAIL") assert billing_user_email is not None, "BILLING_USER_EMAIL must be set" @@ -94,7 +96,8 @@ def auth_api(self): billing_user_workspace_id = cast(Optional[WorkspaceId], self.getenv("BILLING_USER_WORKSPACE_ID")) assert billing_user_workspace_id is not None, "BILLING_USER_WORKSPACE_ID must be set" - self._auth_api = SiAuthApi.login(auth_api_url, billing_user_email, billing_user_password["BILLING_USER_PASSWORD"], billing_user_workspace_id) + print(billing_user_password) + self._auth_api = SiAuthApi.login(auth_api_url, billing_user_email, billing_user_password["BILLING_USER_PASWORD"], billing_user_workspace_id) return self._auth_api diff --git a/component/lambda/functions/si_types.py b/component/lambda/functions/si_types.py index e49d2831fc..e856c7e3a1 100644 --- a/component/lambda/functions/si_types.py +++ b/component/lambda/functions/si_types.py @@ -1,6 +1,7 @@ from typing import NewType -OwnerPk = NewType("OwnerPk", str) -WorkspaceId = NewType("WorkspaceId", str) +Ulid = NewType("Ulid", str) +OwnerPk = NewType("OwnerPk", Ulid) +WorkspaceId = NewType("WorkspaceId", Ulid) SqlTimestamp = NewType("SqlTimestamp", str) IsoTimestamp = NewType("IsoTimestamp", str) \ No newline at end of file diff --git a/component/lambda/functions/upload-billing-hours.py b/component/lambda/functions/upload-billing-hours.py index 7682919f83..1c1c2582b4 100644 --- a/component/lambda/functions/upload-billing-hours.py +++ b/component/lambda/functions/upload-billing-hours.py @@ -68,12 +68,12 @@ def run(self): self.redshift.query_raw( f""" SELECT external_subscription_id, hour_start, resource_count - FROM workspace_operations.owner_resource_hours_with_subscriptions - CROSS JOIN workspace_operations.workspace_update_events_summary - WHERE external_subscription_id IS NOT NULL - AND DATEADD(HOUR, {first_hour_start}, last_complete_hour_start) <= hour_start - AND hour_start < DATEADD(HOUR, {last_hour_end}, last_complete_hour_start) - ORDER BY hour_start DESC + FROM workspace_operations.owner_resource_hours_with_subscriptions + CROSS JOIN workspace_operations.workspace_update_events_summary + WHERE external_subscription_id IS NOT NULL + AND DATEADD(HOUR, {first_hour_start}, last_complete_hour_start) <= hour_start + AND hour_start < DATEADD(HOUR, {last_hour_end}, last_complete_hour_start) + ORDER BY hour_start DESC """ ), ) diff --git a/component/lambda/functions/workspace_delegations_population.py b/component/lambda/functions/workspace_delegations_population.py index 905250850e..8a9b2d5e61 100755 --- a/component/lambda/functions/workspace_delegations_population.py +++ b/component/lambda/functions/workspace_delegations_population.py @@ -124,8 +124,15 @@ def get_owner_lago_subscriptions(self, owner_pk: OwnerPk): def insert_missing_workspaces(self, current_timestamp: SqlTimestamp): missing_workspace_inserts = [ - [workspace_id, self.start_inserting_workspace(workspace_id, self.get_workspace_owner(workspace_id), current_timestamp)] - for [workspace_id] in cast(Iterable[tuple[WorkspaceId]], self.redshift.query(""" + [ + workspace_id, + self.start_inserting_workspace( + workspace_id, + self.auth_api.owner_workspaces(workspace_id)["workspaceOwnerId"], + current_timestamp + ) + ] + for [workspace_id] in cast(Iterable[tuple[WorkspaceId]], self.redshift.query_raw(""" SELECT DISTINCT workspace_id FROM workspace_update_events.workspace_update_events LEFT OUTER JOIN workspace_operations.workspace_owners USING (workspace_id) @@ -162,16 +169,12 @@ def start_inserting_workspace(self, workspace_id: WorkspaceId, workspace_owner_i logging.info(f"Inserting into workspace_owner_subscriptions: {sql}") return self.redshift.start_executing(sql) - def get_workspace_owner(self, workspace_id: WorkspaceId): - owner_workspace_data = self.auth_api.owner_workspaces(workspace_id) - return cast(OwnerPk, owner_workspace_data.get("workspaceOwnerId")) - def run(self): # Get the current timestamp for record insertion current_timestamp = cast(SqlTimestamp, time.strftime('%Y-%m-%d %H:%M:%S')) - updated_subscriptions = self.update_subscriptions(current_timestamp) inserted_workspaces = self.insert_missing_workspaces(current_timestamp) + updated_subscriptions = self.update_subscriptions(current_timestamp) return { 'statusCode': 200, @@ -195,6 +198,8 @@ class LatestOwnerSubscription(TypedDict): @overload def convert_iso_to_datetime(iso_str: IsoTimestamp) -> SqlTimestamp: ... @overload +def convert_iso_to_datetime(iso_str: None) -> None: ... +@overload def convert_iso_to_datetime(iso_str: Optional[IsoTimestamp]) -> Optional[SqlTimestamp]: ... def convert_iso_to_datetime(iso_str: Optional[IsoTimestamp]): if iso_str is None: @@ -204,6 +209,8 @@ def convert_iso_to_datetime(iso_str: Optional[IsoTimestamp]): @overload def iso_to_days(iso_str: IsoTimestamp) -> str: ... @overload +def iso_to_days(iso_str: None) -> None: ... +@overload def iso_to_days(iso_str: Optional[IsoTimestamp]) -> Optional[str]: ... def iso_to_days(iso_str: Optional[IsoTimestamp]): if iso_str is None: