Coverage for mindsdb / migrations / versions / 2024-11-15_9d559f68d535_add_llm_log_columns.py: 33%

1"""add_llm_log_columns 

2 

3Revision ID: 9d559f68d535 

4Revises: 6c57ed39a82b 

5Create Date: 2024-11-15 11:24:28.808881 

6 

7""" 

8from alembic import op 

9import sqlalchemy as sa 

10import mindsdb.interfaces.storage.db # noqa 

11 

12 

13# revision identifiers, used by Alembic. 

14revision = '9d559f68d535' 

15down_revision = '6c57ed39a82b' 

16branch_labels = None 

17depends_on = None 

18 

19 

20def upgrade(): 

21 with op.batch_alter_table('llm_log', schema=None) as batch_op: 

22 batch_op.alter_column('company_id', nullable=False) 

23 batch_op.alter_column('model_id', nullable=True) 

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

25 batch_op.add_column(sa.Column('cost', sa.Numeric(), nullable=True)) 

26 batch_op.add_column(sa.Column('exception', sa.String(), nullable=True)) 

27 batch_op.add_column(sa.Column('traceback', sa.String(), nullable=True)) 

28 batch_op.add_column(sa.Column('stream', sa.Boolean(), default=False)) 

29 batch_op.add_column(sa.Column('metadata', sa.JSON(), nullable=True)) 

30 

31 

32def downgrade(): 

33 with op.batch_alter_table('llm_log', schema=None) as batch_op: 

34 batch_op.alter_column('company_id', nullable=True) 

35 batch_op.alter_column('model_id', nullable=False) 

36 batch_op.drop_column('model_group') 

37 batch_op.drop_column('cost') 

38 batch_op.drop_column('exception') 

39 batch_op.drop_column('traceback') 

40 batch_op.drop_column('stream') 

41 batch_op.drop_column('metadata')