Coverage for mindsdb / migrations / versions / 2022-09-08_87b2df2b83e1_predictor_status.py: 19%

1"""predictor-status 

2 

3Revision ID: 87b2df2b83e1 

4Revises: 96d5fef10caa 

5Create Date: 2022-09-08 14:47:45.238710 

6 

7""" 

8import json 

9 

10from alembic import op 

11import sqlalchemy as sa 

12from sqlalchemy.sql import text 

13 

14import mindsdb.interfaces.storage.db # noqa 

15 

16 

17# revision identifiers, used by Alembic. 

18revision = '87b2df2b83e1' 

19down_revision = '96d5fef10caa' 

20branch_labels = None 

21depends_on = None 

22 

23 

24def upgrade(): 

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

26 batch_op.add_column(sa.Column('status', sa.String(), nullable=True)) 

27 

28 conn = op.get_bind() 

29 session = sa.orm.Session(bind=conn) 

30 predictors = conn.execute(text(''' 

31 select id, data, update_status, json_ai, code from predictor 

32 ''')).fetchall() 

33 

34 for row in predictors: 

35 try: 

36 data = json.loads(row['data']) 

37 except Exception: 

38 data = None 

39 

40 status = None 

41 # assume older models are complete, only temporary 

42 if 'status' in (data or {}): 

43 status = data['status'] 

44 elif 'error' in (data or {}): 

45 status = 'error' 

46 elif row['update_status'] == 'available': 

47 status = 'complete' 

48 elif row['json_ai'] is None and row['code'] is None: 

49 status = 'generating' 

50 elif data is None: 

51 status = 'error' 

52 elif 'training_log' in (data or {}): 

53 status = 'training' 

54 elif 'error' not in (data or {}): 

55 status = 'complete' 

56 else: 

57 status = 'error' 

58 

59 conn.execute( 

60 text(""" 

61 update predictor 

62 set status = :status 

63 where id = :predictor_id 

64 """), { 

65 'status': status, 

66 'predictor_id': row['id'] 

67 } 

68 ) 

69 session.commit() 

70 

71 

72def downgrade(): 

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

74 batch_op.drop_column('status')