Coverage for /usr/local/lib/python3.12/site-packages/prefect/server/schemas/sorting.py: 100%

1""" 

2Schemas for sorting Prefect REST API objects. 

3""" 

4 

5from collections.abc import Iterable 

6from typing import Any 

7 

8import sqlalchemy as sa 

9 

10from prefect.utilities.collections import AutoEnum 

11 

12# TODO: Consider moving the `as_sql_sort` functions out of here since they are a 

13# database model level function and do not properly separate concerns when 

14# present in the schemas module 

15 

16 

17class FlowRunSort(AutoEnum): 

18 """Defines flow run sorting options.""" 

19 

20 ID_DESC = AutoEnum.auto() 

21 START_TIME_ASC = AutoEnum.auto() 

22 START_TIME_DESC = AutoEnum.auto() 

23 EXPECTED_START_TIME_ASC = AutoEnum.auto() 

24 EXPECTED_START_TIME_DESC = AutoEnum.auto() 

25 NAME_ASC = AutoEnum.auto() 

26 NAME_DESC = AutoEnum.auto() 

27 NEXT_SCHEDULED_START_TIME_ASC = AutoEnum.auto() 

28 END_TIME_DESC = AutoEnum.auto() 

29 

30 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

31 """Return an expression used to sort flow runs""" 

32 from prefect.server.database.dependencies import provide_database_interface 

33 

34 db = provide_database_interface() 

35 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

36 "ID_DESC": [db.FlowRun.id.desc()], 

37 "START_TIME_ASC": [ 

38 sa.func.coalesce( 

39 db.FlowRun.start_time, db.FlowRun.expected_start_time 

40 ).asc() 

41 ], 

42 "START_TIME_DESC": [ 

43 sa.func.coalesce( 

44 db.FlowRun.start_time, db.FlowRun.expected_start_time 

45 ).desc() 

46 ], 

47 "EXPECTED_START_TIME_ASC": [db.FlowRun.expected_start_time.asc()], 

48 "EXPECTED_START_TIME_DESC": [db.FlowRun.expected_start_time.desc()], 

49 "NAME_ASC": [db.FlowRun.name.asc()], 

50 "NAME_DESC": [db.FlowRun.name.desc()], 

51 "NEXT_SCHEDULED_START_TIME_ASC": [ 

52 db.FlowRun.next_scheduled_start_time.asc() 

53 ], 

54 "END_TIME_DESC": [db.FlowRun.end_time.desc()], 

55 } 

56 return sort_mapping[self.value] 

57 

58 

59class TaskRunSort(AutoEnum): 

60 """Defines task run sorting options.""" 

61 

62 ID_DESC = AutoEnum.auto() 

63 EXPECTED_START_TIME_ASC = AutoEnum.auto() 

64 EXPECTED_START_TIME_DESC = AutoEnum.auto() 

65 NAME_ASC = AutoEnum.auto() 

66 NAME_DESC = AutoEnum.auto() 

67 NEXT_SCHEDULED_START_TIME_ASC = AutoEnum.auto() 

68 END_TIME_DESC = AutoEnum.auto() 

69 

70 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

71 """Return an expression used to sort task runs""" 

72 from prefect.server.database.dependencies import provide_database_interface 

73 

74 db = provide_database_interface() 

75 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

76 "ID_DESC": [db.TaskRun.id.desc()], 

77 "EXPECTED_START_TIME_ASC": [db.TaskRun.expected_start_time.asc()], 

78 "EXPECTED_START_TIME_DESC": [db.TaskRun.expected_start_time.desc()], 

79 "NAME_ASC": [db.TaskRun.name.asc()], 

80 "NAME_DESC": [db.TaskRun.name.desc()], 

81 "NEXT_SCHEDULED_START_TIME_ASC": [ 

82 db.TaskRun.next_scheduled_start_time.asc() 

83 ], 

84 "END_TIME_DESC": [db.TaskRun.end_time.desc()], 

85 } 

86 return sort_mapping[self.value] 

87 

88 

89class LogSort(AutoEnum): 

90 """Defines log sorting options.""" 

91 

92 TIMESTAMP_ASC = AutoEnum.auto() 

93 TIMESTAMP_DESC = AutoEnum.auto() 

94 

95 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

96 """Return an expression used to sort task runs""" 

97 from prefect.server.database.dependencies import provide_database_interface 

98 

99 db = provide_database_interface() 

100 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

101 "TIMESTAMP_ASC": [db.Log.timestamp.asc()], 

102 "TIMESTAMP_DESC": [db.Log.timestamp.desc()], 

103 } 

104 return sort_mapping[self.value] 

105 

106 

107class FlowSort(AutoEnum): 

108 """Defines flow sorting options.""" 

109 

110 CREATED_DESC = AutoEnum.auto() 

111 UPDATED_DESC = AutoEnum.auto() 

112 NAME_ASC = AutoEnum.auto() 

113 NAME_DESC = AutoEnum.auto() 

114 

115 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

116 """Return an expression used to sort task runs""" 

117 from prefect.server.database.dependencies import provide_database_interface 

118 

119 db = provide_database_interface() 

120 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

