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

102 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="pixelId", 

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, "pixelId") 

117 else: 

118 # make a regular index 

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

120 type=IndexType.INDEX, columns=["pixelId"]) 

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 

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

132 """Generate schema for all tables. 

133 

134 Parameters 

135 ---------- 

136 mysql_engine : `str`, optional 

137 MySQL engine type to use for new tables. 

138 """ 

139 

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

141 

142 tables = {} 

143 for table_enum in ApdbTables: 

144 

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

146 continue 

147 

148 columns = self._tableColumns(table_enum) 

149 constraints = self._tableIndices(table_enum, info) 

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

151 self._metadata, 

152 *columns, 

153 *constraints, 

154 mysql_engine=mysql_engine, 

155 info=info) 

156 tables[table_enum] = table 

157 

158 return tables 

159 

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

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

162 

163 Parameters 

164 ---------- 

165 drop : `bool`, optional 

166 If True then drop tables before creating new ones. 

167 mysql_engine : `str`, optional 

168 MySQL engine type to use for new tables. 

169 """ 

170 

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

172 _LOG.debug("clear metadata") 

173 self._metadata.clear() 

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

175 self._makeTables(mysql_engine=mysql_engine) 

176 

177 # create all tables (optionally drop first) 

178 if drop: 

179 _LOG.info('dropping all tables') 

180 self._metadata.drop_all() 

181 _LOG.info('creating all tables') 

182 self._metadata.create_all() 

183 

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

185 """Return set of columns in a table 

186 

187 Parameters 

188 ---------- 

189 table_name : `ApdbTables` 

190 Name of the table. 

191 

192 Returns 

193 ------- 

194 column_defs : `list` 

195 List of `Column` objects. 

196 """ 

197 

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

199 # specially below 

200 table_schema = self.tableSchemas[table_name] 

201 pkey_columns = set() 

202 for index in table_schema.indices: 

203 if index.type is IndexType.PRIMARY: 

204 pkey_columns = set(index.columns) 

205 break 

206 

207 # convert all column dicts into alchemy Columns 

208 column_defs = [] 

209 for column in table_schema.columns: 

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

211 if column.default is not None: 

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

213 if column.name in pkey_columns: 

214 kwargs.update(autoincrement=False) 

215 ctype = self._type_map[column.type] 

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

217 

218 return column_defs 

219 

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

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

222 

223 Parameters 

224 ---------- 

225 table_name : `ApdbTables` 

226 Name of the table. 

227 info : `dict` 

228 Additional options passed to SQLAlchemy index constructor. 

229 

230 Returns 

231 ------- 

232 index_defs : `list` 

233 List of SQLAlchemy index/constraint objects. 

234 """ 

235 

236 table_schema = self.tableSchemas[table_name] 

237 

238 # convert all index dicts into alchemy Columns 

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

240 for index in table_schema.indices: 

241 if index.type is IndexType.INDEX: 

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

243 else: 

244 kwargs = {} 

245 if index.name: 

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

247 if index.type is IndexType.PRIMARY: 

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

249 elif index.type is IndexType.UNIQUE: 

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

251 

252 return index_defs 

253 

254 @classmethod 

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

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

257 

258 Parameters 

259 ---------- 

260 engine : `sqlalchemy.engine.Engine` 

261 Database engine. 

262 

263 Returns 

264 ------- 

265 type_object : `object` 

266 Database-specific type definition. 

267 """ 

268 if engine.name == 'mysql': 

269 from sqlalchemy.dialects.mysql import DOUBLE 

270 return DOUBLE(asdecimal=False) 

271 elif engine.name == 'postgresql': 

272 from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION 

273 return DOUBLE_PRECISION 

274 elif engine.name == 'oracle': 

275 from sqlalchemy.dialects.oracle import DOUBLE_PRECISION 

276 return DOUBLE_PRECISION 

277 elif engine.name == 'sqlite': 

278 # all floats in sqlite are 8-byte 

279 from sqlalchemy.dialects.sqlite import REAL 

280 return REAL 

281 else: 

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