From 2b52337da98f8383db55557b1ac82b2d60d9d570 Mon Sep 17 00:00:00 2001
From: montezdesousa <79287829+montezdesousa@users.noreply.github.com>
Date: Sun, 9 Apr 2023 23:10:37 +0100
Subject: [PATCH] Keep hub sources updated and allow only available sources in
 the terminal (#4711)

* loop and update sources

* clean update sources dict

* use only allowed sources and update hub sources on login

* silently update hub

* ruff

* update only allowed sources now

* clean up some funcs

* fix dict on tests

* mypy

* remove unused method

* change class arg name sources_dict

* change attribute name in sources model

* remove allowed from SourcesModel

* change from user_choices to choices

* change default msg
---
 openbb_terminal/core/models/sources_model.py  | 24 +++++--
 openbb_terminal/core/session/current_user.py  | 14 ++--
 openbb_terminal/core/session/hub_model.py     | 19 +++--
 openbb_terminal/core/session/local_model.py   |  9 ++-
 openbb_terminal/core/session/session_model.py | 15 ++++
 .../core/session/sources_handler.py           | 69 ++++++++++++++++++-
 openbb_terminal/core/session/utils.py         | 15 ++++
 openbb_terminal/rich_config.py                |  2 +-
 openbb_terminal/sources_controller.py         | 26 +++----
 .../session/test_session_model.py             | 12 ++++
 10 files changed, 166 insertions(+), 39 deletions(-)

diff --git a/openbb_terminal/core/models/sources_model.py b/openbb_terminal/core/models/sources_model.py
index b90b47d54a3..ad8d885bff7 100644
--- a/openbb_terminal/core/models/sources_model.py
+++ b/openbb_terminal/core/models/sources_model.py
@@ -1,5 +1,5 @@
 import json
-from typing import Dict
+from typing import Dict, List
 
 from pydantic.dataclasses import Field, dataclass
 
@@ -29,15 +29,27 @@ def read_default_sources() -> Dict:
         return {}
 
 
+__allowed = read_default_sources()
+
+
+def get_allowed_sources() -> Dict:
+    """Get allowed sources.
+
+    Returns
+    -------
+    Dict
+        Dictionary with sources
+    """
+    return __allowed
+
+
 @dataclass(config=dict(validate_assignment=True))
 class SourcesModel(BaseModel):
     """Model for sources."""
 
-    sources_dict: Dict = Field(default_factory=lambda: read_default_sources())
+    choices: Dict[str, List[str]] = Field(
+        default_factory=lambda: get_allowed_sources(),
+    )
 
     def __repr__(self):
         return super().__repr__()
-
-    def update(self, other: Dict):
-        """Update sources dict."""
-        self.sources_dict.update(other)
diff --git a/openbb_terminal/core/session/current_user.py b/openbb_terminal/core/session/current_user.py
index 1275c7c7140..fe6cc057e46 100644
--- a/openbb_terminal/core/session/current_user.py
+++ b/openbb_terminal/core/session/current_user.py
@@ -13,14 +13,13 @@ from openbb_terminal.core.models import (
     UserModel,
 )
 from openbb_terminal.core.session.env_handler import read_env
-from openbb_terminal.core.session.sources_handler import read_sources
+from openbb_terminal.core.session.sources_handler import load_file_to_model
 from openbb_terminal.core.session.utils import load_dict_to_model
 
 __env_dict = read_env()
 __credentials = load_dict_to_model(__env_dict, CredentialsModel)  # type: ignore
 __preferences = load_dict_to_model(__env_dict, PreferencesModel)
-__sources = SourcesModel()  # type: ignore
-__sources.update(read_sources(Path(__preferences.USER_DATA_SOURCES_FILE)))
+__sources = load_file_to_model(Path(__preferences.USER_DATA_SOURCES_FILE))
 
 __profile = ProfileModel()
 __local_user = UserModel(  # type: ignore
@@ -64,8 +63,7 @@ def set_default_user():
     env_dict = read_env()
     credentials = load_dict_to_model(env_dict, CredentialsModel)
     preferences = load_dict_to_model(env_dict, PreferencesModel)
-    sources = SourcesModel()
-    sources.update(read_sources(Path(preferences.USER_DATA_SOURCES_FILE)))
+    sources = load_file_to_model(Path(__preferences.USER_DATA_SOURCES_FILE))
     profile = ProfileModel()
     default_user = UserModel(  # type: ignore
         credentials=credentials,
@@ -134,15 +132,15 @@ def set_credential(name: str, value: str):
     set_current_user(updated_user)
 
 
-def set_sources(sources_dict: Dict):
+def set_sources(choices: Dict):
     """Set sources
 
     Parameters
     ----------
-    sources_dict : Dict
+    choices : Dict
         Sources dict
     """
     current_user = get_current_user()
-    updated_sources = dataclasses.replace(current_user.sources, sources_dict=sources_dict)  # type: ignore
+    updated_sources = dataclasses.replace(current_user.sources, choices=choices)  # type: ignore
     updated_user = dataclasses.replace(current_user, sources=updated_sources)  # type: ignore
     set_current_user(updated_user)
diff --git a/openbb_terminal/core/session/hub_model.py b/openbb_terminal/core/session/hub_model.py
index a9d217d61db..403dfe54db9 100644
--- a/openbb_terminal/core/session/hub_model.py
+++ b/openbb_terminal/core/session/hub_model.py
@@ -284,6 +284,7 @@ def upload_user_field(
     auth_header: str,
     base_url: str = BASE_URL,
     timeout: int = TIMEOUT,
+    silent: bool = False,
 ) -> Optional[requests.Response]:
     """Send user field to the server.
 
@@ -299,15 +300,19 @@ def upload_user_field(
         The base url, by default BASE_URL
     timeout : int
         The timeout, by default TIMEOUT
+    silent : bool
+        Whether to print the status, by default False
 
     Returns
     -------
     Optional[requests.Response]
         The response from the put request.
     """
-    data: Dict[str, dict] = {key: value}
     try:
-        console.print("Sending to OpenBB hub...")
+        console_print = console.print if not silent else lambda *args, **kwargs: None
+        data: Dict[str, dict] = {key: value}
+
+        console_print("Sending to OpenBB hub...")
         response = requests.put(
             url=base_url + "user",
             headers={"Authorization": auth_header},
@@ -315,18 +320,18 @@ def upload_user_field(
             timeout=timeout,
         )
         if response.status_code == 200:
-            console.print("[green]Saved remotely.[/green]")
+            console_print("[green]Saved remotely.[/green]")
         else:
-            console.print("[red]Failed to save remotely.[/red]")
+            console_print("[red]Failed to save remotely.[/red]")
         return response
     except requests.exceptions.ConnectionError:
-        console.print(f"\n{CONNECTION_ERROR_MSG}")
+        console_print(f"\n{CONNECTION_ERROR_MSG}")
         return None
     except requests.exceptions.Timeout:
-        console.print(f"\n{CONNECTION_TIMEOUT_MSG}")
+        console_print(f"\n{CONNECTION_TIMEOUT_MSG}")
         return None
     except Exception:
-        console.print("[red]Failed to save remotely.[/red]")
+        console_print("[red]Failed to save remotely.[/red]")
         return None
 
 
diff --git a/openbb_terminal/core/session/local_model.py b/openbb_terminal/core/session/local_model.py
index df77ab574db..f5b5c4f02b3 100644
--- a/openbb_terminal/core/session/local_model.py
+++ b/openbb_terminal/core/session/local_model.py
@@ -8,6 +8,7 @@ from openbb_terminal.core.config.paths import (
     MISCELLANEOUS_DIRECTORY,
     SESSION_FILE_PATH,
 )
+from openbb_terminal.core.models.sources_model import get_allowed_sources
 from openbb_terminal.core.session.current_user import (
     get_current_user,
     get_env_dict,
@@ -15,6 +16,7 @@ from openbb_terminal.core.session.current_user import (
     set_preference,
     set_sources,
 )
+from openbb_terminal.core.session.sources_handler import merge_sources
 from openbb_terminal.rich_config import console
 
 
@@ -200,6 +202,7 @@ def set_sources_from_hub(configs: dict):
         The configurations.
     """
     if configs:
-        sources = configs.get("features_sources", {}) or {}
-        if sources:
-            set_sources(sources)
+        incoming = configs.get("features_sources", {}) or {}
+        if incoming:
+            choices = merge_sources(incoming=incoming, allowed=get_allowed_sources())
+            set_sources(choices)
diff --git a/openbb_terminal/core/session/session_model.py b/openbb_terminal/core/session/session_model.py
index 25095fc1800..9c5bf9b64a1 100644
--- a/openbb_terminal/core/session/session_model.py
+++ b/openbb_terminal/core/session/session_model.py
@@ -21,6 +21,8 @@ from openbb_terminal.core.session.current_user import (
     set_current_user,
     set_default_user,
 )
+from openbb_terminal.core.session.sources_handler import get_updated_hub_sources
+from openbb_terminal.core.session.utils import run_thread
 from openbb_terminal.helper_funcs import system_clear
 from openbb_terminal.loggers import setup_logging
 from openbb_terminal.rich_config import console
@@ -100,6 +102,19 @@ def login(session: dict) -> LoginStatus:
             set_current_user(hub_user)
             Local.apply_configs(configs=configs)
             Local.update_flair()
+
+            # Update user sources in backend
+            updated_sources = get_updated_hub_sources(configs)
+            if updated_sources:
+                run_thread(
+                    Hub.upload_user_field,
+                    {
+                        "key": "features_sources",
+                        "value": updated_sources,
+                        "auth_header": get_current_user().profile.get_auth_header(),
+                        "silent": True,
+                    },
+                )
             return LoginStatus.SUCCESS
         if response.status_code == 401:
             return LoginStatus.UNAUTHORIZED
diff --git a/openbb_terminal/core/session/sources_handler.py b/openbb_terminal/core/session/sources_handler.py
index 50f86927a65..83901db44c2 100644
--- a/openbb_terminal/core/session/sources_handler.py
+++ b/openbb_terminal/core/session/sources_handler.py
@@ -3,6 +3,7 @@ import os
 from pathlib import Path
 from typing import Dict
 
+from openbb_terminal.core.models.sources_model import SourcesModel, get_allowed_sources
 from openbb_terminal.core.sources.utils import extend, flatten
 
 
@@ -26,7 +27,7 @@ def read_sources(path: Path) -> Dict:
         return {}
     except Exception as e:
         print(f"\nFailed to read data sources file: {path}\n{e}\n")
-        print("Using OpenBB defaults.")
+        print("Falling back to OpenBB default sources.")
         return {}
 
 
@@ -45,3 +46,69 @@ def write_sources(sources: Dict, path: Path):
             json.dump(extend(sources), f, indent=4)
     except Exception as e:
         print(f"\nFailed to write data sources file: {path}\n{e}\n")
+
+
+def merge_sources(incoming: Dict, allowed: Dict):
+    """Update sources dict if command and sources allowed.
+
+    Parameters
+    ----------
+    incoming : Dict
+        Dictionary with incoming sources
+    allowed : Dict
+        Dictionary with allowed sources
+
+    Returns
+    -------
+    Dict
+        Updated sources
+    """
+    choices = allowed.copy()
+
+    for cmd_path, incoming_sources in incoming.items():
+        if cmd_path in allowed:
+            allowed_sources = allowed[cmd_path]
+            filtered_incoming = [s for s in incoming_sources if s in allowed_sources]
+            remaining = [s for s in allowed_sources if s not in incoming_sources]
+            choices[cmd_path] = filtered_incoming + remaining
+    return choices
+
+
+def load_file_to_model(path: Path) -> SourcesModel:
+    """Load sources model from file.
+
+    Parameters
+    ----------
+    path : Path
+        Path to file
+
+    Returns
+    -------
+    SourcesModel
+        Sources model
+    """
+    choices = merge_sources(incoming=read_sources(path), allowed=get_allowed_sources())
+    return SourcesModel(choices=choices)  # type: ignore
+
+
+def get_updated_hub_sources(configs: Dict) -> Dict:
+    """Update hub sources if new source or command path available.
+
+    Parameters
+    ----------
+    default : Dict
+        Dictionary with default sources
+    hub : Dict
+        Dictionary with hub sources
+
+    Returns
+    -------
+    Dict
+        Updated hub sources
+    """
+    if configs:
+        incoming = configs.get("features_sources", {}) or {}
+        if incoming:
+            return merge_sources(incoming=incoming, allowed=get_allowed_sources())
+        return {}
+    return {}
diff --git a/openbb_terminal/core/session/utils.py b/openbb_terminal/core/session/utils.py
index 67cbc6f0b80..04582817475 100644
--- a/openbb_terminal/core/session/utils.py
+++ b/openbb_terminal/core/session/utils.py
@@ -1,3 +1,4 @@
+from threading import Thread
 from typing import Any, Type, TypeVar
 
 from pydantic import ValidationError
@@ -41,3 +42,17 @@ def load_dict_to_model(dictionary: dict, model: Type[T]) -> T:
     except Exception:
         print(f"Error loading {model_name}, using defaults.")
         return model()  # type: ignore
+
+
+def run_thread(target, kwargs):
+    """Run a thread.
+
+    Parameters
+    ----------
+    target : function
+        The target function.
+    args : tuple
+        The arguments.
+    """
+    thread = Thread(target=target, kwargs=kwargs)
+    thread.start()
diff --git a/openbb_terminal/rich_config.py b/openbb_terminal/rich_config.py
index d87ad7b5976..34acae8d727 100644
--- a/openbb_terminal/rich_config.py
+++ b/openbb_terminal/rich_config.py
@@ -63,7 +63,7 @@ def get_ordered_list_sources(command_path: str) -> List:
         The list of sources for the given command.
     """
     command_path = command_path[1:] if command_path.startswith("/") else command_path
-    return get_current_user().sources.sources_dict.get(command_path, [])
+    return get_current_user().sources.choices.get(command_path, [])
 
 
 class MenuText:
diff --git a/openbb_terminal/sources_controller.py b/openbb_terminal/sources_controller.py
index 9ed46ca4c60..eebedbe7904 100644
--- a/openbb_terminal/sources_controller.py
+++ b/openbb_terminal/sources_controller.py
@@ -63,7 +63,7 @@ class SourcesController(BaseController):
         See `BaseController.parse_input()` for details.
         """
         cmd_filter = r"((set\s+--cmd\s+|set\s+-c\s+|set\s+|get\s+--cmd\s+|get\s+-c\s+|get\s+).*?("
-        for cmd in get_current_user().sources.sources_dict:
+        for cmd in get_current_user().sources.choices:
             cmd = cmd.replace("/", r"\/")
             cmd_filter += f"{cmd}|"
         cmd_filter += ")*)"
@@ -95,7 +95,7 @@ class SourcesController(BaseController):
             "--cmd",
             action="store",
             dest="cmd",
-            choices=list(get_current_user().sources.sources_dict.keys()),
+            choices=list(get_current_user().sources.choices.keys()),
             required="-h" not in other_args and "--help" not in other_args,
             help="Command that we want to check the available data sources and the default one.",
             metavar="COMMAND",
@@ -104,8 +104,8 @@ class SourcesController(BaseController):
             other_args.insert(0, "-c")
         ns_parser = self.parse_simple_args(parser, other_args)
         if ns_parser:
-            sources_dict = get_current_user().sources.sources_dict
-            cmd_defaults = sources_dict.get(ns_parser.cmd, None)
+            choices = get_current_user().sources.choices
+            cmd_defaults = choices.get(ns_parser.cmd, None)
             if cmd_defaults is None:
                 console.print(f"[red]'{ns_parser.cmd}' is not a valid command.[/red]\n")
             elif len(cmd_defaults) == 0:
@@ -131,7 +131,7 @@ class SourcesController(BaseController):
             "--cmd",
             action="store",
             dest="cmd",
-            choices=list(get_current_user().sources.sources_dict.keys()),
+            choices=list(get_current_user().sources.choices.keys()),
             required="-h" not in other_args and "--help" not in other_args,
             help="Command that we to select the default data source.",
             metavar="COMMAND",
@@ -158,15 +158,15 @@ class SourcesController(BaseController):
 
         ns_parser = self.parse_simple_args(parser, other_args)
         if ns_parser:
-            sources_dict = get_current_user().sources.sources_dict
-            if ns_parser.source in sources_dict[ns_parser.cmd]:
-                sources_dict[ns_parser.cmd].remove(ns_parser.source)
-                sources_dict[ns_parser.cmd].insert(0, ns_parser.source)
-                set_sources(sources_dict)
+            choices = get_current_user().sources.choices
+            if ns_parser.source in choices[ns_parser.cmd]:
+                choices[ns_parser.cmd].remove(ns_parser.source)
+                choices[ns_parser.cmd].insert(0, ns_parser.source)
+                set_sources(choices)
 
                 if is_local():
                     write_sources(
-                        sources=sources_dict,
+                        sources=choices,
                         path=Path(
                             get_current_user().preferences.USER_DATA_SOURCES_FILE
                         ),
@@ -174,7 +174,7 @@ class SourcesController(BaseController):
                 else:
                     upload_user_field(
                         key="features_sources",
-                        value=sources_dict,
+                        value=choices,
                         auth_header=get_current_user().profile.get_auth_header(),
                     )
                     console.print("")
@@ -186,6 +186,6 @@ class SourcesController(BaseController):
                 console.print(
                     f"[red]'{ns_parser.source}' is not a valid data source for "
                     f"'{ns_parser.cmd}' command.[/red]\n"
-                    f"[param]\nAvailable :[/param] {', '.join(sources_dict[ns_parser.cmd])}\n"
+                    f"[param]\nAvailable :[/param] {', '.join(choices[ns_parser.cmd])}\n"
                 )
         self.update_print_help()
diff --git a/tests/openbb_terminal/session/test_session_model.py b/tests/openbb_terminal/session/test_session_model.py
index 93447f8f5ed..f38cbd606bf 100644
--- a/tests/openbb_terminal/session/test_session_model.py
+++ b/tests/openbb_terminal/session/test_session_model.py
@@ -68,6 +68,9 @@ def test_login_no_response(mocker):
     mock_fetch_user_configs = mocker.patch(path + "Hub.fetch_user_configs")
     mock_apply_configs = mocker.patch(path + "Local.apply_configs")
     mock_update_flair = mocker.patch(path + "Local.update_flair")
+    mock_get_updated_hub_sources = mocker.patch(
+        path + "get_updated_hub_sources", return_value={}
+    )
 
     mock_fetch_user_configs.return_value = None
 
@@ -76,6 +79,7 @@ def test_login_no_response(mocker):
     mock_fetch_user_configs.assert_called_once_with(TEST_SESSION)
     mock_apply_configs.assert_not_called()
     mock_update_flair.assert_not_called()
+    mock_get_updated_hub_sources.assert_not_called()
 
 
 def test_login_fail_response(mocker):
@@ -83,6 +87,9 @@ def test_login_fail_response(mocker):
     mock_fetch_user_configs = mocker.patch(path + "Hub.fetch_user_configs")
     mock_apply_configs = mocker.patch(path + "Local.apply_configs")
     mock_update_flair = mocker.patch(path + "Local.update_flair")
+    mock_get_updated_hub_sources = mocker.patch(
+        path + "get_updated_hub_sources", return_value={}
+    )
 
     response = Response()
     response.status_code = 400
@@ -93,6 +100,7 @@ def test_login_fail_response(mocker):
     mock_fetch_user_configs.assert_called_once_with(TEST_SESSION)
     mock_apply_configs.assert_not_called()
     mock_update_flair.assert_not_called()
+    mock_get_updated_hub_sources.assert_not_called()
 
 
 def test_login_success_response(mocker):
@@ -100,6 +108,9 @@ def test_login_success_response(mocker):
     mock_fetch_user_configs = mocker.patch(path + "Hub.fetch_user_configs")
     mock_apply_configs = mocker.patch(path + "Local.apply_configs")
     mock_update_flair = mocker.patch(path + "Local.update_flair")
+    mock_get_updated_hub_sources = mocker.patch(
+        path + "get_updated_hub_sources", return_value={}
+    )
 
     response = Response()
     response.status_code = 200
@@ -113,6 +124,7 @@ def test_login_success_response(mocker):
     mock_fetch_user_configs.assert_called_once_with(TEST_SESSION)
     mock_apply_configs.assert_called_once()
     mock_update_flair.assert_called_once()
+    mock_get_updated_hub_sources.assert_called_once()
 
 
 @pytest.mark.parametrize(