Coverage for mindsdb / migrations / versions / 2022-08-29_473e8f239481_straighten.py: 27%

1"""straighten 

2 

3Revision ID: 473e8f239481 

4Revises: 6a54ba55872e 

5Create Date: 2022-08-29 11:12:11.307317 

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 = '473e8f239481' 

17down_revision = '6a54ba55872e' 

18branch_labels = None 

19depends_on = None 

20 

21 

22def upgrade(): 

23 conn = op.get_bind() 

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

25 conn.execute(text(''' 

26 delete from file 

27 where exists ( 

28 select 1 from file as t2 

29 where file.name = t2.name 

30 and file.company_id = t2.company_id 

31 and file.id < t2.id 

32 ) 

33 ''')) 

34 session.commit() 

35 

36 with op.batch_alter_table('file', schema=None) as batch_op: 

37 batch_op.create_unique_constraint('unique_file_name_company_id', ['name', 'company_id']) 

38 

39 conn.execute(text(''' 

40 delete from integration where engine is null 

41 ''')) 

42 session.commit() 

43 

44 with op.batch_alter_table('integration', schema=None) as batch_op: 

45 batch_op.alter_column( 

46 'engine', 

47 existing_type=sa.VARCHAR(), 

48 nullable=False 

49 ) 

50 

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

52 batch_op.drop_constraint('unique_predictor_name_company_id', type_='unique') 

53 

54 conn.execute(text(''' 

55 delete from semaphor 

56 ''')) 

57 session.commit() 

58 

59 try: 

60 with op.batch_alter_table('semaphor', schema=None) as batch_op: 

61 batch_op.create_unique_constraint('uniq_const', ['entity_type', 'entity_id']) 

62 except Exception: 

63 pass 

64 

65 

66def downgrade(): 

67 with op.batch_alter_table('semaphor', schema=None) as batch_op: 

68 batch_op.drop_constraint('uniq_const', type_='unique') 

69 

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

71 batch_op.create_unique_constraint('unique_predictor_name_company_id', ['name', 'company_id']) 

72 

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

74 batch_op.alter_column( 

75 'engine', 

76 existing_type=sa.VARCHAR(), 

77 nullable=True 

78 ) 

79 

80 with op.batch_alter_table('file', schema=None) as batch_op: 

81 batch_op.drop_constraint('unique_file_name_company_id', type_='unique')