Coverage for mindsdb / migrations / versions / 2022-10-14_43c52d23845a_projects.py: 20%

1"""projects 

2 

3Revision ID: 43c52d23845a 

4Revises: cada7d2be947 

5Create Date: 2022-10-14 09:59:44.589745 

6 

7""" 

8from alembic import op 

9import sqlalchemy as sa 

10from sqlalchemy.sql import text 

11 

12import mindsdb.interfaces.storage.db as db 

13 

14 

15# revision identifiers, used by Alembic. 

16revision = '43c52d23845a' 

17down_revision = 'cada7d2be947' 

18branch_labels = None 

19depends_on = None 

20 

21 

22def upgrade(): 

23 op.create_table( 

24 'project', 

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

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

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

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

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

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

31 sa.PrimaryKeyConstraint('id'), 

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

33 ) 

34 

35 project_table = sa.Table( 

36 'project', 

37 sa.MetaData(), 

38 sa.Column('id', sa.Integer()), 

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

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

41 ) 

42 

43 conn = op.get_bind() 

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

45 

46 conn.execute( 

47 project_table.insert().values( 

48 name='mindsdb' 

49 ) 

50 ) 

51 

52 project_record = conn.execute( 

53 project_table.select().where(project_table.c.name == 'mindsdb') 

54 ).fetchone() 

55 

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

57 batch_op.add_column(sa.Column('project_id', sa.Integer())) 

58 batch_op.create_foreign_key('fk_project_id', 'project', ['project_id'], ['id']) 

59 

60 conn.execute(sa.sql.text(''' 

61 update predictor set project_id = :project_id 

62 '''), {'project_id': project_record.id}) 

63 

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

65 batch_op.alter_column( 

66 'project_id', 

67 existing_type=sa.INTEGER(), 

68 nullable=False 

69 ) 

70 

71 with op.batch_alter_table('view', schema=None) as batch_op: 

72 batch_op.add_column(sa.Column('project_id', sa.Integer())) 

73 batch_op.create_foreign_key('fk_project_id', 'project', ['project_id'], ['id']) 

74 

75 conn.execute(sa.sql.text(''' 

76 update view set project_id = :project_id 

77 '''), {'project_id': project_record.id}) 

78 

79 with op.batch_alter_table('view', schema=None) as batch_op: 

80 batch_op.alter_column( 

81 'project_id', 

82 existing_type=sa.INTEGER(), 

83 nullable=False 

84 ) 

85 

86 views = conn.execute(sa.text(''' 

87 select id, name from view 

88 where exists (select 1 from predictor where view.name = predictor.name) 

89 ''')).fetchall() 

90 

91 for row in views: 

92 conn.execute( 

93 text(""" 

94 update view 

95 set name = :name 

96 where id = :view_id 

97 """), { 

98 'name': f"{row['name']}_view", 

99 'view_id': row['id'] 

100 } 

101 ) 

102 

103 session.commit() 

104 

105 

106def downgrade(): 

107 conn = op.get_bind() 

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

109 

110 view_integration = db.Integration.query.filter_by(name='views').first() 

111 if view_integration is None: 

112 views_integration = db.Integration( 

113 name='views', 

114 data={}, 

115 engine='views', 

116 company_id=None 

117 ) 

118 session.add(views_integration) 

119 session.commit() 

120 

121 with op.batch_alter_table('view', schema=None) as batch_op: 

122 batch_op.drop_constraint('fk_project_id', type_='foreignkey') 

123 batch_op.drop_column('project_id') 

124 

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

126 batch_op.drop_constraint('fk_project_id', type_='foreignkey') 

127 batch_op.drop_column('project_id') 

128 

129 op.drop_table('project')