Coverage for mindsdb / migrations / versions / 2025-05-21_9f150e4f9a05_checkpoint_1.py: 80%

1"""checkpoint_1 

2 

3Revision ID: 9f150e4f9a05 

4Revises: 53502b6d63bf 

5Create Date: 2025-05-21 12:25:55.556388 

6 

7""" 

8import datetime 

9 

10from alembic.operations import Operations 

11import sqlalchemy as sa 

12import mindsdb.interfaces.storage.db # noqa 

13 

14 

15# revision identifiers, used by Alembic. 

16revision = '9f150e4f9a05' 

17down_revision = '53502b6d63bf' 

18branch_labels = None 

19depends_on = None 

20 

21 

22def upgrade(op: Operations = None): 

23 # region skip migration if it is existing app, apply if it is new app 

24 if op is None: 24 ↛ 26line 24 didn't jump to line 26 because the condition on line 24 was never true

25 # 'op' is passed only from migrate.py when applying checkpoin migration 

26 return 

27 connection = op.get_bind() 

28 inspector = sa.inspect(connection) 

29 if 'alembic_version' in inspector.get_table_names(): 29 ↛ 31line 29 didn't jump to line 31 because the condition on line 29 was never true

30 # If version_num exists, then it is existing app 

31 result = connection.execute(sa.text("SELECT version_num FROM alembic_version")) 

32 current_version = result.scalar() 

33 if current_version is not None: 

34 return 

35 # endregion 

36 

37 op.create_table( 

38 'agents', 

39 sa.Column('id', sa.Integer(), nullable=False), 

40 sa.Column('company_id', sa.Integer(), nullable=True), 

41 sa.Column('user_class', sa.Integer(), nullable=True), 

42 sa.Column('name', sa.String(), nullable=False), 

43 sa.Column('project_id', sa.Integer(), nullable=False), 

44 sa.Column('model_name', sa.String(), nullable=True), 

45 sa.Column('provider', sa.String(), nullable=True), 

46 sa.Column('params', sa.JSON(), nullable=True), 

47 sa.Column('updated_at', sa.DateTime(), nullable=True), 

48 sa.Column('created_at', sa.DateTime(), nullable=True), 

49 sa.Column('deleted_at', sa.DateTime(), nullable=True), 

50 sa.PrimaryKeyConstraint('id') 

51 ) 

52 op.create_table( 

53 'chat_bots_history', 

54 sa.Column('id', sa.Integer(), nullable=False), 

55 sa.Column('chat_bot_id', sa.Integer(), nullable=False), 

56 sa.Column('type', sa.String(), nullable=True), 

57 sa.Column('text', sa.String(), nullable=True), 

58 sa.Column('user', sa.String(), nullable=True), 

59 sa.Column('destination', sa.String(), nullable=True), 

60 sa.Column('sent_at', sa.DateTime(), nullable=True), 

61 sa.Column('error', sa.String(), nullable=True), 

62 sa.PrimaryKeyConstraint('id') 

63 ) 

64 op.create_table( 

65 'file', 

66 sa.Column('id', sa.Integer(), nullable=False), 

67 sa.Column('name', sa.String(), nullable=False), 

68 sa.Column('company_id', sa.Integer(), nullable=True), 

69 sa.Column('source_file_path', sa.String(), nullable=False), 

70 sa.Column('file_path', sa.String(), nullable=False), 

71 sa.Column('row_count', sa.Integer(), nullable=False), 

72 sa.Column('columns', mindsdb.interfaces.storage.db.Json(), nullable=False), 

73 sa.Column('created_at', sa.DateTime(), nullable=True), 

74 sa.Column('metadata', sa.JSON(), nullable=True), 

75 sa.Column('updated_at', sa.DateTime(), nullable=True), 

76 sa.PrimaryKeyConstraint('id'), 

77 sa.UniqueConstraint('name', 'company_id', name='unique_file_name_company_id') 

78 ) 