121 "CREATED_DESC": [db.Flow.created.desc()], 

122 "UPDATED_DESC": [db.Flow.updated.desc()], 

123 "NAME_ASC": [db.Flow.name.asc()], 

124 "NAME_DESC": [db.Flow.name.desc()], 

125 } 

126 return sort_mapping[self.value] 

127 

128 

129class DeploymentSort(AutoEnum): 

130 """Defines deployment sorting options.""" 

131 

132 CREATED_DESC = AutoEnum.auto() 

133 UPDATED_DESC = AutoEnum.auto() 

134 NAME_ASC = AutoEnum.auto() 

135 NAME_DESC = AutoEnum.auto() 

136 

137 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

138 """Return an expression used to sort task runs""" 

139 from prefect.server.database.dependencies import provide_database_interface 

140 

141 db = provide_database_interface() 

142 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

143 "CREATED_DESC": [db.Deployment.created.desc()], 

144 "UPDATED_DESC": [db.Deployment.updated.desc()], 

145 "NAME_ASC": [db.Deployment.name.asc()], 

146 "NAME_DESC": [db.Deployment.name.desc()], 

147 } 

148 return sort_mapping[self.value] 

149 

150 

151class ArtifactSort(AutoEnum): 

152 """Defines artifact sorting options.""" 

153 

154 CREATED_DESC = AutoEnum.auto() 

155 UPDATED_DESC = AutoEnum.auto() 

156 ID_DESC = AutoEnum.auto() 

157 KEY_DESC = AutoEnum.auto() 

158 KEY_ASC = AutoEnum.auto() 

159 

160 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

161 """Return an expression used to sort task runs""" 

162 from prefect.server.database.dependencies import provide_database_interface 

163 

164 db = provide_database_interface() 

165 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

166 "CREATED_DESC": [db.Artifact.created.desc()], 

167 "UPDATED_DESC": [db.Artifact.updated.desc()], 

168 "ID_DESC": [db.Artifact.id.desc()], 

169 "KEY_DESC": [db.Artifact.key.desc()], 

170 "KEY_ASC": [db.Artifact.key.asc()], 

171 } 

172 return sort_mapping[self.value] 

173 

174 

175class ArtifactCollectionSort(AutoEnum): 

176 """Defines artifact collection sorting options.""" 

177 

178 CREATED_DESC = AutoEnum.auto() 

179 UPDATED_DESC = AutoEnum.auto() 

180 ID_DESC = AutoEnum.auto() 

181 KEY_DESC = AutoEnum.auto() 

182 KEY_ASC = AutoEnum.auto() 

183 

184 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

185 """Return an expression used to sort task runs""" 

186 from prefect.server.database.dependencies import provide_database_interface 

187 

188 db = provide_database_interface() 

189 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

190 "CREATED_DESC": [db.ArtifactCollection.created.desc()], 

191 "UPDATED_DESC": [db.ArtifactCollection.updated.desc()], 

192 "ID_DESC": [db.ArtifactCollection.id.desc()], 

193 "KEY_DESC": [db.ArtifactCollection.key.desc()], 

194 "KEY_ASC": [db.ArtifactCollection.key.asc()], 

195 } 

196 return sort_mapping[self.value] 

197 

198 

199class VariableSort(AutoEnum): 

200 """Defines variables sorting options.""" 

201 

202 CREATED_DESC = "CREATED_DESC" 

203 UPDATED_DESC = "UPDATED_DESC" 

204 NAME_DESC = "NAME_DESC" 

205 NAME_ASC = "NAME_ASC" 

206 

207 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

208 """Return an expression used to sort task runs""" 

209 from prefect.server.database.dependencies import provide_database_interface 

210 

211 db = provide_database_interface() 

212 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

213 "CREATED_DESC": [db.Variable.created.desc()], 

214 "UPDATED_DESC": [db.Variable.updated.desc()], 

215 "NAME_DESC": [db.Variable.name.desc()], 

216 "NAME_ASC": [db.Variable.name.asc()], 

217 } 

218 return sort_mapping[self.value] 

219 

220 

221class BlockDocumentSort(AutoEnum): 

222 """Defines block document sorting options.""" 

223 

224 NAME_DESC = "NAME_DESC" 

225 NAME_ASC = "NAME_ASC" 

226 BLOCK_TYPE_AND_NAME_ASC = "BLOCK_TYPE_AND_NAME_ASC" 

227 

228 def as_sql_sort(self) -> Iterable[sa.ColumnElement[Any]]: 

229 """Return an expression used to sort task runs""" 

230 from prefect.server.database.dependencies import provide_database_interface 

231 

232 db = provide_database_interface() 

233 sort_mapping: dict[str, Iterable[sa.ColumnElement[Any]]] = { 

234 "NAME_DESC": [db.BlockDocument.name.desc()], 

235 "NAME_ASC": [db.BlockDocument.name.asc()], 

236 "BLOCK_TYPE_AND_NAME_ASC": [ 

237 db.BlockDocument.block_type_name.asc(), 

238 db.BlockDocument.name.asc(), 

239 ], 

240 } 

241 return sort_mapping[self.value]