Coverage for mindsdb / migrations / versions / 2023-09-18_011e6f2dd9c2_backfill_agent_id.py: 30%

1"""backfill_agent_id 

2 

3Revision ID: 011e6f2dd9c2 

4Revises: f16d4ab03091 

5Create Date: 2023-09-18 11:02:36.795544 

6 

7""" 

8from alembic import op 

9import datetime 

10import sqlalchemy as sa 

11 

12 

13# revision identifiers, used by Alembic. 

14revision = '011e6f2dd9c2' 

15down_revision = 'f16d4ab03091' 

16branch_labels = None 

17depends_on = None 

18 

19 

20def upgrade(): 

21 conn = op.get_bind() 

22 chatbots_table = sa.Table( 

23 'chat_bots', 

24 sa.MetaData(), 

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

26 sa.Column('project_id', sa.Integer()), 

27 sa.Column('agent_id', sa.Integer()), 

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

29 sa.Column('model_name', sa.String()) 

30 ) 

31 

32 agents_table = sa.Table( 

33 'agents', 

34 sa.MetaData(), 

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

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

37 sa.Column('user_class', sa.Integer()), 

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

39 sa.Column('project_id', sa.Integer()), 

40 sa.Column('model_name', sa.String()), 

41 sa.Column('updated_at', sa.DateTime()), 

42 sa.Column('created_at', sa.DateTime()) 

43 ) 

44 

45 tasks_table = sa.Table( 

46 'tasks', 

47 sa.MetaData(), 

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

49 sa.Column('user_class', sa.Integer()), 

50 sa.Column('object_type', sa.String()), 

51 sa.Column('object_id', sa.Integer()) 

52 ) 

53 

54 all_chatbots = conn.execute(chatbots_table.select()).fetchall() 

55 for chatbot_row in all_chatbots: 

56 id, project_id, _, name, model_name = chatbot_row 

57 

58 # Get the corresponding task. 

59 task_select = tasks_table.select().where(tasks_table.c.object_type == 'chatbot').where(tasks_table.c.object_id == id) 

60 task_row = conn.execute(task_select).first() 

61 if task_row is None: 

62 continue 

63 company_id, user_class, _, _ = task_row 

64 # Create the new agent. 

65 op.execute(agents_table.insert().values( 

66 company_id=company_id, 

67 user_class=user_class, 

68 name=name, 

69 project_id=project_id, 

70 model_name=model_name, 

71 updated_at=datetime.datetime.now(), 

72 created_at=datetime.datetime.now() 

73 )) 

74 

75 # Get the new agent and associate the chatbot with it. 

76 created_agent = conn.execute(agents_table.select().where(agents_table.c.name == name)).first() 

77 agent_id = created_agent[0] 

78 op.execute(chatbots_table.update().where(chatbots_table.c.id == id).values(agent_id=agent_id)) 

79 

80 

81def downgrade(): 

82 pass