79 op.create_table( 

80 'integration', 

81 sa.Column('id', sa.Integer(), nullable=False), 

82 sa.Column('updated_at', sa.DateTime(), nullable=True), 

83 sa.Column('created_at', sa.DateTime(), nullable=True), 

84 sa.Column('name', sa.String(), nullable=False), 

85 sa.Column('engine', sa.String(), nullable=False), 

86 sa.Column('data', mindsdb.interfaces.storage.db.Json(), nullable=True), 

87 sa.Column('company_id', sa.Integer(), nullable=True), 

88 sa.PrimaryKeyConstraint('id'), 

89 sa.UniqueConstraint('name', 'company_id', name='unique_integration_name_company_id') 

90 ) 

91 op.create_table( 

92 'jobs', 

93 sa.Column('id', sa.Integer(), nullable=False), 

94 sa.Column('company_id', sa.Integer(), nullable=True), 

95 sa.Column('user_class', sa.Integer(), nullable=True), 

96 sa.Column('active', sa.Boolean(), nullable=True), 

97 sa.Column('name', sa.String(), nullable=False), 

98 sa.Column('project_id', sa.Integer(), nullable=False), 

99 sa.Column('query_str', sa.String(), nullable=False), 

100 sa.Column('if_query_str', sa.String(), nullable=True), 

101 sa.Column('start_at', sa.DateTime(), nullable=True), 

102 sa.Column('end_at', sa.DateTime(), nullable=True), 

103 sa.Column('next_run_at', sa.DateTime(), nullable=True), 

104 sa.Column('schedule_str', sa.String(), nullable=True), 

105 sa.Column('deleted_at', sa.DateTime(), nullable=True), 

106 sa.Column('updated_at', sa.DateTime(), nullable=True), 

107 sa.Column('created_at', sa.DateTime(), nullable=True), 

108 sa.PrimaryKeyConstraint('id') 

109 ) 

110 op.create_table( 

111 'jobs_history', 

112 sa.Column('id', sa.Integer(), nullable=False), 

113 sa.Column('company_id', sa.Integer(), nullable=True), 

114 sa.Column('job_id', sa.Integer(), nullable=True), 

115 sa.Column('query_str', sa.String(), nullable=True), 

116 sa.Column('start_at', sa.DateTime(), nullable=True), 

117 sa.Column('end_at', sa.DateTime(), nullable=True), 

118 sa.Column('error', sa.String(), nullable=True), 

119 sa.Column('created_at', sa.DateTime(), nullable=True), 

120 sa.Column('updated_at', sa.DateTime(), nullable=True), 

121 sa.PrimaryKeyConstraint('id'), 

122 sa.UniqueConstraint('job_id', 'start_at', name='uniq_job_history_job_id_start') 

123 ) 

124 op.create_table( 

125 'json_storage', 

126 sa.Column('id', sa.Integer(), nullable=False), 

127 sa.Column('resource_group', sa.String(), nullable=True), 

128 sa.Column('resource_id', sa.Integer(), nullable=True), 

129 sa.Column('name', sa.String(), nullable=True), 

130 sa.Column('content', sa.JSON(), nullable=True), 

131 sa.Column('encrypted_content', sa.LargeBinary(), nullable=True), 

132 sa.Column('company_id', sa.Integer(), nullable=True), 

133 sa.PrimaryKeyConstraint('id') 

134 ) 

135 op.create_table( 

136 'llm_data', 

137 sa.Column('id', sa.Integer(), nullable=False), 

138 sa.Column('input', sa.String(), nullable=False), 

139 sa.Column('output', sa.String(), nullable=False), 

140 sa.Column('model_id', sa.Integer(), nullable=False), 

141 sa.Column('created_at', sa.DateTime(), nullable=True), 

142 sa.Column('updated_at', sa.DateTime(), nullable=True), 

143 sa.PrimaryKeyConstraint('id') 

144 ) 

