Coverage for mindsdb / migrations / versions / 2022-09-29_cada7d2be947_json_storage.py: 23%

1"""json_storage 

2 

3Revision ID: cada7d2be947 

4Revises: 3d5e70105df7 

5Create Date: 2022-09-29 15:52:32.695026 

6 

7""" 

8import json 

9 

10from alembic import op 

11import sqlalchemy as sa 

12 

13import mindsdb.interfaces.storage.db as db 

14from mindsdb.interfaces.storage.fs import RESOURCE_GROUP 

15 

16 

17# revision identifiers, used by Alembic. 

18revision = 'cada7d2be947' 

19down_revision = '3d5e70105df7' 

20branch_labels = None 

21depends_on = None 

22 

23 

24def upgrade(): 

25 op.create_table( 

26 'json_storage', 

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

28 sa.Column('resource_group', sa.String(), nullable=True), 

29 sa.Column('resource_id', sa.Integer(), nullable=True), 

30 sa.Column('name', sa.String(), nullable=True), 

31 sa.Column('content', sa.JSON(), nullable=True), 

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

33 sa.PrimaryKeyConstraint('id') 

34 ) 

35 

36 conn = op.get_bind() 

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

38 predictors = conn.execute(sa.text(''' 

39 select id, json_ai from predictor 

40 ''')).fetchall() 

41 

42 for row in predictors: 

43 try: 

44 jai = json.loads(row['json_ai']) 

45 except Exception: 

46 continue 

47 

48 if jai is None or len(jai) == 0: 

49 continue 

50 

51 record = db.JsonStorage( 

52 resource_group=RESOURCE_GROUP.PREDICTOR, 

53 resource_id=row['id'], 

54 name='json_ai', 

55 content=jai, 

56 company_id=None 

57 ) 

58 session.add(record) 

59 

60 session.commit() 

61 

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

63 batch_op.drop_column('json_ai') 

64 

65 

66def downgrade(): 

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

68 batch_op.add_column(sa.Column('json_ai', sa.JSON, nullable=True)) 

69 

70 conn = op.get_bind() 

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

72 jsons = conn.execute(sa.text(''' 

73 select resource_id, name, content 

74 from json_storage 

75 where resource_group = 'predictor' and name = 'json_ai' 

76 ''')).fetchall() 

77 

78 for row in jsons: 

79 predicrtor_record = ( 

80 session.query(db.Predictor) 

81 .filter_by(company_id=None, id=row['resource_id']) 

82 .first() 

83 ) 

84 if predicrtor_record is None: 

85 continue 

86 predicrtor_record.json_ai = row['content'] 

87 

88 session.commit() 

89 

90 op.drop_table('json_storage')