Coverage for mlos_bench/mlos_bench/storage/sql/alembic/versions/b61aa446e724_support_fractional_seconds_with_mysql.py: 37%

1# 

2# Copyright (c) Microsoft Corporation. 

3# Licensed under the MIT License. 

4# 

5""" 

6Support fractional seconds with MySQL. 

7 

8Revision ID: b61aa446e724 

9Revises: 8928a401115b 

10Create Date: 2025-06-02 17:56:34.746642+00:00 

11""" 

12# pylint: disable=no-member 

13 

14from collections.abc import Sequence 

15 

16import sqlalchemy as sa 

17from alembic import context, op 

18from sqlalchemy.dialects import mysql 

19 

20# revision identifiers, used by Alembic. 

21revision: str = "b61aa446e724" 

22down_revision: str | None = "8928a401115b" 

23branch_labels: str | Sequence[str] | None = None 

24depends_on: str | Sequence[str] | None = None 

25 

26 

27def _mysql_datetime(*, with_fsp: bool = False) -> mysql.DATETIME: 

28 """ 

29 Return a MySQL DATETIME type with fractional seconds precision (fsp=6). 

30 

31 Notes 

32 ----- 

33 Split out to allow single mypy ignore. 

34 See <https://github.com/sqlalchemy/sqlalchemy/pull/12164> for details. 

35 """ 

36 if with_fsp: 

37 return mysql.DATETIME(fsp=6) # type: ignore[no-untyped-call] 

38 return mysql.DATETIME() # type: ignore[no-untyped-call] 

39 

40 

41def upgrade() -> None: 

42 """The schema upgrade script for this revision.""" 

43 bind = context.get_bind() 

44 if bind.dialect.name == "mysql": 

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

46 op.alter_column( 

47 "experiment", 

48 "ts_start", 

49 existing_type=_mysql_datetime(), 

50 type_=sa.DateTime(timezone=True).with_variant( 

51 _mysql_datetime(with_fsp=True), 

52 "mysql", 

53 ), 

54 existing_nullable=True, 

55 ) 

56 op.alter_column( 

57 "experiment", 

58 "ts_end", 

59 existing_type=_mysql_datetime(), 

60 type_=sa.DateTime(timezone=True).with_variant( 

61 _mysql_datetime(with_fsp=True), 

62 "mysql", 

63 ), 

64 existing_nullable=True, 

65 ) 

66 op.alter_column( 

67 "trial", 

68 "ts_start", 

69 existing_type=_mysql_datetime(), 

70 type_=sa.DateTime(timezone=True).with_variant( 

71 _mysql_datetime(with_fsp=True), 

72 "mysql", 

73 ), 

74 existing_nullable=False, 

75 ) 

76 op.alter_column( 

77 "trial", 

78 "ts_end", 

79 existing_type=_mysql_datetime(), 

80 type_=sa.DateTime(timezone=True).with_variant( 

81 _mysql_datetime(with_fsp=True), 

82 "mysql", 

83 ), 

84 existing_nullable=True, 

85 ) 

86 op.alter_column( 

87 "trial_status", 

88 "ts", 

89 existing_type=_mysql_datetime(), 

90 type_=sa.DateTime(timezone=True).with_variant( 

91 _mysql_datetime(with_fsp=True), 

92 "mysql", 

93 ), 

94 existing_nullable=False, 

95 ) 

96 op.alter_column( 

97 "trial_telemetry", 

98 "ts", 

99 existing_type=_mysql_datetime(), 

100 type_=sa.DateTime(timezone=True).with_variant( 

101 _mysql_datetime(with_fsp=True), 

102 "mysql", 

103 ), 

104 existing_nullable=False, 

105 ) 

106 # ### end Alembic commands ### 

107 

108 

109def downgrade() -> None: 

110 """The schema downgrade script for this revision.""" 

111 bind = context.get_bind() 

112 if bind.dialect.name == "mysql": 

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

114 op.alter_column( 

115 "trial_telemetry", 

116 "ts", 

117 existing_type=sa.DateTime(timezone=True).with_variant( 

118 _mysql_datetime(with_fsp=True), 

119 "mysql", 

120 ), 

121 type_=_mysql_datetime(), 

122 existing_nullable=False, 

123 ) 

124 op.alter_column( 

125 "trial_status", 

126 "ts", 

127 existing_type=sa.DateTime(timezone=True).with_variant( 

128 _mysql_datetime(with_fsp=True), 

129 "mysql", 

130 ), 

131 type_=_mysql_datetime(), 

132 existing_nullable=False, 

133 ) 

134 op.alter_column( 

135 "trial", 

136 "ts_end", 

137 existing_type=sa.DateTime(timezone=True).with_variant( 

138 _mysql_datetime(with_fsp=True), 

139 "mysql", 

140 ), 

141 type_=_mysql_datetime(), 

142 existing_nullable=True, 

143 ) 

144 op.alter_column( 

145 "trial", 

146 "ts_start", 

147 existing_type=sa.DateTime(timezone=True).with_variant( 

148 _mysql_datetime(with_fsp=True), 

149 "mysql", 

150 ), 

151 type_=_mysql_datetime(), 

152 existing_nullable=False, 

153 ) 

154 op.alter_column( 

155 "experiment", 

156 "ts_end", 

157 existing_type=sa.DateTime(timezone=True).with_variant( 

158 _mysql_datetime(with_fsp=True), 

159 "mysql", 

160 ), 

161 type_=_mysql_datetime(), 

162 existing_nullable=True, 

163 ) 

164 op.alter_column( 

165 "experiment", 

166 "ts_start", 

167 existing_type=sa.DateTime(timezone=True).with_variant( 

168 _mysql_datetime(with_fsp=True), 

169 "mysql", 

170 ), 

171 type_=_mysql_datetime(), 

172 existing_nullable=True, 

173 ) 

174 # ### end Alembic commands ###