145 op.create_table( 

146 'llm_log', 

147 sa.Column('id', sa.Integer(), nullable=False), 

148 sa.Column('company_id', sa.Integer(), nullable=False), 

149 sa.Column('api_key', sa.String(), nullable=True), 

150 sa.Column('model_id', sa.Integer(), nullable=True), 

151 sa.Column('model_group', sa.String(), nullable=True), 

152 sa.Column('input', sa.JSON(), nullable=True), 

153 sa.Column('output', sa.JSON(), nullable=True), 

154 sa.Column('start_time', sa.DateTime(), nullable=False), 

155 sa.Column('end_time', sa.DateTime(), nullable=True), 

156 sa.Column('cost', sa.Numeric(precision=5, scale=2), nullable=True), 

157 sa.Column('prompt_tokens', sa.Integer(), nullable=True), 

158 sa.Column('completion_tokens', sa.Integer(), nullable=True), 

159 sa.Column('total_tokens', sa.Integer(), nullable=True), 

160 sa.Column('success', sa.Boolean(), nullable=False), 

161 sa.Column('exception', sa.String(), nullable=True), 

162 sa.Column('traceback', sa.String(), nullable=True), 

163 sa.Column('stream', sa.Boolean(), nullable=True, comment="Is this completion done in 'streaming' mode"), 

164 sa.Column('metadata', sa.JSON(), nullable=True), 

165 sa.PrimaryKeyConstraint('id') 

166 ) 

167 op.create_table( 

168 'project', 

169 sa.Column('id', sa.Integer(), nullable=False), 

170 sa.Column('created_at', sa.DateTime(), nullable=True), 

171 sa.Column('updated_at', sa.DateTime(), nullable=True), 

172 sa.Column('deleted_at', sa.DateTime(), nullable=True), 

173 sa.Column('name', sa.String(), nullable=False), 

174 sa.Column('company_id', sa.Integer(), nullable=True), 

175 sa.Column('metadata', sa.JSON(), nullable=True), 

176 sa.PrimaryKeyConstraint('id'), 

177 sa.UniqueConstraint('name', 'company_id', name='unique_project_name_company_id') 

178 ) 

179 op.create_table( 

180 'queries', 

181 sa.Column('id', sa.Integer(), nullable=False), 

182 sa.Column('company_id', sa.Integer(), nullable=True), 

183 sa.Column('sql', sa.String(), nullable=False), 

184 sa.Column('database', sa.String(), nullable=True), 

185 sa.Column('started_at', sa.DateTime(), nullable=True), 

186 sa.Column('finished_at', sa.DateTime(), nullable=True), 

187 sa.Column('parameters', sa.JSON(), nullable=True), 

188 sa.Column('context', sa.JSON(), nullable=True), 

189 sa.Column('processed_rows', sa.Integer(), nullable=True), 

190 sa.Column('error', sa.String(), nullable=True), 

191 sa.Column('updated_at', sa.DateTime(), nullable=True), 

192 sa.Column('created_at', sa.DateTime(), nullable=True), 

193 sa.PrimaryKeyConstraint('id') 

194 ) 

195 op.create_table( 

196 'query_context', 

197 sa.Column('id', sa.Integer(), nullable=False), 

198 sa.Column('company_id', sa.Integer(), nullable=True), 

199 sa.Column('query', sa.String(), nullable=False), 

200 sa.Column('context_name', sa.String(), nullable=False), 

201 sa.Column('values', sa.JSON(), nullable=True), 

202 sa.Column('updated_at', sa.DateTime(), nullable=True), 

203 sa.Column('created_at', sa.DateTime(), nullable=True), 

204 sa.PrimaryKeyConstraint('id') 

205 ) 

206 op.create_table( 

207 'skills', 

208 sa.Column('id', sa.Integer(), nullable=False), 

209 sa.Column('name', sa.String(), nullable=False), 

210 sa.Column('project_id', sa.Integer(), nullable=False), 

211 sa.Column('type', sa.String(), nullable=False), 

212 sa.Column('params', sa.JSON(), nullable=True), 

213 sa.Column('created_at', sa.DateTime(), nullable=True), 

214 sa.Column('updated_at', sa.DateTime(), nullable=True), 

215 sa.Column('deleted_at', sa.DateTime(), nullable=True), 

216 sa.PrimaryKeyConstraint('id') 

217 ) 

