Merge pull request #4958 from systeminit/jkeiser/workspace-loader-aut…

…h-api-fix

Fix workspace loading in workspace_delegations_population (billing lambda)

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

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
 

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)

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
                     """
                 ),
             )

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: