Coverage for python/felis/metadata.py: 16%
124 statements
« prev ^ index » next coverage.py v7.10.6, created at 2025-10-09 03:24 -0700
1"""Build SQLAlchemy metadata from a Felis schema."""
3# This file is part of felis.
4#
5# Developed for the LSST Data Management System.
6# This product includes software developed by the LSST Project
7# (https://www.lsst.org).
8# See the COPYRIGHT file at the top-level directory of this distribution
9# for details of code ownership.
10#
11# This program is free software: you can redistribute it and/or modify
12# it under the terms of the GNU General Public License as published by
13# the Free Software Foundation, either version 3 of the License, or
14# (at your option) any later version.
15#
16# This program is distributed in the hope that it will be useful,
17# but WITHOUT ANY WARRANTY; without even the implied warranty of
18# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
19# GNU General Public License for more details.
20#
21# You should have received a copy of the GNU General Public License
22# along with this program. If not, see <https://www.gnu.org/licenses/>.
24from __future__ import annotations
26import logging
27from typing import Any, Literal
29from lsst.utils.iteration import ensure_iterable
30from sqlalchemy import (
31 CheckConstraint,
32 Column,
33 Constraint,
34 ForeignKeyConstraint,
35 Index,
36 MetaData,
37 PrimaryKeyConstraint,
38 Table,
39 TextClause,
40 UniqueConstraint,
41 text,
42)
43from sqlalchemy.dialects import mysql, postgresql
44from sqlalchemy.types import TypeEngine
46from felis.datamodel import Schema
47from felis.db.variants import make_variant_dict
49from . import datamodel
50from .db import sqltypes
51from .types import FelisType
53__all__ = ("MetaDataBuilder", "get_datatype_with_variants")
55logger = logging.getLogger(__name__)
58def _handle_timestamp_column(column_obj: datamodel.Column, variant_dict: dict[str, TypeEngine[Any]]) -> None:
59 """Handle columns with the timestamp datatype.
61 Parameters
62 ----------
63 column_obj
64 The column object representing the timestamp.
65 variant_dict
66 The dictionary of variant overrides for the datatype.
68 Notes
69 -----
70 This function updates the variant dictionary with the appropriate
71 timestamp type for the column object but only if the precision is set.
72 Otherwise, the default timestamp objects defined in the Felis type system
73 will be used instead.
74 """
75 if column_obj.precision is not None:
76 args: Any = [False, column_obj.precision] # Turn off timezone.
77 variant_dict.update({"postgresql": postgresql.TIMESTAMP(*args), "mysql": mysql.DATETIME(*args)})
80def get_datatype_with_variants(column_obj: datamodel.Column) -> TypeEngine:
81 """Use the Felis type system to get a SQLAlchemy datatype with variant
82 overrides from the information in a Felis column object.
84 Parameters
85 ----------
86 column_obj
87 The column object from which to get the datatype.
89 Returns
90 -------
91 `~sqlalchemy.types.TypeEngine`
92 The SQLAlchemy datatype object.
94 Raises
95 ------
96 ValueError
97 Raised if the column has a sized type but no length or if the datatype
98 is invalid.
99 """
100 variant_dict = make_variant_dict(column_obj)
101 felis_type = FelisType.felis_type(column_obj.datatype.value)
102 datatype_fun = getattr(sqltypes, column_obj.datatype.value, None)
103 if datatype_fun is None:
104 raise ValueError(f"Unknown datatype: {column_obj.datatype.value}")
105 args = []
106 if felis_type.is_sized:
107 # Add length argument for size types.
108 if not column_obj.length:
109 raise ValueError(f"Column {column_obj.name} has sized type '{column_obj.datatype}' but no length")
110 args = [column_obj.length]
111 if felis_type.is_timestamp:
112 _handle_timestamp_column(column_obj, variant_dict)
113 return datatype_fun(*args, **variant_dict)
116_VALID_SERVER_DEFAULTS = ("CURRENT_TIMESTAMP", "NOW()", "LOCALTIMESTAMP", "NULL")
119class MetaDataBuilder:
120 """Build a SQLAlchemy metadata object from a Felis schema.
122 Parameters
123 ----------
124 schema
125 The schema object from which to build the SQLAlchemy metadata.
126 apply_schema_to_metadata
127 Whether to apply the schema name to the metadata object.
128 ignore_constraints
129 Whether to ignore constraints when building the metadata.
130 table_name_postfix
131 A string to append to the table names when building the metadata.
132 """
134 def __init__(
135 self,
136 schema: Schema,
137 apply_schema_to_metadata: bool = True,
138 ignore_constraints: bool = False,
139 table_name_postfix: str = "",
140 ) -> None:
141 """Initialize the metadata builder."""
142 self.schema = schema
143 if not apply_schema_to_metadata:
144 logger.debug("Schema name will not be applied to metadata")
145 self.metadata = MetaData(schema=schema.name if apply_schema_to_metadata else None)
146 self._objects: dict[str, Any] = {}
147 self.ignore_constraints = ignore_constraints
148 self.table_name_postfix = table_name_postfix
150 def build(self) -> MetaData:
151 """Build the SQLAlchemy tables and constraints from the schema.
153 Notes
154 -----
155 This first builds the tables and then makes a second pass to build the
156 constraints. This is necessary because the constraints may reference
157 objects that are not yet created when the tables are built.
159 Returns
160 -------
161 `~sqlalchemy.sql.schema.MetaData`
162 The SQLAlchemy metadata object.
163 """
164 self.build_tables()
165 if not self.ignore_constraints:
166 self.build_constraints()
167 else:
168 logger.warning("Ignoring constraints")
169 return self.metadata
171 def build_tables(self) -> None:
172 """Build the SQLAlchemy tables from the schema."""
173 for table in self.schema.tables:
174 self.build_table(table)
175 if table.primary_key:
176 primary_key = self.build_primary_key(table.primary_key)
177 self._objects[table.id].append_constraint(primary_key)
179 def build_primary_key(self, primary_key_columns: str | list[str]) -> PrimaryKeyConstraint:
180 """Build a SQAlchemy ``PrimaryKeyConstraint`` from a single column ID
181 or a list of them.
183 Parameters
184 ----------
185 primary_key_columns
186 The column ID or list of column IDs from which to build the primary
187 key.
189 Returns
190 -------
191 `~sqlalchemy.sql.schema.PrimaryKeyConstraint`
192 The SQLAlchemy primary key constraint object.
194 Notes
195 -----
196 The ``primary_key_columns`` is a string or a list of strings
197 representing IDs which will be used to find the columnn objects in the
198 builder's internal ID map.
199 """
200 return PrimaryKeyConstraint(
201 *[self._objects[column_id] for column_id in ensure_iterable(primary_key_columns)]
202 )
204 def build_table(self, table_obj: datamodel.Table) -> None:
205 """Build a SQLAlchemy ``Table`` from a Felis table and add it to the
206 metadata.
208 Parameters
209 ----------
210 table_obj
211 The Felis table object from which to build the SQLAlchemy table.
213 Notes
214 -----
215 Several MySQL table options, including the engine and charset, are
216 handled by adding annotations to the table. This is not needed for
217 Postgres, as Felis does not support any table options for this dialect.
218 """
219 # Process mysql table options.
220 optargs = {}
221 if table_obj.mysql_engine:
222 optargs["mysql_engine"] = table_obj.mysql_engine
223 if table_obj.mysql_charset:
224 optargs["mysql_charset"] = table_obj.mysql_charset
226 # Create the SQLAlchemy table object and its columns.
227 name = table_obj.name
228 id = table_obj.id
229 description = table_obj.description
230 columns = [self.build_column(column) for column in table_obj.columns]
231 table = Table(
232 name + self.table_name_postfix,
233 self.metadata,
234 *columns,
235 comment=description,
236 **optargs, # type: ignore[arg-type]
237 )
239 # Create the indexes and add them to the table.
240 indexes = [self.build_index(index) for index in table_obj.indexes]
241 for index in indexes:
242 index._set_parent(table)
243 table.indexes.add(index)
245 self._objects[id] = table
247 def build_column(self, column_obj: datamodel.Column) -> Column:
248 """Build a SQLAlchemy ``Column`` from a Felis column object.
250 Parameters
251 ----------
252 column_obj
253 The column object from which to build the SQLAlchemy column.
255 Returns
256 -------
257 `~sqlalchemy.sql.schema.Column`
258 The SQLAlchemy column object.
259 """
260 # Get basic column attributes.
261 name = column_obj.name
262 id = column_obj.id
263 description = column_obj.description
264 value = column_obj.value
265 nullable = column_obj.nullable
267 # Get datatype, handling variant overrides such as "mysql:datatype".
268 datatype = get_datatype_with_variants(column_obj)
270 # Set autoincrement, depending on if it was provided explicitly.
271 autoincrement: Literal["auto"] | bool = (
272 column_obj.autoincrement if column_obj.autoincrement is not None else "auto"
273 )
275 server_default: str | TextClause | None = None
276 if value is not None:
277 server_default = str(value)
278 if server_default in _VALID_SERVER_DEFAULTS or not isinstance(value, str):
279 # If the server default is a valid keyword or not a string,
280 # use it as is.
281 server_default = text(server_default)
283 if server_default is not None:
284 logger.debug(f"Column '{id}' has default value: {server_default}")
286 column: Column = Column(
287 name,
288 datatype,
289 comment=description,
290 autoincrement=autoincrement,
291 nullable=nullable,
292 server_default=server_default,
293 )
295 self._objects[id] = column
297 return column
299 def build_constraints(self) -> None:
300 """Build the SQLAlchemy constraints from the Felis schema and append
301 them to the associated table in the metadata.
303 Notes
304 -----
305 This is performed as a separate step after building the tables so that
306 all the referenced objects in the constraints will be present and can
307 be looked up by their ID.
308 """
309 for table_obj in self.schema.tables:
310 table = self._objects[table_obj.id]
311 for constraint_obj in table_obj.constraints:
312 constraint = self.build_constraint(constraint_obj)
313 table.append_constraint(constraint)
315 def build_constraint(self, constraint_obj: datamodel.Constraint) -> Constraint:
316 """Build a SQLAlchemy ``Constraint`` from a Felis constraint.
318 Parameters
319 ----------
320 constraint_obj
321 The Felis object from which to build the constraint.
323 Returns
324 -------
325 `~sqlalchemy.sql.schema.Constraint`
326 The SQLAlchemy constraint object.
328 Raises
329 ------
330 ValueError
331 If the constraint type is not recognized.
332 TypeError
333 If the constraint object is not the expected type.
334 """
335 args: dict[str, Any] = {
336 "name": constraint_obj.name or None,
337 "comment": constraint_obj.description or None,
338 "deferrable": constraint_obj.deferrable or None,
339 "initially": constraint_obj.initially or None,
340 }
342 constraint: Constraint
344 if isinstance(constraint_obj, datamodel.ForeignKeyConstraint):
345 fk_obj: datamodel.ForeignKeyConstraint = constraint_obj
346 columns = [self._objects[column_id] for column_id in fk_obj.columns]
347 refcolumns = [self._objects[column_id] for column_id in fk_obj.referenced_columns]
348 if constraint_obj.on_delete is not None:
349 args["ondelete"] = constraint_obj.on_delete
350 if constraint_obj.on_update is not None:
351 args["onupdate"] = constraint_obj.on_update
352 constraint = ForeignKeyConstraint(columns, refcolumns, **args)
353 elif isinstance(constraint_obj, datamodel.CheckConstraint):
354 check_obj: datamodel.CheckConstraint = constraint_obj
355 expression = check_obj.expression
356 constraint = CheckConstraint(expression, **args)
357 elif isinstance(constraint_obj, datamodel.UniqueConstraint):
358 uniq_obj: datamodel.UniqueConstraint = constraint_obj
359 columns = [self._objects[column_id] for column_id in uniq_obj.columns]
360 constraint = UniqueConstraint(*columns, **args)
361 else:
362 raise ValueError(f"Unknown constraint type: {type(constraint_obj)}")
364 self._objects[constraint_obj.id] = constraint
366 return constraint
368 def build_index(self, index_obj: datamodel.Index) -> Index:
369 """Build a SQLAlchemy ``Index`` from a Felis `~felis.datamodel.Index`.
371 Parameters
372 ----------
373 index_obj
374 The Felis object from which to build the SQLAlchemy index.
376 Returns
377 -------
378 `~sqlalchemy.sql.schema.Index`
379 The SQLAlchemy index object.
380 """
381 columns = [self._objects[c_id] for c_id in (index_obj.columns if index_obj.columns else [])]
382 expressions = index_obj.expressions if index_obj.expressions else []
383 index = Index(index_obj.name, *columns, *expressions)
384 self._objects[index_obj.id] = index
385 return index