Coverage for mindsdb / migrations / versions / 2025-01-15_c06c35f7e8e1_project_company.py: 33%

1"""project-company 

2 

3Revision ID: c06c35f7e8e1 

4Revises: f6dc924079fa 

5Create Date: 2025-01-15 14:14:29.295834 

6 

7""" 

8from collections import defaultdict 

9 

10from alembic import op 

11import sqlalchemy as sa 

12import mindsdb.interfaces.storage.db # noqa 

13from mindsdb.utilities import log 

14 

15# revision identifiers, used by Alembic. 

16revision = 'c06c35f7e8e1' 

17down_revision = 'f6dc924079fa' 

18branch_labels = None 

19depends_on = None 

20 

21 

22logger = log.getLogger(__name__) 

23 

24 

25def upgrade(): 

26 

27 """ 

28 convert company_id from null to 0 to make constrain works 

29 duplicated names are renamed 

30 """ 

31 

32 conn = op.get_bind() 

33 table = sa.Table( 

34 'project', 

35 sa.MetaData(), 

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

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

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

39 ) 

40 

41 data = conn.execute( 

42 table 

43 .select() 

44 .where(table.c.company_id == sa.null()) 

45 ).fetchall() 

46 

47 names = defaultdict(list) 

48 for id, name, _ in data: 

49 names[name].append(id) 

50 

51 # get duplicated 

52 for name, ids in names.items(): 

53 if len(ids) == 1: 

54 continue 

55 

56 # rename all except first 

57 for id in ids[1:]: 

58 new_name = f'{name}__{id}' 

59 

60 op.execute( 

61 table 

62 .update() 

63 .where(table.c.id == id) 

64 .values({'name': new_name}) 

65 ) 

66 logger.warning(f'Found duplicated project name: {name}, renamed to: {new_name}') 

67 

68 op.execute( 

69 table 

70 .update() 

71 .where(table.c.company_id == sa.null()) 

72 .values({'company_id': 0}) 

73 ) 

74 

75 

76def downgrade(): 

77 table = sa.Table( 

78 'project', 

79 sa.MetaData(), 

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

81 ) 

82 

83 op.execute( 

84 table 

85 .update() 

86 .where(table.c.company_id == 0) 

87 .values({'company_id': sa.null()}) 

88 )