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

108 statements  

« prev     ^ index     » next       coverage.py v6.5.0, created at 2022-10-12 02:38 -0700

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, DDL, Index, MetaData, PrimaryKeyConstraint, 

34 UniqueConstraint, Table, event) 

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 schema_name : `str`, optional 

68 Name of the schema in YAML files. 

69 prefix : `str`, optional 

70 Prefix to add to all scheam elements. 

71 namespace : `str`, optional 

72 Namespace (or schema name) to use for all APDB tables. 

73 """ 

74 def __init__( 

75 self, 

76 engine: sqlalchemy.engine.Engine, 

77 dia_object_index: str, 

78 htm_index_column: str, 

79 schema_file: str, 

80 schema_name: str = "ApdbSchema", 

81 prefix: str = "", 

82 namespace: Optional[str] = None, 

83 ): 

84 

85 super().__init__(schema_file, schema_name) 

86 

87 self._engine = engine 

88 self._dia_object_index = dia_object_index 

89 self._prefix = prefix 

90 

91 self._metadata = MetaData(self._engine, schema=namespace) 

92 

93 # map YAML column types to SQLAlchemy 

94 self._type_map = dict(double=self._getDoubleType(engine), 

95 float=sqlalchemy.types.Float, 

96 timestamp=sqlalchemy.types.TIMESTAMP, 

97 long=sqlalchemy.types.BigInteger, 

98 int=sqlalchemy.types.Integer, 

99 short=sqlalchemy.types.Integer, 

100 byte=sqlalchemy.types.Integer, 

101 binary=sqlalchemy.types.LargeBinary, 

102 text=sqlalchemy.types.CHAR, 

103 string=sqlalchemy.types.CHAR, 

104 char=sqlalchemy.types.CHAR, 

105 unicode=sqlalchemy.types.CHAR, 

106 boolean=sqlalchemy.types.Boolean) 

107 

108 # Adjust index if needed 

109 if self._dia_object_index == 'pix_id_iov': 

110 objects = self.tableSchemas[ApdbTables.DiaObject] 

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

112 

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

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

115 tableDef = self.tableSchemas.get(table) 

116 if not tableDef: 

117 continue 

118 column = ColumnDef(name=htm_index_column, 

119 type="long", 

120 nullable=False, 

121 default=None, 

122 description="", 

123 unit="", 

124 ucd="") 

125 tableDef.columns.append(column) 

126 

127 if table is ApdbTables.DiaObjectLast: 

128 # use it as a leading PK column 

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

130 else: 

131 # make a regular index 

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

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

134 tableDef.indices.append(index) 

135 

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

137 self._tables = self._makeTables() 

138 

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

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

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

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

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

144 

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

146 """Generate schema for all tables. 

147 

148 Parameters 

149 ---------- 

150 mysql_engine : `str`, optional 

151 MySQL engine type to use for new tables. 

152 """ 

153 

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

155 

156 tables = {} 

157 for table_enum in ApdbTables: 

158 

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

160 continue 

161 

162 columns = self._tableColumns(table_enum) 

163 constraints = self._tableIndices(table_enum, info) 

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

165 self._metadata, 

166 *columns, 

167 *constraints, 

168 mysql_engine=mysql_engine, 

169 info=info) 

170 tables[table_enum] = table 

171 

172 return tables 

173 

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

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

176 

177 Parameters 

178 ---------- 

179 drop : `bool`, optional 

180 If True then drop tables before creating new ones. 

181 mysql_engine : `str`, optional 

182 MySQL engine type to use for new tables. 

183 """ 

184 

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

186 _LOG.debug("clear metadata") 

187 self._metadata.clear() 

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

189 self._makeTables(mysql_engine=mysql_engine) 

190 

191 # Create namespace if it does not exist yet, for now this only makes 

192 # sense for postgres. 

193 if self._metadata.schema: 

194 dialect = self._engine.dialect 

195 quoted_schema = dialect.preparer(dialect).quote_schema(self._metadata.schema) 

196 create_schema = DDL( 

197 "CREATE SCHEMA IF NOT EXISTS %(schema)s", context={"schema": quoted_schema} 

198 ).execute_if(dialect='postgresql') 

199 event.listen(self._metadata, "before_create", create_schema) 

200 

201 # create all tables (optionally drop first) 

202 if drop: 

203 _LOG.info('dropping all tables') 

204 self._metadata.drop_all() 

205 _LOG.info('creating all tables') 

206 self._metadata.create_all() 

207 

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

209 """Return set of columns in a table 

210 

211 Parameters 

212 ---------- 

213 table_name : `ApdbTables` 

214 Name of the table. 

215 

216 Returns 

217 ------- 

218 column_defs : `list` 

219 List of `Column` objects. 

220 """ 

221 

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

223 # specially below 

224 table_schema = self.tableSchemas[table_name] 

225 pkey_columns = set() 

226 for index in table_schema.indices: 

227 if index.type is IndexType.PRIMARY: 

228 pkey_columns = set(index.columns) 

229 break 

230 

231 # convert all column dicts into alchemy Columns 

232 column_defs = [] 

233 for column in table_schema.columns: 

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

235 if column.default is not None: 

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

237 if column.name in pkey_columns: 

238 kwargs.update(autoincrement=False) 

239 ctype = self._type_map[column.type] 

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

241 

242 return column_defs 

243 

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

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

246 

247 Parameters 

248 ---------- 

249 table_name : `ApdbTables` 

250 Name of the table. 

251 info : `dict` 

252 Additional options passed to SQLAlchemy index constructor. 

253 

254 Returns 

255 ------- 

256 index_defs : `list` 

257 List of SQLAlchemy index/constraint objects. 

258 """ 

259 

260 table_schema = self.tableSchemas[table_name] 

261 

262 # convert all index dicts into alchemy Columns 

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

264 for index in table_schema.indices: 

265 if index.type is IndexType.INDEX: 

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

267 else: 

268 kwargs = {} 

269 if index.name: 

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

271 if index.type is IndexType.PRIMARY: 

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

273 elif index.type is IndexType.UNIQUE: 

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

275 

276 return index_defs 

277 

278 @classmethod 

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

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

281 

282 Parameters 

283 ---------- 

284 engine : `sqlalchemy.engine.Engine` 

285 Database engine. 

286 

287 Returns 

288 ------- 

289 type_object : `object` 

290 Database-specific type definition. 

291 """ 

292 if engine.name == 'mysql': 

293 from sqlalchemy.dialects.mysql import DOUBLE 

294 return DOUBLE(asdecimal=False) 

295 elif engine.name == 'postgresql': 

296 from sqlalchemy.dialects.postgresql import DOUBLE_PRECISION 

297 return DOUBLE_PRECISION 

298 elif engine.name == 'oracle': 

299 from sqlalchemy.dialects.oracle import DOUBLE_PRECISION 

300 return DOUBLE_PRECISION 

301 elif engine.name == 'sqlite': 

302 # all floats in sqlite are 8-byte 

303 from sqlalchemy.dialects.sqlite import REAL 

304 return REAL 

305 else: 

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