Coverage for python/lsst/dax/apdb/sql/apdbSqlReplica.py: 93%

1# This file is part of dax_apdb. 

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 

22"""Module defining Apdb class and related methods.""" 

23 

24from __future__ import annotations 

25 

26__all__ = ["ApdbSqlReplica"] 

27 

28import logging 

29from collections.abc import Collection, Iterable, Mapping, Sequence 

30from typing import TYPE_CHECKING, cast 

31 

32import astropy.time 

33import felis.datamodel 

34import numpy 

35import pandas 

36import sqlalchemy 

37from sqlalchemy import sql 

38 

39from ..apdbReplica import ApdbReplica, ApdbTableData, ReplicaChunk 

40from ..apdbSchema import ApdbTables 

41from ..apdbUpdateRecord import ApdbUpdateRecord 

42from ..monitor import MonAgent 

43from ..schema_model import Column, ExtraDataTypes 

44from ..timer import Timer 

45from ..versionTuple import VersionTuple 

46from .apdbSqlSchema import ExtraTables 

47 

48if TYPE_CHECKING: 

49 from .apdbSqlSchema import ApdbSqlSchema 

50 

51 

52_LOG = logging.getLogger(__name__) 

53 

54_MON = MonAgent(__name__) 

55 

56VERSION = VersionTuple(1, 0, 0) 

57"""Version for the code controlling replication tables. This needs to be 

58updated following compatibility rules when schema produced by this code 

59changes. 

60""" 

61 

62 

63class ApdbSqlTableData(ApdbTableData): 

64 """Implementation of ApdbTableData that wraps sqlalchemy Result. 

65 

66 Parameters 

67 ---------- 

68 result : `sqlalchemy.engine.Result` 

69 Result returned from query. 

70 column_defs : `list` [`..schema_model.Column`] 

71 Column definitions, must include all columns appearing in the 

72 ``result``. 

73 """ 

74 

75 def __init__(self, result: sqlalchemy.engine.Result, column_defs: list[Column]): 

76 column_map = {column_def.name: column_def for column_def in column_defs} 

77 self._column_defs = tuple(column_map[column_name] for column_name in result.keys()) 

78 column_types = [] 

79 for column_def in self._column_defs: 

80 if isinstance(column_def.datatype, ExtraDataTypes): 80 ↛ 81line 80 didn't jump to line 81 because the condition on line 80 was never true

81 raise TypeError("Unsupported column type {column_def.datatype} for column {column_def.name}") 

82 column_types.append((column_def.name, column_def.datatype)) 

83 self._column_types = tuple(column_types) 

84 self._rows: list[tuple] = cast(list[tuple], list(result.fetchall())) 

85 

86 def column_names(self) -> Sequence[str]: 

87 return tuple(column_def.name for column_def in self._column_defs) 

88 

89 def column_defs(self) -> Sequence[tuple[str, felis.datamodel.DataType]]: 

90 return self._column_types 

91 

92 def rows(self) -> Collection[tuple]: 

93 return self._rows 

94 

95 def to_pandas(self) -> pandas.DataFrame: 

96 """Convert data to pandas DataFrame. 

97 

98 Returns 

99 ------- 

100 dataframe : `pandas.DataFrame` 

101 Resulting DataFrame. 

102 """ 

103 if not self._rows: 

104 # There could be columns that are not in the configured schema, use 

105 # object column type for them. 

106 column_data = {} 

107 for column_def in self._column_defs: 

108 column_data[column_def.name] = pandas.Series(dtype=column_def.pandas_type) 

109 return pandas.DataFrame(column_data) 

110 

111 # To avoid nested loops convert everything to ndarray. 

112 array = numpy.array(self._rows, dtype=object) 

113 array = array.T 

114 column_data = {} 

115 for i, column_def in enumerate(self._column_defs): 

116 column_data[column_def.name] = pandas.Series(array[i], dtype=column_def.pandas_type) 

117 return pandas.DataFrame(column_data) 

118 

119 

120class ApdbSqlReplica(ApdbReplica): 

121 """Implementation of `ApdbReplica` for SQL backend. 

122 

123 Parameters 

124 ---------- 

125 schema : `ApdbSqlSchema` 

126 Instance of `ApdbSqlSchema` class for APDB database. 

127 engine : `sqlalchemy.engine.Engine` 

128 Engine for database access. 

129 db_schema_version : `VersionTuple` 

130 Version of the database schema. 

131 timer : `bool`, optional 

132 If `True` then log timing information. 

133 """ 

134 

135 def __init__( 

136 self, 

137 schema: ApdbSqlSchema, 

138 engine: sqlalchemy.engine.Engine, 

139 db_schema_version: VersionTuple, 

140 timer: bool = False, 

141 ): 

142 self._schema = schema 

143 self._engine = engine 

144 self._db_schema_version = db_schema_version 

145 

146 self._timer_args: list[MonAgent | logging.Logger] = [_MON] 

147 if timer: 147 ↛ 148line 147 didn't jump to line 148 because the condition on line 147 was never true

148 self._timer_args.append(_LOG) 

149 

150 def close(self) -> None: 

151 # docstring is inherited from a base class 

152 self._engine.dispose() 

153 

154 def _timer(self, name: str, *, tags: Mapping[str, str | int] | None = None) -> Timer: 

155 """Create `Timer` instance given its name.""" 

