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
« 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/>.
22"""Module responsible for APDB schema operations.
23"""
25from __future__ import annotations
27__all__ = ["ApdbSqlSchema"]
29import logging
30from typing import Any, Dict, List, Mapping, Optional, Type
32import sqlalchemy
33from sqlalchemy import (Column, DDL, Index, MetaData, PrimaryKeyConstraint,
34 UniqueConstraint, Table, event)
36from .apdbSchema import ApdbSchema, ApdbTables, ColumnDef, IndexDef, IndexType
39_LOG = logging.getLogger(__name__)
42class ApdbSqlSchema(ApdbSchema):
43 """Class for management of APDB schema.
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
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 ):
85 super().__init__(schema_file, schema_name)
87 self._engine = engine
88 self._dia_object_index = dia_object_index
89 self._prefix = prefix
91 self._metadata = MetaData(self._engine, schema=namespace)
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)
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)
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)
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)
136 # generate schema for all tables, must be called last
137 self._tables = self._makeTables()
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]
145 def _makeTables(self, mysql_engine: str = 'InnoDB') -> Mapping[ApdbTables, Table]:
146 """Generate schema for all tables.
148 Parameters
149 ----------
150 mysql_engine : `str`, optional
151 MySQL engine type to use for new tables.
152 """
154 info: Dict[str, Any] = {}
156 tables = {}
157 for table_enum in ApdbTables:
159 if table_enum is ApdbTables.DiaObjectLast and self._dia_object_index != "last_object_table":
160 continue
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
172 return tables
174 def makeSchema(self, drop: bool = False, mysql_engine: str = 'InnoDB') -> None:
175 """Create or re-create all tables.
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 """
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)
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)
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()
208 def _tableColumns(self, table_name: ApdbTables) -> List[Column]:
209 """Return set of columns in a table
211 Parameters
212 ----------
213 table_name : `ApdbTables`
214 Name of the table.
216 Returns
217 -------
218 column_defs : `list`
219 List of `Column` objects.
220 """
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
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))
242 return column_defs
244 def _tableIndices(self, table_name: ApdbTables, info: Dict) -> List[sqlalchemy.schema.Constraint]:
245 """Return set of constraints/indices in a table
247 Parameters
248 ----------
249 table_name : `ApdbTables`
250 Name of the table.
251 info : `dict`
252 Additional options passed to SQLAlchemy index constructor.
254 Returns
255 -------
256 index_defs : `list`
257 List of SQLAlchemy index/constraint objects.
258 """
260 table_schema = self.tableSchemas[table_name]
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))
276 return index_defs
278 @classmethod
279 def _getDoubleType(cls, engine: sqlalchemy.engine.Engine) -> Type:
280 """DOUBLE type is database-specific, select one based on dialect.
282 Parameters
283 ----------
284 engine : `sqlalchemy.engine.Engine`
285 Database engine.
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)