Coverage for python/lsst/daf/butler/core/simpleQuery.py: 30%

59 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-12 09:01 +0000

1# This file is part of daf_butler. 

2# 

3# Developed for the LSST Data Management System. 

4# This product includes software developed by the LSST Project 

5# (http://www.lsst.org). 

6# See the COPYRIGHT file at the top-level directory of this distribution 

7# for details of code ownership. 

8# 

9# This program is free software: you can redistribute it and/or modify 

10# it under the terms of the GNU General Public License as published by 

11# the Free Software Foundation, either version 3 of the License, or 

12# (at your option) any later version. 

13# 

14# This program is distributed in the hope that it will be useful, 

15# but WITHOUT ANY WARRANTY; without even the implied warranty of 

16# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the 

17# GNU General Public License for more details. 

18# 

19# You should have received a copy of the GNU General Public License 

20# along with this program. If not, see <http://www.gnu.org/licenses/>. 

21 

22from __future__ import annotations 

23 

24__all__ = ("SimpleQuery",) 

25 

26from typing import Any, ClassVar, List, Optional, Tuple, Type, TypeVar, Union 

27 

28import sqlalchemy 

29 

30T = TypeVar("T") 

31 

32 

33class SimpleQuery: 

34 """A struct that combines SQLAlchemy objects. 

35 

36 Represents SELECT, FROM, and WHERE clauses. 

37 """ 

38 

39 def __init__(self) -> None: 

40 self.columns = [] 

41 self.where = [] 

42 self.order_by = [] 

43 self.limit = None 

44 self._from: Optional[sqlalchemy.sql.FromClause] = None 

45 

46 class Select: 

47 """Tag class for SELECT queries. 

48 

49 Used to indicate that a field should be returned in 

50 a SELECT query. 

51 """ 

52 

53 Or: ClassVar[Any] 

54 

55 Select.Or = Union[T, Type[Select]] 

56 """A type annotation for arguments that can take the `Select` type or some 

57 other value. 

58 """ 

59 

60 def join( 

61 self, 

62 table: sqlalchemy.sql.FromClause, 

63 *, 

64 onclause: Optional[sqlalchemy.sql.ColumnElement] = None, 

65 isouter: bool = False, 

66 full: bool = False, 

67 **kwargs: Any, 

68 ) -> None: 

69 """Add a table or subquery join to the query. 

70 

71 Possibly also adding SELECT columns or WHERE expressions at the same 

72 time. 

73 

74 Parameters 

75 ---------- 

76 table : `sqlalchemy.sql.FromClause` 

77 Table or subquery to include. 

78 onclause : `sqlalchemy.sql.ColumnElement`, optional 

79 Expression used to join the new table or subquery to those already 

80 present. Passed directly to `sqlalchemy.sql.FromClause.join`, but 

81 ignored if this is the first call to `SimpleQuery.join`. 

82 isouter : `bool`, optional 

83 If `True`, make this an LEFT OUTER JOIN. Passed directly to 

84 `sqlalchemy.sql.FromClause.join`. 

85 full : `bool`, optional 

86 If `True`, make this a FULL OUTER JOIN. Passed directly to 

87 `sqlalchemy.sql.FromClause.join`. 

88 **kwargs 

89 Additional keyword arguments correspond to columns in the joined 

90 table or subquery. Values may be: 

91 

92 - `Select` (a special tag type) to indicate that this column 

93 should be added to the SELECT clause as a query result; 

94 - `None` to do nothing (equivalent to no keyword argument); 

95 - Any other value to add an equality constraint to the WHERE 

96 clause that constrains this column to the given value. Note 

97 that this cannot be used to add ``IS NULL`` constraints, because 

98 the previous condition for `None` is checked first. 

99 """ 

100 if self._from is None: 

101 self._from = table 

102 elif onclause is not None: 

103 self._from = self._from.join(table, onclause=onclause, isouter=isouter, full=full) 

104 else: 

105 # New table is completely unrelated to all already-included 

106 # tables. We need a cross join here but SQLAlchemy does not 

107 # have a specific method for that. Using join() without 

108 # `onclause` will try to join on FK and will raise an exception 

109 # for unrelated tables, so we have to use `onclause` which is 

110 # always true. 

111 self._from = self._from.join(table, sqlalchemy.sql.literal(True)) 

112 for name, arg in kwargs.items(): 

113 if arg is self.Select: 

114 self.columns.append(table.columns[name].label(name)) 

115 elif arg is not None: 

116 self.where.append(table.columns[name] == arg) 

117 

118 def combine(self) -> sqlalchemy.sql.Select: 

119 """Combine all terms into a single query object. 

120 

121 Returns 

122 ------- 

123 sql : `sqlalchemy.sql.Select` 

124 A SQLAlchemy object representing the full query. 

125 """ 

126 result = sqlalchemy.sql.select(*self.columns) 

127 if self._from is not None: 

128 result = result.select_from(self._from) 

129 if self.where: 

130 result = result.where(sqlalchemy.sql.and_(*self.where)) 

131 if self.order_by: 

132 result = result.order_by(*self.order_by) 

133 if self.limit: 

134 result = result.limit(self.limit[0]) 

135 if self.limit[1] is not None: 

136 result = result.offset(self.limit[1]) 

137 return result 

138 

139 @property 

140 def from_(self) -> sqlalchemy.sql.FromClause: 

141 """Return the FROM clause of the query (`sqlalchemy.sql.FromClause`). 

142 

143 This property cannot be set. To add tables to the FROM clause, call 

144 `join`. 

145 """ 

146 return self._from 

147 

148 def copy(self) -> SimpleQuery: 

149 """Return a copy of this object. 

150 

151 Returns the copy with new lists for the `where` and 

152 `columns` attributes that can be modified without changing the 

153 original. 

154 

155 Returns 

156 ------- 

157 copy : `SimpleQuery` 

158 A copy of ``self``. 

159 """ 

160 result = SimpleQuery() 

161 result.columns = list(self.columns) 

162 result.where = list(self.where) 

163 result.order_by = list(self.order_by) 

164 result.limit = self.limit 

165 result._from = self._from 

166 return result 

167 

168 columns: List[sqlalchemy.sql.ColumnElement] 

169 """The columns in the SELECT clause 

170 (`list` [ `sqlalchemy.sql.ColumnElement` ]). 

171 """ 

172 

173 where: List[sqlalchemy.sql.ColumnElement] 

174 """Boolean expressions that will be combined with AND to form the WHERE 

175 clause (`list` [ `sqlalchemy.sql.ColumnElement` ]). 

176 """ 

177 

178 order_by: List[sqlalchemy.sql.ColumnElement] 

179 """Columns to appear in ORDER BY clause (`list` 

180 [`sqlalchemy.sql.ColumnElement` ]) 

181 """ 

182 

183 limit: Optional[Tuple[int, Optional[int]]] 

184 """Limit on the number of returned rows and optional offset (`tuple`)"""