[CDF-23377] 😬QuickStart AuthVerify (#1232)

* refactor: moved out init build deploy

* refactor: setup shell for tmp group

* feat; modify auth verify

* tmp: wheel

* refactor: fix

* tmp: update wheel

* fix: bug

* tmp: update wheel

* fix; try some sleep

* tmp: update wheel

* fix: try reinitializing

* tmp: wheel

* refactor; switch to accept client ID and credentails

* tmp; wheel\

* Ãtmp: remove wheel

cognite_toolkit/_cdf_tk/commands/auth.py

@@ -16,6 +16,7 @@
 import time
 import warnings
 from collections import defaultdict
+from dataclasses import dataclass
 from time import sleep
 
 import questionary
@@ -34,7 +35,11 @@
 
 from cognite_toolkit._cdf_tk import loaders
 from cognite_toolkit._cdf_tk.client import ToolkitClient
-from cognite_toolkit._cdf_tk.constants import HINT_LEAD_TEXT, TOOLKIT_SERVICE_PRINCIPAL_GROUP_NAME
+from cognite_toolkit._cdf_tk.constants import (
+    HINT_LEAD_TEXT,
+    TOOLKIT_DEMO_GROUP_NAME,
+    TOOLKIT_SERVICE_PRINCIPAL_GROUP_NAME,
+)
 from cognite_toolkit._cdf_tk.exceptions import (
     AuthenticationError,
     AuthorizationError,
@@ -52,6 +57,12 @@
 from ._base import ToolkitCommand
 
 
+@dataclass
+class VerifyAuthResult:
+    toolkit_group_id: int | None = None
+    function_status: str | None = None
+
+
 class AuthCommand(ToolkitCommand):
     def init(self, no_verify: bool = False, dry_run: bool = False) -> None:
         auth_vars = AuthVariables.from_env()
@@ -92,8 +103,23 @@ def verify(
         ToolGlobals: CDFToolConfig,
         dry_run: bool,
         no_prompt: bool = False,
-    ) -> None:
+        demo_user: str | None = None,
+    ) -> VerifyAuthResult:
+        """Authorization verification for the Toolkit.
+
+        Args:
+            ToolGlobals: The Toolkit configuration.
+            dry_run: If the verification should be run in dry-run mode.
+            no_prompt: If the verification should be run without any prompts.
+            demo_user: This is used for demo purposes. If passed a temporary Toolkit group is created
+                and the user is added to the group.
+
+        Returns:
+            VerifyAuthResult: The result of the verification.
+        """
+
         is_interactive = not no_prompt
+        is_demo = demo_user is not None
         if ToolGlobals.project is None:
             raise AuthorizationError("CDF_PROJECT is not set.")
         cdf_project = ToolGlobals.project
@@ -116,23 +142,24 @@ def verify(
             raise AuthorizationError("The current user is not member of any groups in the CDF project.")
 
         loader_capabilities, loaders_by_capability_tuple = self._get_capabilities_by_loader(ToolGlobals)
-        toolkit_group = self._create_toolkit_group(loader_capabilities)
+        toolkit_group = self._create_toolkit_group(loader_capabilities, demo_user)
 
-        print(
-            Panel(
-                "The Cognite Toolkit expects the following:\n"
-                " - The principal used with the Toolkit [yellow]should[/yellow] be connected to "
-                "only ONE CDF Group.\n"
-                f" - This group [red]must[/red] be named {toolkit_group.name!r}.\n"
-                f" - The group {toolkit_group.name!r} [red]must[/red] have capabilities to "
-                f"all resources the Toolkit is managing\n"
-                " - All the capabilities [yellow]should[/yellow] be scoped to all resources.",
-                title="Toolkit Access Group",
-                expand=False,
+        if not is_demo:
+            print(
+                Panel(
+                    "The Cognite Toolkit expects the following:\n"
+                    " - The principal used with the Toolkit [yellow]should[/yellow] be connected to "
+                    "only ONE CDF Group.\n"
+                    f" - This group [red]must[/red] be named {toolkit_group.name!r}.\n"
+                    f" - The group {toolkit_group.name!r} [red]must[/red] have capabilities to "
+                    f"all resources the Toolkit is managing\n"
+                    " - All the capabilities [yellow]should[/yellow] be scoped to all resources.",
+                    title="Toolkit Access Group",
+                    expand=False,
+                )
             )
-        )
-        if is_interactive:
-            Prompt.ask("Press enter key to continue...")
+            if is_interactive:
+                Prompt.ask("Press enter key to continue...")
 
         all_groups = ToolGlobals.toolkit_client.iam.groups.list(all=True)
 
@@ -152,7 +179,7 @@ def verify(
                 is_interactive
                 and missing_capabilities
                 and questionary.confirm("Do you want to update the group with the missing capabilities?").ask()
-            ):
+            ) or is_demo:
                 has_added_capabilities = self._update_missing_capabilities(
                     ToolGlobals, cdf_toolkit_group, missing_capabilities, dry_run
                 )
@@ -170,15 +197,18 @@ def verify(
                 and questionary.confirm("Do you want to update the group with the missing capabilities?").ask()
             ):
                 self._update_missing_capabilities(ToolGlobals, cdf_toolkit_group, missing_capabilities, dry_run)
+        elif is_demo:
+            # We create the group for the demo user
+            cdf_toolkit_group = self._create_toolkit_group_in_cdf(ToolGlobals, toolkit_group)
         else:
             print(f"Group {toolkit_group.name!r} does not exist in the CDF project.")
-            cdf_toolkit_group = self._create_toolkit_group_in_cdf(
+            cdf_toolkit_group = self._create_toolkit_group_in_cdf_interactive(
                 ToolGlobals, toolkit_group, all_groups, is_interactive, dry_run
             )
         if cdf_toolkit_group is None:
-            return None
+            return VerifyAuthResult()
 
-        if not is_user_in_toolkit_group:
+        if not is_demo and not is_user_in_toolkit_group:
             print(
                 Panel(
                     f"To use the Toolkit, for example, 'cdf deploy', [red]you need to switch[/red] "
@@ -187,23 +217,30 @@ def verify(
                     expand=False,
                 )
             )
-            return None
-
-        self.check_count_group_memberships(user_groups)
-
-        self.check_source_id_usage(all_groups, cdf_toolkit_group)
-
-        if extra := self.check_duplicated_names(all_groups, cdf_toolkit_group):
-            if is_interactive and questionary.confirm("Do you want to delete the extra groups?", default=True).ask():
-                try:
-                    ToolGlobals.toolkit_client.iam.groups.delete(extra.as_ids())
-                except CogniteAPIError as e:
-                    raise ResourceDeleteError(f"Unable to delete the extra groups.\n{e}")
-                print(f"  [bold green]OK[/] - Deleted {len(extra)} duplicated groups.")
-
-        self.check_function_service_status(ToolGlobals.toolkit_client, dry_run, has_added_capabilities)
+            return VerifyAuthResult(function_status=None, toolkit_group_id=cdf_toolkit_group.id)
+
+        if not is_demo:
+            self.check_count_group_memberships(user_groups)
+
+            self.check_source_id_usage(all_groups, cdf_toolkit_group)
+
+            if extra := self.check_duplicated_names(all_groups, cdf_toolkit_group):
+                if (
+                    is_interactive
+                    and questionary.confirm("Do you want to delete the extra groups?", default=True).ask()
+                ):
+                    try:
+                        ToolGlobals.toolkit_client.iam.groups.delete(extra.as_ids())
+                    except CogniteAPIError as e:
+                        raise ResourceDeleteError(f"Unable to delete the extra groups.\n{e}")
+                    print(f"  [bold green]OK[/] - Deleted {len(extra)} duplicated groups.")
+
+        function_status = self.check_function_service_status(
+            ToolGlobals.toolkit_client, dry_run, has_added_capabilities
+        )
+        return VerifyAuthResult(cdf_toolkit_group.id, function_status)
 
-    def _create_toolkit_group_in_cdf(
+    def _create_toolkit_group_in_cdf_interactive(
         self,
         ToolGlobals: CDFToolConfig,
         toolkit_group: GroupWrite,
@@ -246,6 +283,14 @@ def _create_toolkit_group_in_cdf(
             )
             if not questionary.confirm("This is NOT recommended. Do you want to continue?", default=False).ask():
                 return None
+
+        return self._create_toolkit_group_in_cdf(ToolGlobals, toolkit_group)
+
+    @staticmethod
+    def _create_toolkit_group_in_cdf(
+        ToolGlobals: CDFToolConfig,
+        toolkit_group: GroupWrite,
+    ) -> Group:
         created = ToolGlobals.toolkit_client.iam.groups.create(toolkit_group)
         print(
             f"  [bold green]OK[/] - Created new group {created.name}. It now has {len(created.capabilities or [])} capabilities."
@@ -342,9 +387,9 @@ def _update_missing_capabilities(
         return True
 
     @staticmethod
-    def _create_toolkit_group(loader_capabilities: list[Capability]) -> GroupWrite:
+    def _create_toolkit_group(loader_capabilities: list[Capability], demo_user: str | None) -> GroupWrite:
         toolkit_group = GroupWrite(
-            name=TOOLKIT_SERVICE_PRINCIPAL_GROUP_NAME,
+            name=TOOLKIT_SERVICE_PRINCIPAL_GROUP_NAME if demo_user is None else TOOLKIT_DEMO_GROUP_NAME,
             capabilities=[
                 *loader_capabilities,
                 # Add project ACL to be able to list and read projects, as the
@@ -354,6 +399,8 @@ def _create_toolkit_group(loader_capabilities: list[Capability]) -> GroupWrite:
                 ),
             ],
         )
+        if demo_user:
+            toolkit_group.members = [demo_user]
         return toolkit_group
 
     @staticmethod
@@ -545,7 +592,9 @@ def _merge_capabilities(capability_list: list[Capability]) -> list[Capability]:
             for (cap_cls, scope), actions in actions_by_scope_and_cls.items()
         ]
 
-    def check_function_service_status(self, client: ToolkitClient, dry_run: bool, has_added_capabilities: bool) -> None:
+    def check_function_service_status(
+        self, client: ToolkitClient, dry_run: bool, has_added_capabilities: bool
+    ) -> str | None:
         print("Checking function service status...")
         has_function_read_access = self.has_function_rights(client, [FunctionsAcl.Action.Read], has_added_capabilities)
         if not has_function_read_access:
@@ -570,21 +619,20 @@ def check_function_service_status(self, client: ToolkitClient, dry_run: bool, ha
             )
             if not has_function_write_access:
                 self.warn(HighSeverityWarning("Cannot activate function service, missing function write access."))
-                return None
+                return function_status.status
             try:
                 client.functions.activate()
             except CogniteAPIError as e:
                 self.warn(HighSeverityWarning(f"Unable to activate function service.\n{e}"))
-                return None
+                return function_status.status
             print(
                 "  [bold green]OK[/] - Function service has been activated. "
                 "This may take up to 2 hours to take effect."
             )
-
         else:
             print("  [bold green]OK[/] - Function service has been activated.")
 
-        return None
+        return function_status.status
 
     def has_function_rights(
         self, client: ToolkitClient, actions: list[FunctionsAcl.Action], has_added_capabilities: bool

cognite_toolkit/_cdf_tk/constants.py

@@ -13,6 +13,7 @@
 # The environment file:
 
 TOOLKIT_SERVICE_PRINCIPAL_GROUP_NAME = "cognite_toolkit_service_principal"
+TOOLKIT_DEMO_GROUP_NAME = "cognite_toolkit_demo"
 
 # This is the default Cognite app registration for Entra with device code enabled
 # to be used with the Toolkit.

cognite_toolkit/_cdf_tk/utils/auth.py

@@ -479,6 +479,7 @@ def __init__(
         cluster: str | None = None,
         project: str | None = None,
         cdf_url: str | None = None,
+        auth_vars: AuthVariables | None = None,
         skip_initialization: bool = False,
     ) -> None:
         self._cache = self._Cache()
@@ -514,7 +515,7 @@ def __init__(
             self._initialize_in_browser()
             return
 
-        self._auth_vars = AuthVariables.from_env(self._environ)
+        self._auth_vars = auth_vars or AuthVariables.from_env(self._environ)
         if not skip_initialization:
             self.initialize_from_auth_variables(self._auth_vars)
         self._login_flow = self._auth_vars.login_flow
@@ -690,6 +691,12 @@ def project(self) -> str:
             raise ValueError("Project is not initialized.")
         return self._project
 
+    @property
+    def cdf_cluster(self) -> str:
+        if self._cluster is None:
+            raise ValueError("Cluster is not initialized.")
+        return self._cluster
+
     @overload
     def environ(self, attr: str, default: str | None = None, fail: Literal[True] = True) -> str: ...
 

cognite_toolkit/demo/_base.py

@@ -2,12 +2,13 @@
 import textwrap
 from pathlib import Path
 
+from cognite.client.data_classes import UserProfile
 from rich import print
 from rich.panel import Panel
 
-from cognite_toolkit._cdf_tk.commands import BuildCommand, DeployCommand, ModulesCommand
+from cognite_toolkit._cdf_tk.commands import AuthCommand, BuildCommand, DeployCommand, ModulesCommand
 from cognite_toolkit._cdf_tk.loaders import LOADER_BY_FOLDER_NAME
-from cognite_toolkit._cdf_tk.utils.auth import CDFToolConfig
+from cognite_toolkit._cdf_tk.utils.auth import AuthVariables, CDFToolConfig
 
 
 class CogniteToolkitDemo:
@@ -42,11 +43,58 @@ def _organization_dir(self) -> Path:
         organization_path.mkdir(exist_ok=True)
         return organization_path
 
-    def quickstart(self) -> None:
+    def quickstart(self, client_id: str | None = None, client_secret: str | None = None) -> None:
         print(Panel("Running Toolkit QuickStart..."))
-        # Lookup user ID to add user ID to the group to run the workflow
         user = self._cdf_tool_config.toolkit_client.iam.user_profiles.me()
+        if sum([client_id is None, client_secret is None]) == 1:
+            raise ValueError("Both client_id and client_secret must be provided or neither.")
+        if client_id is None and client_secret is None:
+            print("Client ID and secret not provided. Assuming user has all the necessary permissions.")
+            self._init_build_deploy(user)
+            return
 
+        group_id: int | None = None
+        try:
+            # Lookup user ID to add user ID to the group to run the workflow
+            auth = AuthCommand()
+            auth_result = auth.verify(
+                self._cdf_tool_config,
+                dry_run=False,
+                no_prompt=True,
+                demo_user=client_id,
+            )
+            group_id = auth_result.toolkit_group_id
+            if auth_result.function_status is None:
+                print(Panel("Unknown function status. If the demo fails, please check that functions are activated"))
+            elif auth_result.function_status == "requested":
+                print(
+                    Panel(
+                        "Function status is requested. Please wait for the function status to be activated before running the demo."
+                    )
+                )
+                return
+            elif auth_result.function_status == "inactive":
+                print(Panel("Function status is inactive. Cannot run demo without functions."))
+                return
+
+            print("Switching to the demo service principal...")
+            self._cdf_tool_config = CDFToolConfig(
+                auth_vars=AuthVariables(
+                    cluster=self._cdf_tool_config.cdf_cluster,
+                    project=self._cdf_tool_config.project,
+                    login_flow="client_credentials",
+                    provider="other",
+                    client_id=client_id,
+                    client_secret=client_secret,
+                    token_url=f"{self._cdf_tool_config.toolkit_client.config.base_url}/oauth2/token",
+                )
+            )
+            self._init_build_deploy(user)
+        finally:
+            if group_id is not None:
+                self._cdf_tool_config.toolkit_client.iam.groups.delete(id=group_id)
+
+    def _init_build_deploy(self, user: UserProfile) -> None:
         modules_cmd = ModulesCommand()
         modules_cmd.run(
             lambda: modules_cmd.init(