parallelize across orgs

brian_dashboard_manager/grafana/organization.py

@@ -79,61 +79,68 @@ def create_organization(request: AdminRequest, name: str) -> Union[Dict, None]:
 
 
 @contextlib.contextmanager
-def get_or_create_service_account(request: AdminRequest, org_id):
+def get_or_create_service_accounts(request: AdminRequest, org_ids: list[int]):
     """
-    Gets a service account for the given organization, or creates one if it does not exist.
+    Gets or creates service accounts for multiple organizations and manages their lifecycle.
 
     :param request: AdminRequest object
-    :param org_id: organization ID
-    :param name: service account name
-    :return: service account definition
+    :param org_ids: List of organization IDs to create service accounts for
+    :return: Dictionary mapping org_id to service account definitions
     """
-    switch_active_organization(request, org_id)
-
-    # get provision service account, if it exists
-
-    service_account = None
-    try:
-        service_accounts = request.get(
-            "api/serviceaccounts/search?query=provision"
-        ).json()
+    service_accounts = {}
+    tokens = {}
 
-        if service_accounts and service_accounts.get("totalCount") > 0:
-            service_account = service_accounts.get("serviceAccounts")[0]
+    # Create service accounts for each org
+    for org_id in org_ids:
+        switch_active_organization(request, org_id)
 
-    except HTTPError as e:
-        if e.response.status_code != 404:
-            raise BrianHTTPError(
-                e.response,
-                client_message=f"Error when listing service accounts for organzation #{org_id}",
-            )
+        # Check for existing service account
+        try:
+            accounts = request.get("api/serviceaccounts/search?query=provision").json()
+            if accounts and accounts.get("totalCount") > 0:
+                _account = accounts.get("serviceAccounts", [])[0]
+                service_accounts[org_id] = _account
+                token = create_service_account_token(request, _account["id"])
+                tokens[org_id] = token
+                continue
+        except HTTPError as e:
+            if e.response.status_code != 404:
+                raise BrianHTTPError(
+                    e.response,
+                    client_message=f"Error when listing service accounts for organization #{org_id}"
+                )
 
-    if not service_account:
+        # Create new service account if none exists
         try:
-            service_account = request.post(
+            account = request.post(
                 "api/serviceaccounts",
                 json={
                     "name": "provision",
                     "role": "Admin",
                     "isDisabled": False,
-                },
+                }
             ).json()
+            service_accounts[org_id] = account
+            token = create_service_account_token(request, account["id"])
+            tokens[org_id] = token
+            logger.info(f"Created provision service account for organization #{org_id}")
         except HTTPError as e:
             raise BrianHTTPError(
                 e.response,
-                client_message=f"Error when creating service account for organization #{org_id}",
+                client_message=f"Error when creating service account for organization #{org_id}"
             )
-        logger.info(f"Created provision service account for organization #{org_id}")
+
     try:
-        yield service_account
+        yield tokens
     finally:
-        try:
-            request.delete(f"api/serviceaccounts/{service_account['id']}")
-        except HTTPError as e:
-            raise BrianHTTPError(
-                e.response,
-                client_message=f"Error when deleting service account for organization #{org_id}",
-            )
+        # Clean up all service accounts
+        for org_id, account in service_accounts.items():
+            try:
+                switch_active_organization(request, org_id)
+                request.delete(f"api/serviceaccounts/{account['id']}")
+                logger.info(f"Deleted service account for organization #{org_id}")
+            except HTTPError as e:
+                logger.error(f"Failed to delete service account for organization #{org_id}: {e}")
 
 
 def create_service_account_token(request: AdminRequest, service_account_id: int):

brian_dashboard_manager/provision/__init__.py

@@ -65,8 +65,7 @@ from brian_dashboard_manager.grafana.folder import (
 )
 from brian_dashboard_manager.grafana.organization import (
     create_organization,
-    create_service_account_token,
-    get_or_create_service_account,
+    get_or_create_service_accounts,
     get_organizations,
 )
 from brian_dashboard_manager.grafana.organization import (
@@ -355,20 +354,16 @@ def provision_organization(
     organization: Organization,
     folders: Sequence[DashboardFolder],
     data: BrianData,
+    token: str,
 ):
-
-    request = AdminRequest(**config)
     logger.info(f"--- Provisioning org {organization.name} (ID #{organization.id}) ---")
-
-    with get_or_create_service_account(request, organization.id) as account:
-        token = create_service_account_token(request, account["id"])
-        organization.token_params = (config["hostname"], token)
-        datasource = provision_datasource(config, organization.request)
-        data = dataclasses.replace(
-            data, datasource_name=datasource.get("name", "PollerInfluxDB")
-        )
-        provision_all_dashboards(folders, data, organization, config=config)
-        set_home_dashboard(organization=organization)
+    organization.token_params = (config["hostname"], token)
+    datasource = provision_datasource(config, organization.request)
+    data = dataclasses.replace(
+        data, datasource_name=datasource.get("name", "PollerInfluxDB")
+    )
+    provision_all_dashboards(folders, data, organization, config=config)
+    set_home_dashboard(organization=organization)
 
 
 def provision_datasource(config, request):
@@ -551,8 +546,12 @@ def provision(config):
     to_populate = provision_organizations(
         config.get("organizations", DEFAULT_ORGANIZATIONS), request
     )
-    with register_errors():
-        for org in to_populate:
-            provision_organization(config, org, FOLDERS, data)
+    with register_errors(), get_or_create_service_accounts(request, [org.id for org in to_populate]) as tokens:
+        with ThreadPoolExecutor() as executor:
+            futs = []
+            for org in to_populate:
+                futs.append(executor.submit(provision_organization, config, org, FOLDERS, data, tokens[org.id]))
+            wait(futs)
+
         logger.info(f"Time to complete: {time.time() - start}")
         return HAS_ERRORS

test/grafana/test_organization.py

 import responses
 from brian_dashboard_manager.grafana.organization import (
     create_organization,
-    get_or_create_service_account,
+    get_or_create_service_accounts,
     get_organizations,
 )
 
@@ -34,7 +34,7 @@ def test_create_organization(mock_grafana):
 def test_creates_and_deletes_service_accounts(mock_grafana):
     mock_grafana.create_organization({"id": 1, "name": "Testorg1"})
     assert not mock_grafana.service_accounts
-    with get_or_create_service_account(mock_grafana.request, 1):
+    with get_or_create_service_accounts(mock_grafana.request, [1]):
         assert mock_grafana.service_accounts
     assert not mock_grafana.service_accounts
 
@@ -50,6 +50,6 @@ def test_doesnt_create_existing_service_account_but_does_delete_it(mock_grafana)
         }
     )
     assert mock_grafana.service_accounts
-    with get_or_create_service_account(mock_grafana.request, 1):
+    with get_or_create_service_accounts(mock_grafana.request, [1]):
         pass
     assert not mock_grafana.service_accounts