Coverage for python/lsst/dax/apdb/sql/modelToSql.py: 81%
120 statements
« prev ^ index » next coverage.py v7.16.2, created at 2026-11-08 01:40 -0800
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/>.
22from __future__ import annotations
24__all__ = ["GUID", "ModelToSql"]
26import uuid
27from collections.abc import Iterable, Mapping
28from typing import Any
30import felis.datamodel
31import sqlalchemy
32from sqlalchemy.dialects.postgresql import UUID
34from .. import schema_model
37#
38# Copied from daf_butler.
39#
40class GUID(sqlalchemy.TypeDecorator):
41 """Platform-independent GUID type.
43 Uses PostgreSQL's UUID type, otherwise uses CHAR(32), storing as
44 stringified hex values.
45 """
47 impl = sqlalchemy.CHAR
49 cache_ok = True
51 def load_dialect_impl(self, dialect: sqlalchemy.engine.Dialect) -> sqlalchemy.types.TypeEngine:
52 if dialect.name == "postgresql": 52 ↛ 53line 52 didn't jump to line 53 because the condition on line 52 was never true
53 return dialect.type_descriptor(UUID())
54 else:
55 return dialect.type_descriptor(sqlalchemy.CHAR(32))
57 def process_bind_param(self, value: Any, dialect: sqlalchemy.engine.Dialect) -> str | None:
58 if value is None: 58 ↛ 59line 58 didn't jump to line 59 because the condition on line 58 was never true
59 return value
61 # Coerce input to UUID type, in general having UUID on input is the
62 # only thing that we want but there is code right now that uses ints.
63 if isinstance(value, int): 63 ↛ 64line 63 didn't jump to line 64 because the condition on line 63 was never true
64 value = uuid.UUID(int=value)
65 elif isinstance(value, bytes): 65 ↛ 66line 65 didn't jump to line 66 because the condition on line 65 was never true
66 value = uuid.UUID(bytes=value)
67 elif isinstance(value, str): 67 ↛ 69line 67 didn't jump to line 69 because the condition on line 67 was never true
68 # hexstring
69 value = uuid.UUID(hex=value)
70 elif not isinstance(value, uuid.UUID): 70 ↛ 71line 70 didn't jump to line 71 because the condition on line 70 was never true
71 raise TypeError(f"Unexpected type of a bind value: {type(value)}")
73 if dialect.name == "postgresql": 73 ↛ 74line 73 didn't jump to line 74 because the condition on line 73 was never true
74 return str(value)
75 else:
76 return f"{value.int:032x}"
78 def process_result_value(
79 self, value: str | uuid.UUID | None, dialect: sqlalchemy.engine.Dialect
80 ) -> uuid.UUID | None:
81 if value is None: 81 ↛ 82line 81 didn't jump to line 82 because the condition on line 81 was never true
82 return value
83 elif isinstance(value, uuid.UUID): 83 ↛ 85line 83 didn't jump to line 85 because the condition on line 83 was never true
84 # sqlalchemy 2 converts to UUID internally
85 return value
86 else:
87 return uuid.UUID(hex=value)
90class ModelToSql:
91 """Class which implements schema model conversion to SQLAlchemy format.
93 Parameters
94 ----------
95 metadata : `sqlalchemy.schema.MetaData`
96 Metadata object for created tables.
97 prefix : `str`, optional
98 Prefix to add to all schema elements.
99 """
101 def __init__(
102 self,
103 metadata: sqlalchemy.schema.MetaData,
104 prefix: str = "",
105 ):
106 self._metadata = metadata
107 self._prefix = prefix
109 # Map model column types to SQLAlchemy.
110 self._type_map: dict[felis.datamodel.DataType | schema_model.ExtraDataTypes, type] = {
111 felis.datamodel.DataType.double: sqlalchemy.types.Double,
112 felis.datamodel.DataType.float: sqlalchemy.types.REAL,
113 felis.datamodel.DataType.timestamp: sqlalchemy.types.TIMESTAMP,
114 felis.datamodel.DataType.long: sqlalchemy.types.BigInteger,
115 felis.datamodel.DataType.int: sqlalchemy.types.Integer,
116 felis.datamodel.DataType.short: sqlalchemy.types.SmallInteger,
117 felis.datamodel.DataType.byte: sqlalchemy.types.SmallInteger, # Byte types are not very portable
118 felis.datamodel.DataType.binary: sqlalchemy.types.LargeBinary,
119 felis.datamodel.DataType.text: sqlalchemy.types.Text,
120 felis.datamodel.DataType.string: sqlalchemy.types.VARCHAR,
121 felis.datamodel.DataType.char: sqlalchemy.types.CHAR,
122 felis.datamodel.DataType.unicode: sqlalchemy.types.NVARCHAR,
123 felis.datamodel.DataType.boolean: sqlalchemy.types.Boolean,
124 schema_model.ExtraDataTypes.UUID: GUID,
125 }
127 def make_tables(self, tables: Iterable[schema_model.Table]) -> Mapping[str, sqlalchemy.schema.Table]:
128 """Generate sqlalchemy table schema from the list of modedls.
130 Parameters
131 ----------
132 tables : `~collections.abc.Iterable` [`schema_model.Table`]
133 List of table models.
135 Returns
136 -------
137 tables : `~collections.abc.Mapping` [`str`, `sqlalchemy.schema.Table`]
138 SQLAlchemy table definitions indexed by identifier of the table
139 model.
140 """
141 # Order tables based on their FK dependencies.
142 tables = self._topo_sort(tables)
144 table_map: dict[str, sqlalchemy.schema.Table] = {}
145 for table in tables:
146 columns = self._table_columns(table)
147 constraints = self._table_constraints(table, table_map)
148 sa_table = sqlalchemy.schema.Table(
149 self._prefix + table.name,
150 self._metadata,
151 *columns,
152 *constraints,
153 schema=self._metadata.schema,
154 )
155 table_map[table.id] = sa_table
157 return table_map
159 def _table_columns(self, table: schema_model.Table) -> list[sqlalchemy.schema.Column]:
160 """Return set of columns in a table
162 Parameters
163 ----------
164 table : `schema_model.Table`
165 Table model.
167 Returns
168 -------
169 column_defs : `list` [`sqlalchemy.schema.Column`]
170 List of columns.
171 """
172 column_defs: list[sqlalchemy.schema.Column] = []
173 for column in table.columns:
174 kwargs: dict[str, Any] = {"nullable": column.nullable}
175 if isinstance(column.value, bool):
176 # str(False) is stored by SQLite as text, which reads back
177 # as true.
178 kwargs.update(server_default=sqlalchemy.true() if column.value else sqlalchemy.false())
179 elif column.value is not None:
180 kwargs.update(server_default=str(column.value))
181 if column in table.primary_key and column.autoincrement is None:
182 kwargs.update(autoincrement=False)
183 else:
184 kwargs.update(autoincrement=column.autoincrement)
185 ctype: Any = self._type_map[column.datatype]
186 if column.length is not None:
187 if ctype not in (sqlalchemy.types.Text, sqlalchemy.types.TIMESTAMP): 187 ↛ 189line 187 didn't jump to line 189 because the condition on line 187 was always true
188 ctype = ctype(length=column.length)
189 if ctype is sqlalchemy.types.TIMESTAMP:
190 # Use TIMESTAMP WITH TIMEZONE.
191 ctype = ctype(timezone=True)
192 column_defs.append(sqlalchemy.schema.Column(column.name, ctype, **kwargs))
194 return column_defs
196 def _table_constraints(
197 self,
198 table: schema_model.Table,
199 table_map: Mapping[str, sqlalchemy.schema.Table],
200 ) -> list[sqlalchemy.schema.SchemaItem]:
201 """Return set of constraints/indices in a table.
203 Parameters
204 ----------
205 table : `schema_model.Table`
206 Table model.
207 table_map : `~collections.abc.Mapping`
208 MApping of table ID to sqlalchemy table definition for tables
209 that already exist, this must include all tables referenced by
210 foreign keys in ``table``.
212 Returns
213 -------
214 constraints : `list` [`sqlalchemy.schema.SchemaItem`]
215 List of SQLAlchemy index/constraint objects.
216 """
217 constraints: list[sqlalchemy.schema.SchemaItem] = []
218 if table.primary_key:
219 # It is very useful to have named PK.
220 name = self._prefix + table.name + "_pk"
221 constraints.append(
222 sqlalchemy.schema.PrimaryKeyConstraint(*[column.name for column in table.primary_key])
223 )
224 for index in table.indexes:
225 if index.expressions: 225 ↛ 226line 225 didn't jump to line 226 because the condition on line 225 was never true
226 raise TypeError(f"Expression indices are not supported: {table}")
227 name = self._prefix + index.name if index.name else ""
228 constraints.append(sqlalchemy.schema.Index(name, *[column.name for column in index.columns]))
229 for constraint in table.constraints:
230 constr_name: str | None = None
231 if constraint.name: 231 ↛ 233line 231 didn't jump to line 233 because the condition on line 231 was always true
232 constr_name = self._prefix + constraint.name
233 if isinstance(constraint, schema_model.UniqueConstraint): 233 ↛ 234line 233 didn't jump to line 234 because the condition on line 233 was never true
234 constraints.append(
235 sqlalchemy.schema.UniqueConstraint(
236 *[column.name for column in constraint.columns], name=constr_name
237 )
238 )
239 elif isinstance(constraint, schema_model.ForeignKeyConstraint): 239 ↛ 254line 239 didn't jump to line 254 because the condition on line 239 was always true
240 column_names = [col.name for col in constraint.columns]
241 foreign_table = table_map[constraint.referenced_table.id]
242 refcolumns = [foreign_table.columns[col.name] for col in constraint.referenced_columns]
243 constraints.append(
244 sqlalchemy.schema.ForeignKeyConstraint(
245 columns=column_names,
246 refcolumns=refcolumns,
247 name=constr_name,
248 deferrable=constraint.deferrable,
249 initially=constraint.initially,
250 onupdate=constraint.onupdate,
251 ondelete=constraint.ondelete,
252 )
253 )
254 elif isinstance(constraint, schema_model.CheckConstraint):
255 constraints.append(
256 sqlalchemy.schema.CheckConstraint(
257 constraint.expression,
258 name=constr_name,
259 deferrable=constraint.deferrable,
260 initially=constraint.initially,
261 )
262 )
263 else:
264 raise TypeError(f"Unknown constraint type: {constraint}")
266 return constraints
268 @staticmethod
269 def _topo_sort(table_iter: Iterable[schema_model.Table]) -> list[schema_model.Table]:
270 """Toplogical sorting of tables."""
271 result: list[schema_model.Table] = []
272 result_ids: set[str] = set()
273 tables = list(table_iter)
275 # Map of table ID to foreign table IDs.
276 referenced_tables: dict[str, set[str]] = {}
277 for table in tables:
278 referenced_tables[table.id] = set()
279 for constraint in table.constraints:
280 if isinstance(constraint, schema_model.ForeignKeyConstraint): 280 ↛ 279line 280 didn't jump to line 279 because the condition on line 280 was always true
281 referenced_tables[table.id].add(constraint.referenced_table.id)
283 while True:
284 keep = []
285 changed = False
286 for table in tables:
287 if referenced_tables[table.id].issubset(result_ids): 287 ↛ 292line 287 didn't jump to line 292 because the condition on line 287 was always true
288 changed = True
289 result.append(table)
290 result_ids.add(table.id)
291 else:
292 keep.append(table)
293 tables = keep
294 if not changed:
295 break
297 # If nothing can be removed it means cycle.
298 if tables: 298 ↛ 299line 298 didn't jump to line 299 because the condition on line 298 was never true
299 raise ValueError(f"Dependency cycle in foreign keys: {tables}")
301 return result