From 5e8442e531ae1d71627893723c0cc8d5f0075bbc Mon Sep 17 00:00:00 2001
From: Romazes <romazes20@gmail.com>
Date: Wed, 18 Sep 2024 23:42:44 +0300
Subject: [PATCH] feat: support `account-id` config

---
 lean/models/api.py           |  1 +
 lean/models/configuration.py | 31 +++++++++++++++++++++++++++++++
 lean/models/json_module.py   | 12 +++++++++++-
 3 files changed, 43 insertions(+), 1 deletion(-)

diff --git a/lean/models/api.py b/lean/models/api.py
index 14851de3..66c3cb30 100644
--- a/lean/models/api.py
+++ b/lean/models/api.py
@@ -23,6 +23,7 @@
 
 class QCAuth0Authorization(WrappedBaseModel):
     authorization: Optional[Dict[str, str]]
+    accountIds: Optional[List[str]]
 
 class ProjectEncryptionKey(WrappedBaseModel):
     id: str
diff --git a/lean/models/configuration.py b/lean/models/configuration.py
index 867abb29..949fd0c3 100644
--- a/lean/models/configuration.py
+++ b/lean/models/configuration.py
@@ -121,6 +121,8 @@ def factory(config_json_object) -> 'Configuration':
             return BrokerageEnvConfiguration.factory(config_json_object)
         elif config_json_object["type"] == "oauth-token":
             return AuthConfiguration.factory(config_json_object)
+        elif config_json_object["type"] == "input-account-id":
+            return AccountIdsConfiguration.factory(config_json_object)
         else:
             raise ValueError(
                 f'Undefined input method type {config_json_object["type"]}')
@@ -419,6 +421,35 @@ def ask_user_for_input(self, default_value, logger: Logger, hide_input: bool = F
         raise ValueError(f'user input not allowed with {self.__class__.__name__}')
 
 
+class AccountIdsConfiguration(InternalInputUserInput):
+
+    def __init__(self, config_json_object):
+        super().__init__(config_json_object)
+
+    def factory(config_json_object) -> 'AccountIdsConfiguration':
+        """Creates an instance of the child classes.
+
+        :param config_json_object: the json object dict with configuration info
+        :return: An instance of AuthConfiguration.
+        """
+        if config_json_object["type"] == "input-account-id":
+            return AccountIdsConfiguration(config_json_object)
+        else:
+            raise ValueError(
+                f'Undefined input method type {config_json_object["type"]}')
+
+    def ask_user_for_input(self, default_value, logger: Logger, hide_input: bool = False):
+        """Prompts user to provide input while validating the type of input
+        against the expected type
+
+        :param default_value: The default to prompt to the user.
+        :param logger: The instance of logger class.
+        :param hide_input: Whether to hide the input (not used for this type of input, which is never hidden).
+        :return: The value provided by the user.
+        """
+        raise ValueError(f'user input not allowed with {self.__class__.__name__}')
+
+
 class FilterEnvConfiguration(BrokerageEnvConfiguration):
     """This class adds extra filters to user filters."""
 
diff --git a/lean/models/json_module.py b/lean/models/json_module.py
index 8b18f4ef..ec1d8b9c 100644
--- a/lean/models/json_module.py
+++ b/lean/models/json_module.py
@@ -21,8 +21,9 @@
 from lean.components.util.logger import Logger
 from lean.constants import MODULE_TYPE, MODULE_PLATFORM, MODULE_CLI_PLATFORM
 from lean.container import container
+from lean.models.logger import Option
 from lean.models.configuration import BrokerageEnvConfiguration, Configuration, InternalInputUserInput, \
-    PathParameterUserInput, AuthConfiguration
+    PathParameterUserInput, AuthConfiguration, AccountIdsConfiguration
 from copy import copy
 from abc import ABC
 
@@ -212,6 +213,15 @@ def config_build(self,
                 logger.debug(f'auth: {auth_authorizations}')
                 configuration._value = auth_authorizations.authorization
                 continue
+            elif isinstance(configuration, AccountIdsConfiguration):
+                account_ids = get_authorization(container.api_client.auth0, self._display_name.lower(), logger).accountIds
+                if account_ids and len(account_ids) > 0:
+                    logger.debug(f'accountIds: {account_ids}')
+                    options = [Option(id=data_type, label=data_type) for data_type in account_ids]
+                    account_id = container.logger.prompt_list(f"Chosen {self._display_name.lower()} account ID", options)
+                    logger.debug(f'user choose account ID: {account_id}')
+                    configuration._value = account_id
+                continue
 
             property_name = self.convert_lean_key_to_variable(configuration._id)
             # Only ask for user input if the config wasn't given as an option