Migration of users for multiworkspace environment (#1500)

Author: jcchr

platform/services/initial_user/app/account_service_client.py

@@ -35,8 +35,9 @@ def handle_rpcerror(self, rpc_err: RpcError, email: str) -> None:
 
     def migrate_accounts(
         self, uh_connection: UsersHandlerConnection, postgresql_connection: PostgreSQLConnection, organization_id: str
-    ) -> None:
+    ) -> bool:
         users = uh_connection.get_accounts_for_migration()
+        number_of_migrated_users = 0
         for index, user in enumerate(users):
             is_org_admin = index == 0
             if user["mail"] is None:
@@ -88,9 +89,13 @@ def migrate_accounts(
                     organization_id=organization_id,
                     is_admin=is_org_admin,
                 )
+
+                number_of_migrated_users += 1
             except RpcError as rpc_err:
                 self.handle_rpcerror(rpc_err, user["mail"])
 
+        return number_of_migrated_users == len(users)
+
     @staticmethod
     def update_user_id_spicedb(
         postgresql_connection: PostgreSQLConnection, user_mail: str | None, user_uid: str

platform/services/initial_user/app/run_migration.py

@@ -42,10 +42,11 @@ def main() -> None:
     SpiceDB().link_organization_to_workspace_in_spicedb(workspace_id=workspace_id, organization_id=organization_id)
     psql_connection = PostgreSQLConnection()
     try:  # https://www.psycopg.org/docs/usage.html#with-statement
-        acc_svc.migrate_accounts(
+        first_migration: bool = acc_svc.migrate_accounts(
             uh_connection=uh_connection, postgresql_connection=psql_connection, organization_id=organization_id
         )
-        psql_connection.update_workspace_object_id(workspace_id=workspace_id)
+        if first_migration:
+            psql_connection.update_workspace_object_id(workspace_id=workspace_id)
     finally:
         psql_connection.close()