Coverage for mindsdb / migrations / versions / 2024-07-19_45eb2eb61f70_add_provider_to_agent.py: 37%

1"""add provider to agent 

2 

3Revision ID: 45eb2eb61f70 

4Revises: 459a4cd24933 

5Create Date: 2024-07-19 00:48:47.629700 

6 

7""" 

8from alembic import op 

9import sqlalchemy as sa 

10from sqlalchemy.sql import table, select, update 

11import mindsdb.interfaces.storage.db # noqa 

12 

13 

14# revision identifiers, used by Alembic. 

15revision = '45eb2eb61f70' 

16down_revision = 'bfc6f44f5bc9' 

17branch_labels = None 

18depends_on = None 

19 

20 

21def upgrade(): 

22 # ### commands auto generated by Alembic - please adjust! ### 

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

24 batch_op.add_column(sa.Column('provider', sa.String(), nullable=True)) 

25 

26 # code for migrating 'provider' from 'params' to its own column 

27 agents = table('agents', 

28 sa.Column('id', sa.Integer, primary_key=True), 

29 sa.Column('params', sa.JSON), 

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

31 sa.Column('provider', sa.String())) 

32 

33 conn = op.get_bind() 

34 for agent in conn.execute(select(agents)): 

35 if agent.params and 'provider' in agent.params: 

36 conn.execute(update(agents).where(agents.c.id == agent.id).values(provider=agent.params['provider'])) 

37 if agent.model_name is None: 

38 conn.execute(update(agents).where(agents.c.id == agent.id).values(model_name=agent.params['model_name'])) 

39 # ### end Alembic commands ### 

40 

41 

42def downgrade(): 

43 # ### commands auto generated by Alembic - please adjust! ### 

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

45 batch_op.drop_column('provider') 

46 

47 # ### end Alembic commands ###