Coverage for mindsdb / interfaces / variables / variables_controller.py: 54%

1import os 

2from typing import Callable 

3 

4from mindsdb_sql_parser import Function, Constant, Variable 

5 

6from mindsdb.utilities import log 

7from mindsdb.interfaces.storage.fs import RESOURCE_GROUP 

8from mindsdb.interfaces.storage.json import get_json_storage 

9from mindsdb.utilities.context import context as ctx 

10 

11 

12logger = log.getLogger(__name__) 

13 

14 

15ENV_VAR_PREFIX = "MDB_" 

16 

17 

18class VariablesController: 

19 def __init__(self) -> None: 

20 self._storage = get_json_storage(resource_id=0, resource_group=RESOURCE_GROUP.SYSTEM) 

21 self._store_key = "variables" 

22 self._data = None 

23 

24 def _get_data(self) -> dict: 

25 if self._data is None: 25 ↛ 29line 25 didn't jump to line 29 because the condition on line 25 was always true

26 self._data = self._storage.get(self._store_key) 

27 if self._data is None: 27 ↛ 29line 27 didn't jump to line 29 because the condition on line 27 was always true

28 self._data = {} 

29 return {**self._data, **ctx.params} 

30 

31 def get_value(self, name: str): 

32 data = self._get_data() 

33 if name not in data: 33 ↛ 34line 33 didn't jump to line 34 because the condition on line 33 was never true

34 raise ValueError(f"Variable {name} is not defined") 

35 return data[name] 

36 

37 def set_value(self, name: str, value): 

38 data = self._get_data() 

39 data[name] = value 

40 self._storage.set(self._store_key, data) 

41 

42 def _from_env(self, name: Constant) -> str: 

43 # gets variable value from environment. 

44 # available names are restricted by ENV_VAR_PREFIX to don't provide access to arbitrary venv variable 

45 

46 var_name = name.value 

47 if not var_name.startswith(ENV_VAR_PREFIX): 

48 raise ValueError(f"Can access only to variable names starting with {ENV_VAR_PREFIX}") 

49 if var_name not in os.environ: 

50 raise ValueError(f"Environment variable {var_name} is not defined") 

51 return os.environ[var_name] 

52 

53 def _get_function(self, name: str) -> Callable: 

54 if name == "from_env": 

55 return self._from_env 

56 raise ValueError(f"Function {name} is not found") 

57 

58 def set_variable(self, name: str, value): 

59 # store new value for variable in database 

60 # if value is a function - extract value using this function 

61 

62 name = name.lower() 

63 if isinstance(value, Function): 

64 fnc = self._get_function(value.op) 

65 value = fnc(*value.args) 

66 

67 elif isinstance(value, Constant): 

68 value = value.value 

69 

70 else: 

71 # ignore 

72 return 

73 

74 self.set_value(name, value) 

75 

76 def fill_parameters(self, var): 

77 # recursively check input and fill Variables if they exist there 

78 

79 if isinstance(var, Variable): 

80 return self.get_value(var.value.lower()) 

81 if isinstance(var, Function): 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true

82 fnc = self._get_function(var.op) 

83 return fnc(*var.args) 

84 elif isinstance(var, dict): 

85 return {key: self.fill_parameters(value) for key, value in var.items()} 

86 elif isinstance(var, list): 86 ↛ 87line 86 didn't jump to line 87 because the condition on line 86 was never true

87 return [self.fill_parameters(value) for value in var] 

88 return var 

89 

90 

91variables_controller = VariablesController()