Coverage for python/lsst/dax/apdb/apdbSqlSchema.py: 13%

Shortcuts on this page

r m x p   toggle line displays

j k   next/prev highlighted chunk

0   (zero) top of page

1   (one) first highlighted chunk

103 statements  

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 responsible for APDB schema operations. 

23""" 

24 

25from __future__ import annotations 

26 

27__all__ = ["ApdbSqlSchema"] 

28 

29import logging 

30from typing import Any, Dict, List, Mapping, Optional, Type 

31 

32import sqlalchemy 

33from sqlalchemy import (Column, Index, MetaData, PrimaryKeyConstraint, 

34 UniqueConstraint, Table) 

35 

36from .apdbSchema import ApdbSchema, ApdbTables, ColumnDef, IndexDef, IndexType 

37 

38 

39_LOG = logging.getLogger(__name__) 

40 

41 

42class ApdbSqlSchema(ApdbSchema): 

43 """Class for management of APDB schema. 

44 

45 Attributes 

46 ---------- 

47 objects : `sqlalchemy.Table` 

48 DiaObject table instance 

49 objects_last : `sqlalchemy.Table` 

50 DiaObjectLast table instance, may be None 

51 sources : `sqlalchemy.Table` 

52 DiaSource table instance 

53 forcedSources : `sqlalchemy.Table` 

54 DiaForcedSource table instance 

55 

56 Parameters 

57 ---------- 

58 engine : `sqlalchemy.engine.Engine` 

59 SQLAlchemy engine instance 

60 dia_object_index : `str` 

61 Indexing mode for DiaObject table, see `ApdbSqlConfig.dia_object_index` 

62 for details. 

63 htm_index_column : `str` 

64 Name of a HTM index column for DiaObject and DiaSource tables. 

65 schema_file : `str` 

66 Name of the YAML schema file. 

67 extra_schema_file : `str`, optional 

68 Name of the YAML schema file with extra column definitions. 

69 prefix : `str`, optional 

70 Prefix to add to all scheam elements. 

71 """ 

72 def __init__(self, engine: sqlalchemy.engine.Engine, dia_object_index: str, htm_index_column: str, 

73 schema_file: str, extra_schema_file: Optional[str] = None, prefix: str = ""): 

74 

75 super().__init__(schema_file, extra_schema_file) 

76 

77 self._engine = engine 

78 self._dia_object_index = dia_object_index 

79 self._prefix = prefix 

80 

81 self._metadata = MetaData(self._engine) 

82 

83 # map YAML column types to SQLAlchemy 

84 self._type_map = dict(DOUBLE=self._getDoubleType(engine), 

85 FLOAT=sqlalchemy.types.Float, 

86 DATETIME=sqlalchemy.types.TIMESTAMP, 

87 BIGINT=sqlalchemy.types.BigInteger, 

88 INTEGER=sqlalchemy.types.Integer, 

89 INT=sqlalchemy.types.Integer, 

90 TINYINT=sqlalchemy.types.Integer, 

91 BLOB=sqlalchemy.types.LargeBinary, 

92 CHAR=sqlalchemy.types.CHAR, 

93 BOOL=sqlalchemy.types.Boolean) 

94 

95 # Adjust index if needed 

96 if self._dia_object_index == 'pix_id_iov': 

97 objects = self.tableSchemas[ApdbTables.DiaObject] 

98 objects.primary_key.columns.insert(0, htm_index_column) 

99 

100 # Add pixelId column and index to tables that need it 

101 for table in (ApdbTables.DiaObject, ApdbTables.DiaObjectLast, ApdbTables.DiaSource): 

102 tableDef = self.tableSchemas.get(table) 

103 if not tableDef: 

104 continue 

105 column = ColumnDef(name=htm_index_column, 

106 type="BIGINT", 

107 nullable=False, 

108 default=None, 

109 description="", 

110 unit="", 

111 ucd="") 

112 tableDef.columns.append(column) 

113 

114 if table is ApdbTables.DiaObjectLast: 

115 # use it as a leading PK column 

116 tableDef.primary_key.columns.insert(0, htm_index_column) 

117 else: 

118 # make a regular index 

119 index = IndexDef(name=f"IDX_{tableDef.name}_{htm_index_column}", 

120 type=IndexType.INDEX, columns=[htm_index_column]) 

121 tableDef.indices.append(index) 

122 

123 # generate schema for all tables, must be called last 

124 self._tables = self._makeTables() 

125 

126 self.objects = self._tables[ApdbTables.DiaObject] 

127 self.objects_last = self._tables.get(ApdbTables.DiaObjectLast) 

128 self.sources = self._tables[ApdbTables.DiaSource] 

129 self.forcedSources = self._tables[ApdbTables.DiaForcedSource] 

130 self.ssObjects = self._tables[ApdbTables.SSObject] 

131 

132 def _makeTables(self, mysql_engine: str = 'InnoDB') -> Mapping[ApdbTables, Table]: 

133 """Generate schema for all tables. 

134 

135 Parameters 

136 ---------- 

137 mysql_engine : `str`, optional 

138 MySQL engine type to use for new tables. 

