Coverage for mindsdb / migrations / versions / 2022-09-06_96d5fef10caa_data_integration_id.py: 30%

1"""data_integration_id 

2 

3Revision ID: 96d5fef10caa 

4Revises: 473e8f239481 

5Create Date: 2022-09-06 15:20:02.382203 

6 

7""" 

8from alembic import op 

9import sqlalchemy as sa 

10from sqlalchemy.sql import text 

11 

12import mindsdb.interfaces.storage.db # noqa 

13 

14 

15# revision identifiers, used by Alembic. 

16revision = '96d5fef10caa' 

17down_revision = '473e8f239481' 

18branch_labels = None 

19depends_on = None 

20 

21 

22def upgrade(): 

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

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

25 

26 conn = op.get_bind() 

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

28 result = conn.execute(text(''' 

29 select 1 from integration where name = 'lightwood'; 

30 ''')).fetchall() 

31 if len(result) == 0: 

32 conn.execute(text(''' 

33 insert into integration (name, engine, data) 

34 values ('lightwood', 'lightwood', '{}') 

35 ''')) 

36 conn.execute(text(''' 

37 update predictor set data_integration_id = integration_id 

38 where exists (select 1 from integration where integration.id = predictor.integration_id); 

39 ''')) 

40 conn.execute(text(''' 

41 update predictor set integration_id = (select id from integration where name = 'lightwood'); 

42 ''')) 

43 session.commit() 

44 

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

46 batch_op.alter_column( 

47 'integration_id', 

48 existing_type=sa.INTEGER(), 

49 nullable=False 

50 ) 

51 batch_op.create_foreign_key('fk_data_integration_id', 'integration', ['data_integration_id'], ['id']) 

52 

53 

54def downgrade(): 

55 conn = op.get_bind() 

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

57 conn.execute(text(''' 

58 update predictor set integration_id = data_integration_id; 

59 ''')) 

60 session.commit() 

61 

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

63 batch_op.drop_constraint('fk_data_integration_id', type_='foreignkey') 

64 batch_op.alter_column( 

65 'integration_id', 

66 existing_type=sa.INTEGER(), 

67 nullable=True 

68 ) 

69 batch_op.drop_column('data_integration_id')