Coverage for mindsdb / interfaces / model / functions.py: 76%

1from typing import Optional, List 

2 

3from sqlalchemy import null, func 

4 

5import mindsdb.interfaces.storage.db as db 

6from mindsdb.utilities.context import context as ctx 

7import mindsdb.utilities.profiler as profiler 

8 

9 

10class PredictorRecordNotFound(Exception): 

11 def __init__(self, **kwargs): 

12 name = kwargs.get("name") or "-" 

13 predictor_id = kwargs.get("id") or "-" 

14 super().__init__(f"Predictor not found: name='{name}' id='{predictor_id}'") 

15 

16 

17class MultiplePredictorRecordsFound(Exception): 

18 def __init__(self, **kwargs): 

19 name = kwargs.get("name") or "-" 

20 predictor_id = kwargs.get("id") or "-" 

21 super().__init__(f"Found multiple predictor with: name='{name}' id='{predictor_id}'") 

22 

23 

24@profiler.profile() 

25def get_integration_record(name: str) -> db.Integration: 

26 company_id = ctx.company_id 

27 if company_id is None: 27 ↛ 30line 27 didn't jump to line 30 because the condition on line 27 was always true

28 company_id = null() 

29 

30 record = db.session.query(db.Integration).filter_by(company_id=company_id, name=name).first() 

31 return record 

32 

33 

34@profiler.profile() 

35def get_project_record(name: str) -> db.Project: 

36 company_id = ctx.company_id if ctx.company_id is not None else "0" 

37 

38 project_record = ( 

39 db.session.query(db.Project) 

40 .filter( 

41 (func.lower(db.Project.name) == name) 

42 & (db.Project.company_id == company_id) 

43 & (db.Project.deleted_at == null()) 

44 ) 

45 .first() 

46 ) 

47 return project_record 

48 

49 

50@profiler.profile() 

51def get_project_records() -> List[db.Project]: 

52 company_id = ctx.company_id if ctx.company_id is not None else "0" 

53 

54 return ( 

55 db.session.query(db.Project) 

56 .filter((db.Project.company_id == company_id) & (db.Project.deleted_at == null())) 

57 .all() 

58 ) 

59 

60 

61@profiler.profile() 

62def get_predictor_integration(record: db.Predictor) -> db.Integration: 

63 integration_record = db.session.query(db.Integration).filter_by(id=record.integration_id).first() 

64 return integration_record 

65 

66 

67@profiler.profile() 

68def get_predictor_project(record: db.Predictor) -> db.Project: 

69 project_record = db.session.query(db.Project).filter_by(id=record.project_id).first() 

70 return project_record 

71 

72 

73@profiler.profile() 

74def get_model_records( 

75 integration_id=None, 

76 active=True, 

77 deleted_at=null(), 

78 project_name: Optional[str] = None, 

79 ml_handler_name: Optional[str] = None, 

80 **kwargs, 

81): 

82 kwargs["company_id"] = ctx.company_id 

83 if kwargs["company_id"] is None: 83 ↛ 86line 83 didn't jump to line 86 because the condition on line 83 was always true

84 kwargs["company_id"] = null() 

85 

86 if deleted_at is not None: 

87 kwargs["deleted_at"] = deleted_at 

88 if active is not None: 

89 kwargs["active"] = active 

90 

91 if project_name is not None: 91 ↛ 92line 91 didn't jump to line 92 because the condition on line 91 was never true

92 project_record = get_project_record(name=project_name) 

93 if project_record is None: 

94 return [] 

95 kwargs["project_id"] = project_record.id 

96 

97 if ml_handler_name is not None: 

98 ml_handler_record = get_integration_record(name=ml_handler_name) 

99 if ml_handler_record is None: 99 ↛ 101line 99 didn't jump to line 101 because the condition on line 99 was never true

100 # raise Exception(f'unknown ml handler: {ml_handler_name}') 

101 return [] 

102 kwargs["integration_id"] = ml_handler_record.id 

103 

104 if integration_id is not None: 104 ↛ 105line 104 didn't jump to line 105 because the condition on line 104 was never true

105 kwargs["integration_id"] = integration_id 

106 

107 return db.session.query(db.Predictor).filter_by(**kwargs).all() 

108 

109 

110@profiler.profile() 

111def get_model_record( 

112 except_absent=False, 

113 ml_handler_name: Optional[str] = None, 

114 project_name: Optional[str] = None, 

115 active: bool = True, 

116 deleted_at=null(), 

117 version: Optional[int] = None, 

118 **kwargs, 

119): 

120 kwargs["company_id"] = ctx.company_id 

121 if kwargs["company_id"] is None: 121 ↛ 124line 121 didn't jump to line 124 because the condition on line 121 was always true

122 kwargs["company_id"] = null() 

123 

124 kwargs["deleted_at"] = deleted_at 

125 if active is not None: 

126 kwargs["active"] = active 

127 if version is not None: 

128 kwargs["version"] = version 

129 

130 if project_name is not None: 

131 project_record = get_project_record(name=project_name) 

132 if project_record is None: 

133 return None 

134 kwargs["project_id"] = project_record.id 

135 

136 if ml_handler_name is not None: 136 ↛ 137line 136 didn't jump to line 137 because the condition on line 136 was never true

137 ml_handler_record = get_integration_record(name=ml_handler_name) 

138 if ml_handler_record is None: 

139 # raise Exception(f'unknown ml handler: {ml_handler_name}') 

140 return [] 

141 kwargs["integration_id"] = ml_handler_record.id 

142 

143 records = db.session.query(db.Predictor).filter_by(**kwargs).all() 

144 if len(records) > 1: 144 ↛ 145line 144 didn't jump to line 145 because the condition on line 144 was never true

145 raise MultiplePredictorRecordsFound(**kwargs) 

146 if len(records) == 0: 

147 if except_absent is True: 

148 raise PredictorRecordNotFound(**kwargs) 

149 else: 

150 return None 

151 return records[0]