218 op.create_table( 

219 'tasks', 

220 sa.Column('id', sa.Integer(), nullable=False), 

221 sa.Column('company_id', sa.Integer(), nullable=True), 

222 sa.Column('user_class', sa.Integer(), nullable=True), 

223 sa.Column('object_type', sa.String(), nullable=False), 

224 sa.Column('object_id', sa.Integer(), nullable=False), 

225 sa.Column('last_error', sa.String(), nullable=True), 

226 sa.Column('active', sa.Boolean(), nullable=True), 

227 sa.Column('reload', sa.Boolean(), nullable=True), 

228 sa.Column('run_by', sa.String(), nullable=True), 

229 sa.Column('alive_time', sa.DateTime(timezone=True), nullable=True), 

230 sa.Column('updated_at', sa.DateTime(), nullable=True), 

231 sa.Column('created_at', sa.DateTime(), nullable=True), 

232 sa.PrimaryKeyConstraint('id') 

233 ) 

234 op.create_table( 

235 'triggers', 

236 sa.Column('id', sa.Integer(), nullable=False), 

237 sa.Column('name', sa.String(), nullable=False), 

238 sa.Column('project_id', sa.Integer(), nullable=False), 

239 sa.Column('database_id', sa.Integer(), nullable=False), 

240 sa.Column('table_name', sa.String(), nullable=False), 

241 sa.Column('query_str', sa.String(), nullable=False), 

242 sa.Column('columns', sa.String(), nullable=True), 

243 sa.Column('updated_at', sa.DateTime(), nullable=True), 

244 sa.Column('created_at', sa.DateTime(), nullable=True), 

245 sa.PrimaryKeyConstraint('id') 

246 ) 

247 op.create_table( 

248 'agent_skills', 

249 sa.Column('agent_id', sa.Integer(), nullable=False), 

250 sa.Column('skill_id', sa.Integer(), nullable=False), 

251 sa.Column('parameters', sa.JSON(), nullable=True), 

252 sa.ForeignKeyConstraint(['agent_id'], ['agents.id'], ), 

253 sa.ForeignKeyConstraint(['skill_id'], ['skills.id'], ), 

254 sa.PrimaryKeyConstraint('agent_id', 'skill_id') 

255 ) 

256 op.create_table( 

257 'chat_bots', 

258 sa.Column('id', sa.Integer(), nullable=False), 

259 sa.Column('name', sa.String(), nullable=False), 

260 sa.Column('project_id', sa.Integer(), nullable=False), 

261 sa.Column('agent_id', sa.Integer(), nullable=True), 

262 sa.Column('model_name', sa.String(), nullable=True), 

263 sa.Column('database_id', sa.Integer(), nullable=True), 

264 sa.Column('params', sa.JSON(), nullable=True), 

265 sa.Column('updated_at', sa.DateTime(), nullable=True), 

266 sa.Column('created_at', sa.DateTime(), nullable=True), 

267 sa.Column('webhook_token', sa.String(), nullable=True), 

268 sa.ForeignKeyConstraint(['agent_id'], ['agents.id'], name='fk_agent_id'), 

269 sa.PrimaryKeyConstraint('id') 

270 ) 