139 """ 

140 

141 info: Dict[str, Any] = {} 

142 

143 tables = {} 

144 for table_enum in ApdbTables: 

145 

146 if table_enum is ApdbTables.DiaObjectLast and self._dia_object_index != "last_object_table": 

147 continue 

148 

149 columns = self._tableColumns(table_enum) 

150 constraints = self._tableIndices(table_enum, info) 

151 table = Table(table_enum.table_name(self._prefix), 

152 self._metadata, 

153 *columns, 

154 *constraints, 

155 mysql_engine=mysql_engine, 

156 info=info) 

157 tables[table_enum] = table 

158 

159 return tables 

160 

161 def makeSchema(self, drop: bool = False, mysql_engine: str = 'InnoDB') -> None: 

162 """Create or re-create all tables. 

163 

164 Parameters 

165 ---------- 

166 drop : `bool`, optional 

167 If True then drop tables before creating new ones. 

168 mysql_engine : `str`, optional 

169 MySQL engine type to use for new tables. 

170 """ 

171 

172 # re-make table schema for all needed tables with possibly different options 

173 _LOG.debug("clear metadata") 

174 self._metadata.clear() 

175 _LOG.debug("re-do schema mysql_engine=%r", mysql_engine) 

176 self._makeTables(mysql_engine=mysql_engine) 

177 

178 # create all tables (optionally drop first) 

179 if drop: 

180 _LOG.info('dropping all tables') 

181 self._metadata.drop_all() 

182 _LOG.info('creating all tables') 

183 self._metadata.create_all() 

184 

185 def _tableColumns(self, table_name: ApdbTables) -> List[Column]: 

186 """Return set of columns in a table 

187 

188 Parameters 

189 ---------- 

190 table_name : `ApdbTables` 

191 Name of the table. 

192 

193 Returns 

194 ------- 

195 column_defs : `list` 

196 List of `Column` objects. 

197 """ 

198 

199 # get the list of columns in primary key, they are treated somewhat 

200 # specially below 

201 table_schema = self.tableSchemas[table_name] 

202 pkey_columns = set() 

203 for index in table_schema.indices: 

204 if index.type is IndexType.PRIMARY: 

205 pkey_columns = set(index.columns) 

206 break 

207 

208 # convert all column dicts into alchemy Columns 

209 column_defs = [] 

210 for column in table_schema.columns: 

211 kwargs: Dict[str, Any] = dict(nullable=column.nullable) 

212 if column.default is not None: 

213 kwargs.update(server_default=str(column.default)) 

214 if column.name in pkey_columns: 

215 kwargs.update(autoincrement=False) 

216 ctype = self._type_map[column.type] 

217 column_defs.append(Column(column.name, ctype, **kwargs)) 

218 

219 return column_defs 

220 

221 def _tableIndices(self, table_name: ApdbTables, info: Dict) -> List[sqlalchemy.schema.Constraint]: 

222 """Return set of constraints/indices in a table 

223 

224 Parameters 

225 ---------- 

226 table_name : `ApdbTables` 

227 Name of the table. 

228 info : `dict` 

229 Additional options passed to SQLAlchemy index constructor. 

230 

231 Returns 

232 ------- 

233 index_defs : `list` 

234 List of SQLAlchemy index/constraint objects. 

235 """ 

236 

237 table_schema = self.tableSchemas[table_name] 

238 

239 # convert all index dicts into alchemy Columns 

240 index_defs: List[sqlalchemy.schema.Constraint] = [] 

241 for index in table_schema.indices: 

242 if index.type is IndexType.INDEX: 

243 index_defs.append(Index(self._prefix + index.name, *index.columns, info=info)) 

244 else: 

245 kwargs = {} 

246 if index.name: 

247 kwargs['name'] = self._prefix + index.name 

248 if index.type is IndexType.PRIMARY: 

249 index_defs.append(PrimaryKeyConstraint(*index.columns, **kwargs)) 

250 elif index.type is IndexType.UNIQUE: 

251 index_defs.append(UniqueConstraint(*index.columns, **kwargs)) 

252 

253 return index_defs 

254 

255 @classmethod 

256 def _getDoubleType(cls, engine: sqlalchemy.engine.Engine) -> Type: 

257 """DOUBLE type is database-specific, select one based on dialect. 

258 

259 Parameters 

260 ---------- 

261 engine : `sqlalchemy.engine.Engine` 

262 Database engine. 

263 

264 Returns 

265 ------- 

266 type_object : `object` 

267 Database-specific type definition. 

268 """ 

269 if engine.name == 'mysql': 

270 from sqlalchemy.dialects.mysql import DOUBLE 

271 return DOUBLE(asdecimal=False) 

272 elif engine.name == 'postgresql': 

273 from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION 

274 return DOUBLE_PRECISION 

275 elif engine.name == 'oracle': 

276 from sqlalchemy.dialects.oracle import DOUBLE_PRECISION 

277 return DOUBLE_PRECISION 

278 elif engine.name == 'sqlite': 

279 # all floats in sqlite are 8-byte 

280 from sqlalchemy.dialects.sqlite import REAL 

281 return REAL 

282 else: 

283 raise TypeError('cannot determine DOUBLE type, unexpected dialect: ' + engine.name)