Coverage for mindsdb / migrations / versions / 2022-11-07_1e60096fc817_predictor_version.py: 29%

1"""predictor_version 

2 

3Revision ID: 1e60096fc817 

4Revises: 43c52d23845a 

5Create Date: 2022-11-07 16:43:47.301692 

6 

7""" 

8from alembic import op 

9import sqlalchemy as sa 

10from sqlalchemy.sql import text 

11 

12 

13# revision identifiers, used by Alembic. 

14revision = '1e60096fc817' 

15down_revision = '43c52d23845a' 

16branch_labels = None 

17depends_on = None 

18 

19 

20def upgrade(): 

21 # ### commands auto generated by Alembic - please adjust! ### 

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

23 batch_op.add_column(sa.Column('label', sa.String(), nullable=True)) 

24 batch_op.add_column(sa.Column('version', sa.Integer(), nullable=True)) 

25 

26 # update current predictor versions 

27 conn = op.get_bind() 

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

29 

30 key0 = (None, None, None) 

31 models = conn.execute(sa.text(''' 

32 select company_id, project_id, name, id 

33 from predictor 

34 order by company_id, project_id, name, created_at 

35 ''')).fetchall() 

36 for model in models: 

37 key = (model['company_id'], model['project_id'], model['name'].lower()) 

38 

39 # it is different name or project or company 

40 if key != key0: 

41 version = 1 

42 key0 = key 

43 

44 conn.execute( 

45 text(""" 

46 update predictor 

47 set version = :version 

48 where id = :id 

49 """), { 

50 'version': version, 

51 'id': model['id'] 

52 } 

53 ) 

54 version += 1 

55 

56 session.commit() 

57 # ### end Alembic commands ### 

58 

59 

60def downgrade(): 

61 # ### commands auto generated by Alembic - please adjust! ### 

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

63 batch_op.drop_column('version') 

64 batch_op.drop_column('label') 

65 

66 # ### end Alembic commands ###