271 op.create_table( 

272 'predictor', 

273 sa.Column('id', sa.Integer(), nullable=False), 

274 sa.Column('updated_at', sa.DateTime(), nullable=True), 

275 sa.Column('created_at', sa.DateTime(), nullable=True), 

276 sa.Column('deleted_at', sa.DateTime(), nullable=True), 

277 sa.Column('name', sa.String(), nullable=True), 

278 sa.Column('data', mindsdb.interfaces.storage.db.Json(), nullable=True), 

279 sa.Column('to_predict', mindsdb.interfaces.storage.db.Array(), nullable=True), 

280 sa.Column('company_id', sa.Integer(), nullable=True), 

281 sa.Column('mindsdb_version', sa.String(), nullable=True), 

282 sa.Column('native_version', sa.String(), nullable=True), 

283 sa.Column('integration_id', sa.Integer(), nullable=True), 

284 sa.Column('data_integration_ref', mindsdb.interfaces.storage.db.Json(), nullable=True), 

285 sa.Column('fetch_data_query', sa.String(), nullable=True), 

286 sa.Column('learn_args', mindsdb.interfaces.storage.db.Json(), nullable=True), 

287 sa.Column('update_status', sa.String(), nullable=True), 

288 sa.Column('status', sa.String(), nullable=True), 

289 sa.Column('active', sa.Boolean(), nullable=True), 

290 sa.Column('training_data_columns_count', sa.Integer(), nullable=True), 

291 sa.Column('training_data_rows_count', sa.Integer(), nullable=True), 

292 sa.Column('training_start_at', sa.DateTime(), nullable=True), 

293 sa.Column('training_stop_at', sa.DateTime(), nullable=True), 

294 sa.Column('label', sa.String(), nullable=True), 

295 sa.Column('version', sa.Integer(), nullable=True), 

296 sa.Column('code', sa.String(), nullable=True), 

297 sa.Column('lightwood_version', sa.String(), nullable=True), 

298 sa.Column('dtype_dict', mindsdb.interfaces.storage.db.Json(), nullable=True), 

299 sa.Column('project_id', sa.Integer(), nullable=False), 

300 sa.Column('training_phase_current', sa.Integer(), nullable=True), 

301 sa.Column('training_phase_total', sa.Integer(), nullable=True), 

302 sa.Column('training_phase_name', sa.String(), nullable=True), 

303 sa.Column('training_metadata', sa.JSON(), nullable=False), 

304 sa.ForeignKeyConstraint(['integration_id'], ['integration.id'], name='fk_integration_id'), 

305 sa.ForeignKeyConstraint(['project_id'], ['project.id'], name='fk_project_id'), 

306 sa.PrimaryKeyConstraint('id') 

307 ) 

308 with op.batch_alter_table('predictor', schema=None) as batch_op: 

309 batch_op.create_index('predictor_index', ['company_id', 'name', 'version', 'active', 'deleted_at'], unique=True) 

310 

311 op.create_table( 

312 'view', 

313 sa.Column('id', sa.Integer(), nullable=False), 

314 sa.Column('name', sa.String(), nullable=False), 

315 sa.Column('company_id', sa.Integer(), nullable=True), 

316 sa.Column('query', sa.String(), nullable=False), 

317 sa.Column('project_id', sa.Integer(), nullable=False), 

318 sa.ForeignKeyConstraint(['project_id'], ['project.id'], name='fk_project_id'), 

319 sa.PrimaryKeyConstraint('id'), 

320 sa.UniqueConstraint('name', 'company_id', name='unique_view_name_company_id') 

321 ) 

322 op.create_table( 

323 'knowledge_base', 

324 sa.Column('id', sa.Integer(), nullable=False), 

325 sa.Column('name', sa.String(), nullable=False), 

326 sa.Column('project_id', sa.Integer(), nullable=False), 

327 sa.Column('params', sa.JSON(), nullable=True), 

328 sa.Column('vector_database_id', sa.Integer(), nullable=True), 

329 sa.Column('vector_database_table', sa.String(), nullable=True), 

330 sa.Column('embedding_model_id', sa.Integer(), nullable=True), 

331 sa.Column('query_id', sa.Integer(), nullable=True), 

332 sa.Column('created_at', sa.DateTime(), nullable=True), 

333 sa.Column('updated_at', sa.DateTime(), nullable=True), 

334 sa.ForeignKeyConstraint(['embedding_model_id'], ['predictor.id'], name='fk_knowledge_base_embedding_model_id'), 

335 sa.ForeignKeyConstraint(['vector_database_id'], ['integration.id'], name='fk_knowledge_base_vector_database_id'), 

336 sa.PrimaryKeyConstraint('id'), 

337 sa.UniqueConstraint('name', 'project_id', name='unique_knowledge_base_name_project_id') 

338 ) 

339 

340 # Insert default project 

341 op.bulk_insert( 

342 sa.table( 

343 'project', 

344 sa.Column('name', sa.String()), 

345 sa.Column('company_id', sa.Integer()), 

346 sa.Column('metadata', sa.JSON()), 

347 sa.Column('created_at', sa.DateTime()) 

348 ), 

349 [{ 

350 'name': 'mindsdb', 

351 'company_id': 0, 

352 'metadata': {'is_default': True}, 

353 'created_at': datetime.datetime.now() 

354 }] 

355 ) 

356 

357 

358def downgrade(): 

359 # do nothging, since it is checkpoint migration 

360 pass