156 return Timer(name, *self._timer_args, tags=tags) 

157 

158 def schemaVersion(self) -> VersionTuple: 

159 # Docstring inherited from base class. 

160 return self._db_schema_version 

161 

162 @classmethod 

163 def apdbReplicaImplementationVersion(cls) -> VersionTuple: 

164 # Docstring inherited from base class. 

165 return VERSION 

166 

167 def getReplicaChunks(self) -> list[ReplicaChunk] | None: 

168 # docstring is inherited from a base class 

169 if not self._schema.replication_enabled: 

170 return None 

171 

172 table = self._schema.get_table(ExtraTables.ApdbReplicaChunks) 

173 assert table is not None, "replication_enabled=True means it must be defined" 

174 query = sql.select( 

175 table.columns["apdb_replica_chunk"], table.columns["last_update_time"], table.columns["unique_id"] 

176 ).order_by(table.columns["last_update_time"]) 

177 with self._timer("chunks_select_time") as timer: 

178 with self._engine.connect() as conn: 

179 result = conn.execution_options(stream_results=True, max_row_buffer=10000).execute(query) 

180 ids = [] 

181 for row in result: 

182 last_update_time = astropy.time.Time(row[1].timestamp(), format="unix_tai") 

183 ids.append(ReplicaChunk(id=row[0], last_update_time=last_update_time, unique_id=row[2])) 

184 timer.add_values(row_count=len(ids)) 

185 return ids 

186 

187 def deleteReplicaChunks(self, chunks: Iterable[int]) -> None: 

188 # docstring is inherited from a base class 

189 if not self._schema.replication_enabled: 189 ↛ 190line 189 didn't jump to line 190 because the condition on line 189 was never true

190 raise ValueError("APDB is not configured for replication") 

191 

192 table = self._schema.get_table(ExtraTables.ApdbReplicaChunks) 

193 chunk_list = list(chunks) 

194 where_clause = table.columns["apdb_replica_chunk"].in_(chunk_list) 

195 stmt = table.delete().where(where_clause) 

196 with self._timer("chunks_delete_time") as timer: 

197 with self._engine.begin() as conn: 

198 conn.execute(stmt) 

199 timer.add_values(row_count=len(chunk_list)) 

200 

201 def getTableDataChunks(self, table: ApdbTables, chunks: Iterable[int]) -> ApdbTableData: 

202 # docstring is inherited from a base class 

203 for chunk_table, table_enum in ExtraTables.replica_chunk_tables().items(): 

204 if table is table_enum: 

205 return self._get_chunks(chunks, table, chunk_table) 

206 raise ValueError(f"Table {table} does not support replica chunks.") 

207 

208 def _get_chunks( 

209 self, 

210 chunks: Iterable[int], 

211 table_enum: ApdbTables, 

212 chunk_table_enum: ExtraTables, 

213 ) -> ApdbTableData: 

214 """Return catalog of records for given insert identifiers, common 

215 implementation for all DIA tables. 

216 """ 

217 if not self._schema.replication_enabled: 

218 raise ValueError("APDB is not configured for replication") 

219 

220 table = self._schema.get_table(table_enum) 

221 chunk_table = self._schema.get_table(chunk_table_enum) 

222 

223 join = table.join(chunk_table) 

224 chunk_id_column = chunk_table.columns["apdb_replica_chunk"] 

225 apdb_columns = self._schema.get_apdb_columns(table_enum) 

226 where_clause = chunk_id_column.in_(chunks) 

227 query = sql.select(chunk_id_column, *apdb_columns).select_from(join).where(where_clause) 

228 

229 table_schema = self._schema.tableSchemas[table_enum] 

230 chunk_column_def = Column( 

231 name="apdb_replica_chunk", datatype=felis.datamodel.DataType.long, id="", nullable=False 

232 ) 

233 column_defs = table_schema.columns + [chunk_column_def] 

234 

235 # execute select 

236 with self._timer("table_chunk_select_time", tags={"table": table.name}) as timer: 

237 with self._engine.begin() as conn: 

238 result = conn.execution_options(stream_results=True, max_row_buffer=10000).execute(query) 

239 table_data = ApdbSqlTableData(result, column_defs) 

240 timer.add_values(row_count=len(table_data.rows())) 

241 return table_data 

242 

243 def getUpdateRecordChunks(self, chunks: Iterable[int]) -> Sequence[ApdbUpdateRecord]: 

244 # docstring is inherited from a base class 

245 if not self._schema.replication_enabled: 245 ↛ 246line 245 didn't jump to line 246 because the condition on line 245 was never true

246 raise ValueError("APDB is not configured for replication") 

247 

248 try: 

249 table = self._schema.get_table(ExtraTables.ApdbUpdateRecordChunks) 

250 except ValueError: 

251 # Table does not exist yet. 

252 return [] 

253 query = table.select().where(table.columns["apdb_replica_chunk"].in_(chunks)) 

254 

255 records = [] 

256 with self._timer("select_update_record_time", tags={"table": table.name}) as timer: 

257 with self._engine.begin() as conn: 

258 result = conn.execution_options(stream_results=True, max_row_buffer=10000).execute(query) 

259 for row in result: 

260 records.append( 

261 ApdbUpdateRecord.from_json(row.update_time_ns, row.update_order, row.update_payload) 

262 ) 

263 timer.add_values(row_count=len(records)) 

264 

265 records.sort() 

266 return records