Coverage for python / felis / datamodel.py: 31%

696 statements  

« prev     ^ index     » next       coverage.py v7.12.0, created at 2025-12-13 00:54 -0800

1"""Define Pydantic data models for Felis.""" 

2 

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/>. 

23 

24from __future__ import annotations 

25 

26import json 

27import logging 

28import sys 

29from collections.abc import Sequence 

30from enum import StrEnum, auto 

31from typing import IO, Annotated, Any, Generic, Literal, TypeAlias, TypeVar 

32 

33import yaml 

34from astropy import units as units # type: ignore 

35from astropy.io.votable import ucd # type: ignore 

36from lsst.resources import ResourcePath, ResourcePathExpression 

37from pydantic import ( 

38 BaseModel, 

39 ConfigDict, 

40 Field, 

41 PrivateAttr, 

42 ValidationError, 

43 ValidationInfo, 

44 field_serializer, 

45 field_validator, 

46 model_validator, 

47) 

48from pydantic_core import InitErrorDetails 

49 

50from .db._dialects import get_supported_dialects, string_to_typeengine 

51from .db._sqltypes import get_type_func 

52from .types import Boolean, Byte, Char, Double, FelisType, Float, Int, Long, Short, String, Text, Unicode 

53 

54logger = logging.getLogger(__name__) 

55 

56__all__ = ( 

57 "BaseObject", 

58 "CheckConstraint", 

59 "Column", 

60 "ColumnOverrides", 

61 "ColumnResourceRef", 

62 "Constraint", 

63 "DataType", 

64 "ForeignKeyConstraint", 

65 "Index", 

66 "Resource", 

67 "Schema", 

68 "SchemaVersion", 

69 "Table", 

70 "UniqueConstraint", 

71) 

72 

73CONFIG = ConfigDict( 

74 populate_by_name=True, # Populate attributes by name. 

75 extra="forbid", # Do not allow extra fields. 

76 str_strip_whitespace=True, # Strip whitespace from string fields. 

77 use_enum_values=False, # Do not use enum values during serialization. 

78) 

79"""Pydantic model configuration as described in: 

80https://docs.pydantic.dev/2.0/api/config/#pydantic.config.ConfigDict 

81""" 

82 

83DESCR_MIN_LENGTH = 3 

84"""Minimum length for a description field.""" 

85 

86DescriptionStr: TypeAlias = Annotated[str, Field(min_length=DESCR_MIN_LENGTH)] 

87"""Type for a description, which must be three or more characters long.""" 

88 

89 

90class BaseObject(BaseModel): 

91 """Base model. 

92 

93 All classes representing objects in the Felis data model should inherit 

94 from this class. 

95 """ 

96 

97 model_config = CONFIG 

98 """Pydantic model configuration.""" 

99 

100 name: str 

101 """Name of the database object.""" 

102 

103 id: str = Field(alias="@id") 

104 """Unique identifier of the database object.""" 

105 

106 description: DescriptionStr | None = None 

107 """Description of the database object.""" 

108 

109 votable_utype: str | None = Field(None, alias="votable:utype") 

110 """VOTable utype (usage-specific or unique type) of the object.""" 

111 

112 @model_validator(mode="after") 

113 def check_description(self, info: ValidationInfo) -> BaseObject: 

114 """Check that the description is present if required. 

115 

116 Parameters 

117 ---------- 

118 info 

119 Validation context used to determine if the check is enabled. 

120 

121 Returns 

122 ------- 

123 `BaseObject` 

124 The object being validated. 

125 """ 

126 context = info.context 

127 if not context or not context.get("check_description", False): 

128 return self 

129 if self.description is None or self.description == "": 

130 raise ValueError("Description is required and must be non-empty") 

131 if len(self.description) < DESCR_MIN_LENGTH: 

132 raise ValueError(f"Description must be at least {DESCR_MIN_LENGTH} characters long") 

133 return self 

134 

135 

136class DataType(StrEnum): 

137 """``Enum`` representing the data types supported by Felis.""" 

138 

139 boolean = auto() 

140 byte = auto() 

141 short = auto() 

142 int = auto() 

143 long = auto() 

144 float = auto() 

145 double = auto() 

146 char = auto() 

147 string = auto() 

148 unicode = auto() 

149 text = auto() 

150 binary = auto() 

151 timestamp = auto() 

152 

153 

154def validate_ivoa_ucd(ivoa_ucd: str) -> str: 

155 """Validate IVOA UCD values. 

156 

157 Parameters 

158 ---------- 

159 ivoa_ucd 

160 IVOA UCD value to check. 

161 

162 Returns 

163 ------- 

164 `str` 

165 The IVOA UCD value if it is valid. 

166 

167 Raises 

168 ------ 

169 ValueError 

170 If the IVOA UCD value is invalid. 

171 """ 

172 if ivoa_ucd is not None: 

173 try: 

174 ucd.parse_ucd(ivoa_ucd, check_controlled_vocabulary=True, has_colon=";" in ivoa_ucd) 

175 except ValueError as e: 

176 raise ValueError(f"Invalid IVOA UCD: {e}") 

177 return ivoa_ucd 

178 

179 

180class Column(BaseObject): 

181 """Column model.""" 

182 

183 datatype: DataType 

184 """Datatype of the column.""" 

185 

186 length: int | None = Field(None, gt=0) 

187 """Length of the column.""" 

188 

189 precision: int | None = Field(None, ge=0) 

190 """The numerical precision of the column. 

191 

192 For timestamps, this is the number of fractional digits retained in the 

193 seconds field. 

194 """ 

195 

196 nullable: bool = True 

197 """Whether the column can be ``NULL``.""" 

198 

199 value: str | int | float | bool | None = None 

200 """Default value of the column.""" 

201 

202 autoincrement: bool | None = None 

203 """Whether the column is autoincremented.""" 

204 

205 ivoa_ucd: str | None = Field(None, alias="ivoa:ucd") 

206 """IVOA UCD of the column.""" 

207 

208 fits_tunit: str | None = Field(None, alias="fits:tunit") 

209 """FITS TUNIT of the column.""" 

210 

211 ivoa_unit: str | None = Field(None, alias="ivoa:unit") 

212 """IVOA unit of the column.""" 

213 

214 tap_column_index: int | None = Field(None, alias="tap:column_index") 

215 """TAP_SCHEMA column index of the column.""" 

216 

217 tap_principal: int | None = Field(0, alias="tap:principal", ge=0, le=1) 

218 """Whether this is a TAP_SCHEMA principal column.""" 

219 

220 votable_arraysize: int | str | None = Field(None, alias="votable:arraysize") 

221 """VOTable arraysize of the column.""" 

222 

223 tap_std: int | None = Field(0, alias="tap:std", ge=0, le=1) 

224 """TAP_SCHEMA indication that this column is defined by an IVOA standard. 

225 """ 

226 

227 votable_xtype: str | None = Field(None, alias="votable:xtype") 

228 """VOTable xtype (extended type) of the column.""" 

229 

230 votable_datatype: str | None = Field(None, alias="votable:datatype") 

231 """VOTable datatype of the column.""" 

232 

233 mysql_datatype: str | None = Field(None, alias="mysql:datatype") 

234 """MySQL datatype override on the column.""" 

235 

236 postgresql_datatype: str | None = Field(None, alias="postgresql:datatype") 

237 """PostgreSQL datatype override on the column.""" 

238 

239 _is_resource_ref: bool = PrivateAttr(False) 

240 """Whether this column is a resource reference column.""" 

241 

242 @model_validator(mode="after") 

243 def check_value(self) -> Column: 

244 """Check that the default value is valid. 

245 

246 Returns 

247 ------- 

248 `Column` 

249 The column being validated. 

250 """ 

251 if (value := self.value) is not None: 

252 if value is not None and self.autoincrement is True: 

253 raise ValueError("Column cannot have both a default value and be autoincremented") 

254 felis_type = FelisType.felis_type(self.datatype) 

255 if felis_type.is_numeric: 

256 if felis_type in (Byte, Short, Int, Long) and not isinstance(value, int): 

257 raise ValueError("Default value must be an int for integer type columns") 

258 elif felis_type in (Float, Double) and not isinstance(value, float): 

259 raise ValueError("Default value must be a decimal number for float and double columns") 

260 elif felis_type in (String, Char, Unicode, Text): 

261 if not isinstance(value, str): 

262 raise ValueError("Default value must be a string for string columns") 

263 if not len(value): 

264 raise ValueError("Default value must be a non-empty string for string columns") 

265 elif felis_type is Boolean and not isinstance(value, bool): 

266 raise ValueError("Default value must be a boolean for boolean columns") 

267 return self 

268 

269 @field_validator("ivoa_ucd") 

270 @classmethod 

271 def check_ivoa_ucd(cls, ivoa_ucd: str) -> str: 

272 """Check that IVOA UCD values are valid. 

273 

274 Parameters 

275 ---------- 

276 ivoa_ucd 

277 IVOA UCD value to check. 

278 

279 Returns 

280 ------- 

281 `str` 

282 The IVOA UCD value if it is valid. 

283 """ 

284 return validate_ivoa_ucd(ivoa_ucd) 

285 

286 @model_validator(mode="after") 

287 def check_units(self) -> Column: 

288 """Check that the ``fits:tunit`` or ``ivoa:unit`` field has valid 

289 units according to astropy. Only one may be provided. 

290 

291 Returns 

292 ------- 

293 `Column` 

294 The column being validated. 

295 

296 Raises 

297 ------ 

298 ValueError 

299 Raised if both FITS and IVOA units are provided, or if the unit is 

300 invalid. 

301 """ 

302 fits_unit = self.fits_tunit 

303 ivoa_unit = self.ivoa_unit 

304 

305 if fits_unit and ivoa_unit: 

306 raise ValueError("Column cannot have both FITS and IVOA units") 

307 unit = fits_unit or ivoa_unit 

308 

309 if unit is not None: 

310 try: 

311 units.Unit(unit) 

312 except ValueError as e: 

313 raise ValueError(f"Invalid unit: {e}") 

314 

315 return self 

316 

317 @model_validator(mode="before") 

318 @classmethod 

319 def check_length(cls, values: dict[str, Any]) -> dict[str, Any]: 

320 """Check that a valid length is provided for sized types. 

321 

322 Parameters 

323 ---------- 

324 values 

325 Values of the column. 

326 

327 Returns 

328 ------- 

329 `dict` [ `str`, `Any` ] 

330 The values of the column. 

331 

332 Raises 

333 ------ 

334 ValueError 

335 Raised if a length is not provided for a sized type. 

336 """ 

337 datatype = values.get("datatype") 

338 if datatype is None: 

339 # Skip this validation if datatype is not provided 

340 return values 

341 length = values.get("length") 

342 felis_type = FelisType.felis_type(datatype) 

343 if felis_type.is_sized and length is None: 

344 raise ValueError( 

345 f"Length must be provided for type '{datatype}'" 

346 + (f" in column '{values['@id']}'" if "@id" in values else "") 

347 ) 

348 elif not felis_type.is_sized and length is not None: 

349 logger.warning( 

350 f"The datatype '{datatype}' does not support a specified length" 

351 + (f" in column '{values['@id']}'" if "@id" in values else "") 

352 ) 

353 return values 

354 

355 @model_validator(mode="after") 

356 def check_redundant_datatypes(self, info: ValidationInfo) -> Column: 

357 """Check for redundant datatypes on columns. 

358 

359 Parameters 

360 ---------- 

361 info 

362 Validation context used to determine if the check is enabled. 

363 

364 Returns 

365 ------- 

366 `Column` 

367 The column being validated. 

368 

369 Raises 

370 ------ 

371 ValueError 

372 Raised if a datatype override is redundant. 

373 """ 

374 context = info.context 

375 if not context or not context.get("check_redundant_datatypes", False): 

376 return self 

377 if all( 

378 getattr(self, f"{dialect}:datatype", None) is not None 

379 for dialect in get_supported_dialects().keys() 

380 ): 

381 return self 

382 

383 datatype = self.datatype 

384 length: int | None = self.length or None 

385 

386 datatype_func = get_type_func(datatype) 

387 felis_type = FelisType.felis_type(datatype) 

388 if felis_type.is_sized: 

389 datatype_obj = datatype_func(length) 

390 else: 

391 datatype_obj = datatype_func() 

392 

393 for dialect_name, dialect in get_supported_dialects().items(): 

394 db_annotation = f"{dialect_name}_datatype" 

395 if datatype_string := self.model_dump().get(db_annotation): 

396 db_datatype_obj = string_to_typeengine(datatype_string, dialect, length) 

397 if datatype_obj.compile(dialect) == db_datatype_obj.compile(dialect): 

398 raise ValueError( 

399 "'{}: {}' is a redundant override of 'datatype: {}' in column '{}'{}".format( 

400 db_annotation, 

401 datatype_string, 

402 self.datatype, 

403 self.id, 

404 "" if length is None else f" with length {length}", 

405 ) 

406 ) 

407 else: 

408 logger.debug( 

409 f"Type override of 'datatype: {self.datatype}' " 

410 f"with '{db_annotation}: {datatype_string}' in column '{self.id}' " 

411 f"compiled to '{datatype_obj.compile(dialect)}' and " 

412 f"'{db_datatype_obj.compile(dialect)}'" 

413 ) 

414 return self 

415 

416 @model_validator(mode="after") 

417 def check_precision(self) -> Column: 

418 """Check that precision is only valid for timestamp columns. 

419 

420 Returns 

421 ------- 

422 `Column` 

423 The column being validated. 

424 """ 

425 if self.precision is not None and self.datatype != "timestamp": 

426 raise ValueError("Precision is only valid for timestamp columns") 

427 return self 

428 

429 @model_validator(mode="before") 

430 @classmethod 

431 def check_votable_arraysize(cls, values: dict[str, Any], info: ValidationInfo) -> dict[str, Any]: 

432 """Set the default value for the ``votable_arraysize`` field, which 

433 corresponds to ``arraysize`` in the IVOA VOTable standard. 

434 

435 Parameters 

436 ---------- 

437 values 

438 Values of the column. 

439 info 

440 Validation context used to determine if the check is enabled. 

441 

442 Returns 

443 ------- 

444 `dict` [ `str`, `Any` ] 

445 The values of the column. 

446 

447 Notes 

448 ----- 

449 Following the IVOA VOTable standard, an ``arraysize`` of 1 should not 

450 be used. 

451 """ 

452 if values.get("name", None) is None or values.get("datatype", None) is None: 

453 # Skip bad column data that will not validate 

454 return values 

455 context = info.context if info.context else {} 

456 arraysize = values.get("votable:arraysize", None) 

457 if arraysize is None: 

458 length = values.get("length", None) 

459 datatype = values.get("datatype") 

460 if length is not None and length > 1: 

461 # Following the IVOA standard, arraysize of 1 is disallowed 

462 if datatype == "char": 

463 arraysize = str(length) 

464 elif datatype in ("string", "unicode", "binary"): 

465 if context.get("force_unbounded_arraysize", False): 

466 arraysize = "*" 

467 logger.debug( 

468 f"Forced VOTable's 'arraysize' to '*' on column '{values['name']}' with datatype " 

469 + f"'{values['datatype']}' and length '{length}'" 

470 ) 

471 else: 

472 arraysize = f"{length}*" 

473 elif datatype in ("timestamp", "text"): 

474 arraysize = "*" 

475 if arraysize is not None: 

476 values["votable:arraysize"] = arraysize 

477 logger.debug( 

478 f"Set default 'votable:arraysize' to '{arraysize}' on column '{values['name']}'" 

479 + f" with datatype '{values['datatype']}' and length '{values.get('length', None)}'" 

480 ) 

481 else: 

482 logger.debug(f"Using existing 'votable:arraysize' of '{arraysize}' on column '{values['name']}'") 

483 if isinstance(values["votable:arraysize"], int): 

484 logger.warning( 

485 f"Usage of an integer value for 'votable:arraysize' in column '{values['name']}' is " 

486 + "deprecated" 

487 ) 

488 values["votable:arraysize"] = str(arraysize) 

489 return values 

490 

491 @field_serializer("datatype") 

492 def serialize_datatype(self, value: DataType) -> str: 

493 """Convert `DataType` to string when serializing to JSON/YAML. 

494 

495 Parameters 

496 ---------- 

497 value 

498 The `DataType` value to serialize. 

499 

500 Returns 

501 ------- 

502 `str` 

503 The serialized `DataType` value. 

504 """ 

505 return str(value) 

506 

507 @field_validator("datatype", mode="before") 

508 @classmethod 

509 def deserialize_datatype(cls, value: str) -> DataType: 

510 """Convert string back into `DataType` when loading from JSON/YAML. 

511 

512 Parameters 

513 ---------- 

514 value 

515 The string value to deserialize. 

516 

517 Returns 

518 ------- 

519 `DataType` 

520 The deserialized `DataType` value. 

521 """ 

522 return DataType(value) 

523 

524 @model_validator(mode="after") 

525 def check_votable_xtype(self) -> Column: 

526 """Set the default value for the ``votable_xtype`` field, which 

527 corresponds to an Extended Datatype or ``xtype`` in the IVOA VOTable 

528 standard. 

529 

530 Returns 

531 ------- 

532 `Column` 

533 The column being validated. 

534 

535 Notes 

536 ----- 

537 This is currently only set automatically for the Felis ``timestamp`` 

538 datatype. 

539 """ 

540 if self.datatype == DataType.timestamp and self.votable_xtype is None: 

541 self.votable_xtype = "timestamp" 

542 return self 

543 

544 

545class Constraint(BaseObject): 

546 """Table constraint model.""" 

547 

548 deferrable: bool = False 

549 """Whether this constraint will be declared as deferrable.""" 

550 

551 initially: Literal["IMMEDIATE", "DEFERRED"] | None = None 

552 """Value for ``INITIALLY`` clause; only used if `deferrable` is 

553 `True`.""" 

554 

555 @model_validator(mode="after") 

556 def check_deferrable(self) -> Constraint: 

557 """Check that the ``INITIALLY`` clause is only used if `deferrable` is 

558 `True`. 

559 

560 Returns 

561 ------- 

562 `Constraint` 

563 The constraint being validated. 

564 """ 

565 if self.initially is not None and not self.deferrable: 

566 raise ValueError("INITIALLY clause can only be used if deferrable is True") 

567 return self 

568 

569 

570class CheckConstraint(Constraint): 

571 """Table check constraint model.""" 

572 

573 type: Literal["Check"] = Field("Check", alias="@type") 

574 """Type of the constraint.""" 

575 

576 expression: str 

577 """Expression for the check constraint.""" 

578 

579 @field_serializer("type") 

580 def serialize_type(self, value: str) -> str: 

581 """Ensure '@type' is included in serialized output. 

582 

583 Parameters 

584 ---------- 

585 value 

586 The value to serialize. 

587 

588 Returns 

589 ------- 

590 `str` 

591 The serialized value. 

592 """ 

593 return value 

594 

595 

596class UniqueConstraint(Constraint): 

597 """Table unique constraint model.""" 

598 

599 type: Literal["Unique"] = Field("Unique", alias="@type") 

600 """Type of the constraint.""" 

601 

602 columns: list[str] 

603 """Columns in the unique constraint.""" 

604 

605 @field_serializer("type") 

606 def serialize_type(self, value: str) -> str: 

607 """Ensure '@type' is included in serialized output. 

608 

609 Parameters 

610 ---------- 

611 value 

612 The value to serialize. 

613 

614 Returns 

615 ------- 

616 `str` 

617 The serialized value. 

618 """ 

619 return value 

620 

621 

622class ForeignKeyConstraint(Constraint): 

623 """Table foreign key constraint model. 

624 

625 This constraint is used to define a foreign key relationship between two 

626 tables in the schema. There must be at least one column in the 

627 `columns` list, and at least one column in the `referenced_columns` list 

628 or a validation error will be raised. 

629 

630 Notes 

631 ----- 

632 These relationships will be reflected in the TAP_SCHEMA ``keys`` and 

633 ``key_columns`` data. 

634 """ 

635 

636 type: Literal["ForeignKey"] = Field("ForeignKey", alias="@type") 

637 """Type of the constraint.""" 

638 

639 columns: list[str] = Field(min_length=1) 

640 """The columns comprising the foreign key.""" 

641 

642 referenced_columns: list[str] = Field(alias="referencedColumns", min_length=1) 

643 """The columns referenced by the foreign key.""" 

644 

645 on_delete: Literal["CASCADE", "SET NULL", "SET DEFAULT", "RESTRICT", "NO ACTION"] | None = None 

646 """Action to take when the referenced row is deleted.""" 

647 

648 on_update: Literal["CASCADE", "SET NULL", "SET DEFAULT", "RESTRICT", "NO ACTION"] | None = None 

649 """Action to take when the referenced row is updated.""" 

650 

651 @field_serializer("type") 

652 def serialize_type(self, value: str) -> str: 

653 """Ensure '@type' is included in serialized output. 

654 

655 Parameters 

656 ---------- 

657 value 

658 The value to serialize. 

659 

660 Returns 

661 ------- 

662 `str` 

663 The serialized value. 

664 """ 

665 return value 

666 

667 @model_validator(mode="after") 

668 def check_column_lengths(self) -> ForeignKeyConstraint: 

669 """Check that the `columns` and `referenced_columns` lists have the 

670 same length. 

671 

672 Returns 

673 ------- 

674 `ForeignKeyConstraint` 

675 The foreign key constraint being validated. 

676 

677 Raises 

678 ------ 

679 ValueError 

680 Raised if the `columns` and `referenced_columns` lists do not have 

681 the same length. 

682 """ 

683 if len(self.columns) != len(self.referenced_columns): 

684 raise ValueError( 

685 "Columns and referencedColumns must have the same length for a ForeignKey constraint" 

686 ) 

687 return self 

688 

689 

690_ConstraintType = Annotated[ 

691 CheckConstraint | ForeignKeyConstraint | UniqueConstraint, Field(discriminator="type") 

692] 

693"""Type alias for a constraint type.""" 

694 

695 

696class Index(BaseObject): 

697 """Table index model. 

698 

699 An index can be defined on either columns or expressions, but not both. 

700 """ 

701 

702 columns: list[str] | None = None 

703 """Columns in the index.""" 

704 

705 expressions: list[str] | None = None 

706 """Expressions in the index.""" 

707 

708 @model_validator(mode="before") 

709 @classmethod 

710 def check_columns_or_expressions(cls, values: dict[str, Any]) -> dict[str, Any]: 

711 """Check that columns or expressions are specified, but not both. 

712 

713 Parameters 

714 ---------- 

715 values 

716 Values of the index. 

717 

718 Returns 

719 ------- 

720 `dict` [ `str`, `Any` ] 

721 The values of the index. 

722 

723 Raises 

724 ------ 

725 ValueError 

726 Raised if both columns and expressions are specified, or if neither 

727 are specified. 

728 """ 

729 if "columns" in values and "expressions" in values: 

730 raise ValueError("Defining columns and expressions is not valid") 

731 elif "columns" not in values and "expressions" not in values: 

732 raise ValueError("Must define columns or expressions") 

733 return values 

734 

735 

736ColumnRef: TypeAlias = str 

737"""Type alias for a column reference.""" 

738 

739 

740class ColumnGroup(BaseObject): 

741 """Column group model.""" 

742 

743 columns: list[ColumnRef | Column] = Field(..., min_length=1) 

744 """Columns in the group.""" 

745 

746 ivoa_ucd: str | None = Field(None, alias="ivoa:ucd") 

747 """IVOA UCD of the column.""" 

748 

749 table: Table | None = Field(None, exclude=True) 

750 """Reference to the parent table.""" 

751 

752 @field_validator("ivoa_ucd") 

753 @classmethod 

754 def check_ivoa_ucd(cls, ivoa_ucd: str) -> str: 

755 """Check that IVOA UCD values are valid. 

756 

757 Parameters 

758 ---------- 

759 ivoa_ucd 

760 IVOA UCD value to check. 

761 

762 Returns 

763 ------- 

764 `str` 

765 The IVOA UCD value if it is valid. 

766 """ 

767 return validate_ivoa_ucd(ivoa_ucd) 

768 

769 @model_validator(mode="after") 

770 def check_unique_columns(self) -> ColumnGroup: 

771 """Check that the columns list contains unique items. 

772 

773 Returns 

774 ------- 

775 `ColumnGroup` 

776 The column group being validated. 

777 """ 

778 column_ids = [col if isinstance(col, str) else col.id for col in self.columns] 

779 if len(column_ids) != len(set(column_ids)): 

780 raise ValueError("Columns in the group must be unique") 

781 return self 

782 

783 def _dereference_columns(self) -> None: 

784 """Dereference ColumnRef to Column objects.""" 

785 if self.table is None: 

786 raise ValueError("ColumnGroup must have a reference to its parent table") 

787 

788 dereferenced_columns: list[ColumnRef | Column] = [] 

789 for col in self.columns: 

790 if isinstance(col, str): 

791 # Dereference ColumnRef to Column object 

792 try: 

793 col_obj = self.table._find_column_by_id(col) 

794 except KeyError as e: 

795 raise ValueError(f"Column '{col}' not found in table '{self.table.name}'") from e 

796 dereferenced_columns.append(col_obj) 

797 else: 

798 dereferenced_columns.append(col) 

799 

800 self.columns = dereferenced_columns 

801 

802 @field_serializer("columns") 

803 def serialize_columns(self, columns: list[ColumnRef | Column]) -> list[str]: 

804 """Serialize columns as their IDs. 

805 

806 Parameters 

807 ---------- 

808 columns 

809 The columns to serialize. 

810 

811 Returns 

812 ------- 

813 `list` [ `str` ] 

814 The serialized column IDs. 

815 """ 

816 return [col if isinstance(col, str) else col.id for col in columns] 

817 

818 

819class ColumnOverrides(BaseModel): 

820 """Allowed overrides for a referenced column. 

821 

822 All fields are optional; missing values mean "inherit from the base 

823 column". 

824 """ 

825 

826 datatype: DataType | None = None 

827 """New datatype for the column.""" 

828 

829 length: int | None = None 

830 """New length for the column.""" 

831 

832 description: str | None = None 

833 """New description for the column.""" 

834 

835 nullable: bool = True 

836 """New nullable flag for the column.""" 

837 

838 tap_principal: int | None = Field(default=None, alias="tap:principal") 

839 """Override for the TAP_SCHEMA 'principal' flag.""" 

840 

841 tap_column_index: int | None = Field(default=None, alias="tap:column_index") 

842 """Override for the TAP_SCHEMA column index.""" 

843 

844 @field_serializer("datatype") 

845 def serialize_datatype(self, value: DataType | None) -> str | None: 

846 """Convert `DataType` to string when serializing to JSON/YAML. 

847 

848 Parameters 

849 ---------- 

850 value 

851 The `DataType` value to serialize, or None. 

852 

853 Returns 

854 ------- 

855 `str` | None 

856 The serialized `DataType` value, or None if the input was None. 

857 """ 

858 if value is None: 

859 return None 

860 return str(value) 

861 

862 @field_validator("datatype", mode="before") 

863 @classmethod 

864 def deserialize_datatype(cls, value: str | None) -> DataType | None: 

865 """Convert string back into `DataType` when loading from JSON/YAML. 

866 

867 Parameters 

868 ---------- 

869 value 

870 The string value to deserialize, or None. 

871 

872 Returns 

873 ------- 

874 `DataType` | None 

875 The deserialized `DataType` value, or None if the input was None. 

876 """ 

877 if value is None: 

878 return None 

879 return DataType(value) 

880 

881 

882class ColumnResourceRef(BaseModel): 

883 """A column which is dervived from an external resource.""" 

884 

885 ref_name: str | None = None 

886 """Name of the referenced column in the resource 

887 (if different from the key).""" 

888 

889 overrides: ColumnOverrides | None = None 

890 """Optional overrides of the referenced column's attributes.""" 

891 

892 

893# Type aliases for the nested mapping structure of referenced columns 

894ResourceColumnMap: TypeAlias = dict[str, ColumnResourceRef | None] 

895ResourceTableMap: TypeAlias = dict[str, ResourceColumnMap] 

896ResourceMap: TypeAlias = dict[str, ResourceTableMap] 

897 

898 

899class Table(BaseObject): 

900 """Table model.""" 

901 

902 primary_key: str | list[str] | None = Field(None, alias="primaryKey") 

903 """Primary key of the table.""" 

904 

905 tap_table_index: int | None = Field(None, alias="tap:table_index") 

906 """IVOA TAP_SCHEMA table index of the table.""" 

907 

908 mysql_engine: str | None = Field("MyISAM", alias="mysql:engine") 

909 """MySQL engine to use for the table.""" 

910 

911 mysql_charset: str | None = Field(None, alias="mysql:charset") 

912 """MySQL charset to use for the table.""" 

913 

914 columns: list[Column] = Field(default_factory=list) 

915 """Columns in the table.""" 

916 

917 column_refs: ResourceMap = Field(default_factory=dict, alias="columnRefs") 

918 """Referenced columns from external resources.""" 

919 

920 column_groups: list[ColumnGroup] = Field(default_factory=list, alias="columnGroups") 

921 """Column groups in the table.""" 

922 

923 constraints: list[_ConstraintType] = Field(default_factory=list) 

924 """Constraints on the table.""" 

925 

926 indexes: list[Index] = Field(default_factory=list) 

927 """Indexes on the table.""" 

928 

929 @field_validator("columns", mode="after") 

930 @classmethod 

931 def check_unique_column_names(cls, columns: list[Column]) -> list[Column]: 

932 """Check that column names are unique. 

933 

934 Parameters 

935 ---------- 

936 columns 

937 The columns to check. 

938 

939 Returns 

940 ------- 

941 `list` [ `Column` ] 

942 The columns if they are unique. 

943 

944 Raises 

945 ------ 

946 ValueError 

947 Raised if column names are not unique. 

948 """ 

949 if len(columns) != len(set(column.name for column in columns)): 

950 raise ValueError("Column names must be unique") 

951 return columns 

952 

953 @model_validator(mode="after") 

954 def check_tap_table_index(self, info: ValidationInfo) -> Table: 

955 """Check that the table has a TAP table index. 

956 

957 Parameters 

958 ---------- 

959 info 

960 Validation context used to determine if the check is enabled. 

961 

962 Returns 

963 ------- 

964 `Table` 

965 The table being validated. 

966 

967 Raises 

968 ------ 

969 ValueError 

970 Raised If the table is missing a TAP table index. 

971 """ 

972 context = info.context 

973 if not context or not context.get("check_tap_table_indexes", False): 

974 return self 

975 if self.tap_table_index is None: 

976 raise ValueError("Table is missing a TAP table index") 

977 return self 

978 

979 @model_validator(mode="after") 

980 def check_tap_principal(self, info: ValidationInfo) -> Table: 

981 """Check that at least one column is flagged as 'principal' for TAP 

982 purposes. 

983 

984 Parameters 

985 ---------- 

986 info 

987 Validation context used to determine if the check is enabled. 

988 

989 Returns 

990 ------- 

991 `Table` 

992 The table being validated. 

993 

994 Raises 

995 ------ 

996 ValueError 

997 Raised if the table is missing a column flagged as 'principal'. 

998 """ 

999 context = info.context 

1000 if not context or not context.get("check_tap_principal", False): 

1001 return self 

1002 for col in self.columns: 

1003 if col.tap_principal == 1: 

1004 return self 

1005 raise ValueError(f"Table '{self.name}' is missing at least one column designated as 'tap:principal'") 

1006 

1007 def _find_column_by_id(self, id: str) -> Column: 

1008 """Find a column by ID. 

1009 

1010 Parameters 

1011 ---------- 

1012 id 

1013 The ID of the column to find. 

1014 

1015 Returns 

1016 ------- 

1017 `Column` 

1018 The column with the given ID. 

1019 

1020 Raises 

1021 ------ 

1022 ValueError 

1023 Raised if the column is not found. 

1024 """ 

1025 for column in self.columns: 

1026 if column.id == id: 

1027 return column 

1028 raise KeyError(f"Column '{id}' not found in table '{self.name}'") 

1029 

1030 def _find_column_by_name(self, name: str) -> Column: 

1031 for column in self.columns: 

1032 if column.name == name: 

1033 return column 

1034 raise KeyError(f"Column '{name}' not found in table '{self.name}'") 

1035 

1036 @model_validator(mode="after") 

1037 def dereference_column_groups(self: Table) -> Table: 

1038 """Dereference columns in column groups. 

1039 

1040 Returns 

1041 ------- 

1042 `Table` 

1043 The table with dereferenced column groups. 

1044 """ 

1045 for group in self.column_groups: 

1046 group.table = self 

1047 group._dereference_columns() 

1048 return self 

1049 

1050 @field_serializer("columns") 

1051 def _serialize_columns(self, columns: list[Column]) -> list[dict[str, Any]]: 

1052 """Serialize only non-resource columns.""" 

1053 return [ 

1054 col.model_dump( 

1055 by_alias=True, 

1056 exclude_none=True, 

1057 exclude_defaults=True, 

1058 ) 

1059 for col in columns 

1060 if not col._is_resource_ref 

1061 ] 

1062 

1063 

1064class SchemaVersion(BaseModel): 

1065 """Schema version model.""" 

1066 

1067 current: str 

1068 """The current version of the schema.""" 

1069 

1070 compatible: list[str] = Field(default_factory=list) 

1071 """The compatible versions of the schema.""" 

1072 

1073 read_compatible: list[str] = Field(default_factory=list) 

1074 """The read compatible versions of the schema.""" 

1075 

1076 

1077class SchemaIdVisitor: 

1078 """Visit a schema and build the map of IDs to objects. 

1079 

1080 Notes 

1081 ----- 

1082 Duplicates are added to a set when they are encountered, which can be 

1083 accessed via the ``duplicates`` attribute. The presence of duplicates will 

1084 not throw an error. Only the first object with a given ID will be added to 

1085 the map, but this should not matter, since a ``ValidationError`` will be 

1086 thrown by the ``model_validator`` method if any duplicates are found in the 

1087 schema. 

1088 """ 

1089 

1090 def __init__(self) -> None: 

1091 """Create a new SchemaVisitor.""" 

1092 self.schema: Schema | None = None 

1093 self.duplicates: set[str] = set() 

1094 

1095 def add(self, obj: BaseObject) -> None: 

1096 """Add an object to the ID map. 

1097 

1098 Parameters 

1099 ---------- 

1100 obj 

1101 The object to add to the ID map. 

1102 """ 

1103 if hasattr(obj, "id"): 

1104 obj_id = getattr(obj, "id") 

1105 if self.schema is not None: 

1106 if obj_id in self.schema._id_map: 

1107 self.duplicates.add(obj_id) 

1108 else: 

1109 self.schema._id_map[obj_id] = obj 

1110 

1111 def visit_schema(self, schema: Schema) -> None: 

1112 """Visit the objects in a schema and build the ID map. 

1113 

1114 Parameters 

1115 ---------- 

1116 schema 

1117 The schema object to visit. 

1118 

1119 Notes 

1120 ----- 

1121 This will set an internal variable pointing to the schema object. 

1122 """ 

1123 self.schema = schema 

1124 self.duplicates.clear() 

1125 self.add(self.schema) 

1126 for table in self.schema.tables: 

1127 self.visit_table(table) 

1128 

1129 def visit_table(self, table: Table) -> None: 

1130 """Visit a table object. 

1131 

1132 Parameters 

1133 ---------- 

1134 table 

1135 The table object to visit. 

1136 """ 

1137 self.add(table) 

1138 for column in table.columns: 

1139 self.visit_column(column) 

1140 for constraint in table.constraints: 

1141 self.visit_constraint(constraint) 

1142 

1143 def visit_column(self, column: Column) -> None: 

1144 """Visit a column object. 

1145 

1146 Parameters 

1147 ---------- 

1148 column 

1149 The column object to visit. 

1150 """ 

1151 self.add(column) 

1152 

1153 def visit_constraint(self, constraint: Constraint) -> None: 

1154 """Visit a constraint object. 

1155 

1156 Parameters 

1157 ---------- 

1158 constraint 

1159 The constraint object to visit. 

1160 """ 

1161 self.add(constraint) 

1162 

1163 

1164T = TypeVar("T", bound=BaseObject) 

1165 

1166 

1167def _strip_ids(data: Any) -> Any: 

1168 """Recursively strip '@id' fields from a dictionary or list. 

1169 

1170 Parameters 

1171 ---------- 

1172 data 

1173 The data to strip IDs from, which can be a dictionary, list, or any 

1174 other type. Other types will be returned unchanged. 

1175 """ 

1176 if isinstance(data, dict): 

1177 data.pop("@id", None) 

1178 for k, v in data.items(): 

1179 data[k] = _strip_ids(v) 

1180 return data 

1181 elif isinstance(data, list): 

1182 return [_strip_ids(item) for item in data] 

1183 else: 

1184 return data 

1185 

1186 

1187def _append_error( 

1188 errors: list[InitErrorDetails], 

1189 loc: tuple, 

1190 input_value: Any, 

1191 error_message: str, 

1192 error_type: str = "value_error", 

1193) -> None: 

1194 """Append an error to the errors list. 

1195 

1196 Parameters 

1197 ---------- 

1198 errors : list[InitErrorDetails] 

1199 The list of errors to append to. 

1200 loc : tuple 

1201 The location of the error in the schema. 

1202 input_value : Any 

1203 The input value that caused the error. 

1204 error_message : str 

1205 The error message to include in the context. 

1206 """ 

1207 errors.append( 

1208 { 

1209 "type": error_type, 

1210 "loc": loc, 

1211 "input": input_value, 

1212 "ctx": {"error": error_message}, 

1213 } 

1214 ) 

1215 

1216 

1217class Resource(BaseModel): 

1218 """A resource definition referencing an external schema.""" 

1219 

1220 uri: str = Field(..., description="Resource URI or path") 

1221 """URI of the schema resource which may be a local path, ``resource://``, 

1222 or remote URL.""" 

1223 

1224 

1225class Schema(BaseObject, Generic[T]): 

1226 """Database schema model. 

1227 

1228 This represents a database schema, which contains one or more tables. 

1229 """ 

1230 

1231 version: SchemaVersion | str | None = None 

1232 """The version of the schema.""" 

1233 

1234 tables: Sequence[Table] 

1235 """The tables in the schema.""" 

1236 

1237 resources: dict[str, Resource] = Field(default_factory=dict) 

1238 """External resources referenced by this schema.""" 

1239 

1240 _id_map: dict[str, Any] = PrivateAttr(default_factory=dict) 

1241 """Map of IDs to objects.""" 

1242 

1243 _resource_map: dict[str, Schema] = PrivateAttr(default_factory=dict) 

1244 """Map of resource names to loaded schemas.""" 

1245 

1246 @model_validator(mode="after") 

1247 def _load_resources(self: Schema, info: ValidationInfo) -> Schema: 

1248 """Load external resources referenced by this schema into an internal 

1249 mapping of resource names to their `Schema` objects. 

1250 

1251 Returns 

1252 ------- 

1253 `Schema` 

1254 The schema being validated. 

1255 

1256 Raises 

1257 ------ 

1258 ValueError 

1259 Raised if a resource cannot be loaded. 

1260 """ 

1261 if info.context: 

1262 context = info.context.copy() 

1263 # Ignore this flag for loading the resources themselves 

1264 context.pop("dereference_resources", None) 

1265 else: 

1266 context = {} 

1267 

1268 for resource_name, resource in self.resources.items(): 

1269 uri = resource.uri 

1270 try: 

1271 loaded_schema = Schema.from_uri(uri, context=context) 

1272 self._resource_map[resource_name] = loaded_schema 

1273 logger.debug(f"Loaded resource '{resource_name}' from URI '{uri}'") 

1274 except Exception as e: 

1275 raise ValueError(f"Failed to load resource '{resource_name}' from URI '{uri}': {e}") from e 

1276 return self 

1277 

1278 def _find_table_by_name(self, name: str) -> Table: 

1279 """Find a table by name. 

1280 

1281 Parameters 

1282 ---------- 

1283 name 

1284 The name of the table to find. 

1285 

1286 Returns 

1287 ------- 

1288 `Table` 

1289 The table with the given name. 

1290 

1291 Raises 

1292 ------ 

1293 KeyError 

1294 Raised if the table is not found. 

1295 """ 

1296 for table in self.tables: 

1297 if table.name == name: 

1298 return table 

1299 raise KeyError(f"Table '{name}' not found in schema '{self.name}'") 

1300 

1301 @model_validator(mode="after") 

1302 def _dereference_resource_columns(self: Schema, info: ValidationInfo) -> Schema: 

1303 """Dereference columns from external resources and add them to the 

1304 tables in this schema. 

1305 """ 

1306 context = info.context 

1307 if context is not None and context.get("dereference_resources", False): 

1308 dereference_resources = True 

1309 else: 

1310 dereference_resources = False 

1311 for table in self.tables: 

1312 if column_refs := table.column_refs: 

1313 for resource_name, tables in column_refs.items(): 

1314 resource_schema = self._resource_map.get(resource_name) 

1315 if resource_schema is None: 

1316 raise ValueError(f"Schema resource '{resource_name}' was not found in resources") 

1317 self._process_column_refs( 

1318 table, 

1319 tables, 

1320 resource_schema, 

1321 dereference_resources, 

1322 ) 

1323 if dereference_resources and len(table.column_refs) > 0: 

1324 # Clear column refs in table if fully dereferencing 

1325 logger.debug( 

1326 f"Clearing columnRefs in table '{table.name}' after dereferencing resource columns" 

1327 ) 

1328 table.column_refs = {} 

1329 return self 

1330 

1331 def _process_column_refs( 

1332 self, 

1333 table: Table, 

1334 ref_tables: ResourceTableMap, 

1335 resource_schema: Schema, 

1336 dereference_resources: bool = False, 

1337 ) -> None: 

1338 """Process column references from an external resource and add them 

1339 to the given table. 

1340 """ 

1341 for table_name, columns in ref_tables.items(): 

1342 try: 

1343 resource_table = resource_schema._find_table_by_name(table_name) 

1344 except KeyError as e: 

1345 raise ValueError( 

1346 f"Table '{table_name}' not found in resource '{resource_schema.name}'" 

1347 ) from e 

1348 for local_column_name, column_ref in columns.items(): 

1349 if column_ref is not None and column_ref.ref_name is not None: 

1350 # Use specified ref_name 

1351 ref_column_name = column_ref.ref_name 

1352 else: 

1353 # Use the local column name if no ref_name 

1354 # specified 

1355 ref_column_name = local_column_name 

1356 

1357 # Check if referenced column exists in resource 

1358 try: 

1359 base_column = resource_table._find_column_by_name(ref_column_name) 

1360 except KeyError: 

1361 # The ref_name is specified but column is not 

1362 # found 

1363 if column_ref is not None and column_ref.ref_name is not None: 

1364 raise ValueError( 

1365 f"Column '{ref_column_name}' not found in table '{table_name}' " 

1366 f"from resource '{resource_schema.name}'" 

1367 ) 

1368 # The ref_name is not specified and the local 

1369 # column name is not found 

1370 raise ValueError( 

1371 f"Column '{local_column_name}' not found in table '{table_name}' " 

1372 f"from resource '{resource_schema.name}' and no ref_name provided" 

1373 ) 

1374 

1375 # Create a copy of the base column and apply 

1376 # overrides 

1377 column_copy = base_column.model_copy() 

1378 # Set the local name (key from the mapping) 

1379 column_copy.name = local_column_name 

1380 if not dereference_resources: 

1381 column_copy._is_resource_ref = True 

1382 

1383 if column_ref is not None and column_ref.overrides is not None: 

1384 overrides = column_ref.overrides 

1385 for field_name, override_value in overrides.model_dump().items(): 

1386 if override_value is not None: 

1387 setattr(column_copy, field_name, override_value) 

1388 

1389 # Manually set the ID of the copied column as 

1390 # ID generation has already occurred 

1391 column_copy.id = f"{table.id}.{local_column_name}" 

1392 table.columns.append(column_copy) 

1393 logger.debug( 

1394 f"Dereferenced column '{local_column_name}' from table '{table_name}' " 

1395 f"in resource '{resource_schema.name}' into table '{table.name}'" 

1396 ) 

1397 

1398 @model_validator(mode="before") 

1399 @classmethod 

1400 def generate_ids(cls, values: dict[str, Any], info: ValidationInfo) -> dict[str, Any]: 

1401 """Generate IDs for objects that do not have them. 

1402 

1403 Parameters 

1404 ---------- 

1405 values 

1406 The values of the schema. 

1407 info 

1408 Validation context used to determine if ID generation is enabled. 

1409 

1410 Returns 

1411 ------- 

1412 `dict` [ `str`, `Any` ] 

1413 The values of the schema with generated IDs. 

1414 """ 

1415 context = info.context 

1416 if not context or not context.get("id_generation", False): 

1417 logger.debug("Skipping ID generation") 

1418 return values 

1419 schema_name = values["name"] 

1420 if "@id" not in values: 

1421 values["@id"] = f"#{schema_name}" 

1422 logger.debug(f"Generated ID '{values['@id']}' for schema '{schema_name}'") 

1423 if "tables" in values: 

1424 for table in values["tables"]: 

1425 if "@id" not in table: 

1426 table["@id"] = f"#{table['name']}" 

1427 logger.debug(f"Generated ID '{table['@id']}' for table '{table['name']}'") 

1428 if "columns" in table: 

1429 for column in table["columns"]: 

1430 if "@id" not in column: 

1431 column["@id"] = f"#{table['name']}.{column['name']}" 

1432 logger.debug(f"Generated ID '{column['@id']}' for column '{column['name']}'") 

1433 if "columnGroups" in table: 

1434 for column_group in table["columnGroups"]: 

1435 if "@id" not in column_group: 

1436 column_group["@id"] = f"#{table['name']}.{column_group['name']}" 

1437 logger.debug( 

1438 f"Generated ID '{column_group['@id']}' for column group " 

1439 f"'{column_group['name']}'" 

1440 ) 

1441 if "constraints" in table: 

1442 for constraint in table["constraints"]: 

1443 if "@id" not in constraint: 

1444 constraint["@id"] = f"#{constraint['name']}" 

1445 logger.debug( 

1446 f"Generated ID '{constraint['@id']}' for constraint '{constraint['name']}'" 

1447 ) 

1448 if "indexes" in table: 

1449 for index in table["indexes"]: 

1450 if "@id" not in index: 

1451 index["@id"] = f"#{index['name']}" 

1452 logger.debug(f"Generated ID '{index['@id']}' for index '{index['name']}'") 

1453 return values 

1454 

1455 @field_validator("tables", mode="after") 

1456 @classmethod 

1457 def check_unique_table_names(cls, tables: list[Table]) -> list[Table]: 

1458 """Check that table names are unique. 

1459 

1460 Parameters 

1461 ---------- 

1462 tables 

1463 The tables to check. 

1464 

1465 Returns 

1466 ------- 

1467 `list` [ `Table` ] 

1468 The tables if they are unique. 

1469 

1470 Raises 

1471 ------ 

1472 ValueError 

1473 Raised if table names are not unique. 

1474 """ 

1475 if len(tables) != len(set(table.name for table in tables)): 

1476 raise ValueError("Table names must be unique") 

1477 return tables 

1478 

1479 @model_validator(mode="after") 

1480 def check_tap_table_indexes(self, info: ValidationInfo) -> Schema: 

1481 """Check that the TAP table indexes are unique. 

1482 

1483 Parameters 

1484 ---------- 

1485 info 

1486 The validation context used to determine if the check is enabled. 

1487 

1488 Returns 

1489 ------- 

1490 `Schema` 

1491 The schema being validated. 

1492 """ 

1493 context = info.context 

1494 if not context or not context.get("check_tap_table_indexes", False): 

1495 return self 

1496 table_indicies = set() 

1497 for table in self.tables: 

1498 table_index = table.tap_table_index 

1499 if table_index is not None: 

1500 if table_index in table_indicies: 

1501 raise ValueError(f"Duplicate 'tap:table_index' value {table_index} found in schema") 

1502 table_indicies.add(table_index) 

1503 return self 

1504 

1505 @model_validator(mode="after") 

1506 def check_unique_constraint_names(self: Schema) -> Schema: 

1507 """Check for duplicate constraint names in the schema. 

1508 

1509 Returns 

1510 ------- 

1511 `Schema` 

1512 The schema being validated. 

1513 

1514 Raises 

1515 ------ 

1516 ValueError 

1517 Raised if duplicate constraint names are found in the schema. 

1518 """ 

1519 constraint_names = set() 

1520 duplicate_names = [] 

1521 

1522 for table in self.tables: 

1523 for constraint in table.constraints: 

1524 constraint_name = constraint.name 

1525 if constraint_name in constraint_names: 

1526 duplicate_names.append(constraint_name) 

1527 else: 

1528 constraint_names.add(constraint_name) 

1529 

1530 if duplicate_names: 

1531 raise ValueError(f"Duplicate constraint names found in schema: {duplicate_names}") 

1532 

1533 return self 

1534 

1535 @model_validator(mode="after") 

1536 def check_unique_index_names(self: Schema) -> Schema: 

1537 """Check for duplicate index names in the schema. 

1538 

1539 Returns 

1540 ------- 

1541 `Schema` 

1542 The schema being validated. 

1543 

1544 Raises 

1545 ------ 

1546 ValueError 

1547 Raised if duplicate index names are found in the schema. 

1548 """ 

1549 index_names = set() 

1550 duplicate_names = [] 

1551 

1552 for table in self.tables: 

1553 for index in table.indexes: 

1554 index_name = index.name 

1555 if index_name in index_names: 

1556 duplicate_names.append(index_name) 

1557 else: 

1558 index_names.add(index_name) 

1559 

1560 if duplicate_names: 

1561 raise ValueError(f"Duplicate index names found in schema: {duplicate_names}") 

1562 

1563 return self 

1564 

1565 @model_validator(mode="after") 

1566 def create_id_map(self: Schema) -> Schema: 

1567 """Create a map of IDs to objects. 

1568 

1569 Returns 

1570 ------- 

1571 `Schema` 

1572 The schema with the ID map created. 

1573 

1574 Raises 

1575 ------ 

1576 ValueError 

1577 Raised if duplicate identifiers are found in the schema. 

1578 """ 

1579 if self._id_map: 

1580 logger.debug("Ignoring call to create_id_map() - ID map was already populated") 

1581 return self 

1582 visitor: SchemaIdVisitor = SchemaIdVisitor() 

1583 visitor.visit_schema(self) 

1584 if len(visitor.duplicates): 

1585 raise ValueError( 

1586 "Duplicate IDs found in schema:\n " + "\n ".join(visitor.duplicates) + "\n" 

1587 ) 

1588 logger.debug("Created ID map with %d entries", len(self._id_map)) 

1589 return self 

1590 

1591 def _validate_column_id( 

1592 self: Schema, 

1593 column_id: str, 

1594 loc: tuple, 

1595 errors: list[InitErrorDetails], 

1596 ) -> None: 

1597 """Validate a column ID from a constraint and append errors if invalid. 

1598 

1599 Parameters 

1600 ---------- 

1601 schema : Schema 

1602 The schema being validated. 

1603 column_id : str 

1604 The column ID to validate. 

1605 loc : tuple 

1606 The location of the error in the schema. 

1607 errors : list[InitErrorDetails] 

1608 The list of errors to append to. 

1609 """ 

1610 if column_id not in self: 

1611 _append_error( 

1612 errors, 

1613 loc, 

1614 column_id, 

1615 f"Column ID '{column_id}' not found in schema", 

1616 ) 

1617 elif not isinstance(self[column_id], Column): 

1618 _append_error( 

1619 errors, 

1620 loc, 

1621 column_id, 

1622 f"ID '{column_id}' does not refer to a Column object", 

1623 ) 

1624 

1625 def _validate_foreign_key_column( 

1626 self: Schema, 

1627 column_id: str, 

1628 table: Table, 

1629 loc: tuple, 

1630 errors: list[InitErrorDetails], 

1631 ) -> None: 

1632 """Validate a foreign key column ID from a constraint and append errors 

1633 if invalid. 

1634 

1635 Parameters 

1636 ---------- 

1637 schema : Schema 

1638 The schema being validated. 

1639 column_id : str 

1640 The foreign key column ID to validate. 

1641 loc : tuple 

1642 The location of the error in the schema. 

1643 errors : list[InitErrorDetails] 

1644 The list of errors to append to. 

1645 """ 

1646 try: 

1647 table._find_column_by_id(column_id) 

1648 except KeyError: 

1649 _append_error( 

1650 errors, 

1651 loc, 

1652 column_id, 

1653 f"Column '{column_id}' not found in table '{table.name}'", 

1654 ) 

1655 

1656 @model_validator(mode="after") 

1657 def check_constraints(self: Schema) -> Schema: 

1658 """Check constraint objects for validity. This needs to be deferred 

1659 until after the schema is fully loaded and the ID map is created. 

1660 

1661 Raises 

1662 ------ 

1663 pydantic.ValidationError 

1664 Raised if any constraints are invalid. 

1665 

1666 Returns 

1667 ------- 

1668 `Schema` 

1669 The schema being validated. 

1670 """ 

1671 errors: list[InitErrorDetails] = [] 

1672 

1673 for table_index, table in enumerate(self.tables): 

1674 for constraint_index, constraint in enumerate(table.constraints): 

1675 column_ids: list[str] = [] 

1676 referenced_column_ids: list[str] = [] 

1677 

1678 if isinstance(constraint, ForeignKeyConstraint): 

1679 column_ids += constraint.columns 

1680 referenced_column_ids += constraint.referenced_columns 

1681 elif isinstance(constraint, UniqueConstraint): 

1682 column_ids += constraint.columns 

1683 # No extra checks are required on CheckConstraint objects. 

1684 

1685 # Validate the foreign key columns 

1686 for column_id in column_ids: 

1687 self._validate_column_id( 

1688 column_id, 

1689 ( 

1690 "tables", 

1691 table_index, 

1692 "constraints", 

1693 constraint_index, 

1694 "columns", 

1695 column_id, 

1696 ), 

1697 errors, 

1698 ) 

1699 # Check that the foreign key column is within the source 

1700 # table. 

1701 self._validate_foreign_key_column( 

1702 column_id, 

1703 table, 

1704 ( 

1705 "tables", 

1706 table_index, 

1707 "constraints", 

1708 constraint_index, 

1709 "columns", 

1710 column_id, 

1711 ), 

1712 errors, 

1713 ) 

1714 

1715 # Validate the primary key (reference) columns 

1716 for referenced_column_id in referenced_column_ids: 

1717 self._validate_column_id( 

1718 referenced_column_id, 

1719 ( 

1720 "tables", 

1721 table_index, 

1722 "constraints", 

1723 constraint_index, 

1724 "referenced_columns", 

1725 referenced_column_id, 

1726 ), 

1727 errors, 

1728 ) 

1729 

1730 if errors: 

1731 raise ValidationError.from_exception_data("Schema validation failed", errors) 

1732 

1733 return self 

1734 

1735 def __getitem__(self, id: str) -> BaseObject: 

1736 """Get an object by its ID. 

1737 

1738 Parameters 

1739 ---------- 

1740 id 

1741 The ID of the object to get. 

1742 

1743 Raises 

1744 ------ 

1745 KeyError 

1746 Raised if the object with the given ID is not found in the schema. 

1747 """ 

1748 if id not in self: 

1749 raise KeyError(f"Object with ID '{id}' not found in schema") 

1750 return self._id_map[id] 

1751 

1752 def __contains__(self, id: str) -> bool: 

1753 """Check if an object with the given ID is in the schema. 

1754 

1755 Parameters 

1756 ---------- 

1757 id 

1758 The ID of the object to check. 

1759 """ 

1760 return id in self._id_map 

1761 

1762 def find_object_by_id(self, id: str, obj_type: type[T]) -> T: 

1763 """Find an object with the given type by its ID. 

1764 

1765 Parameters 

1766 ---------- 

1767 id 

1768 The ID of the object to find. 

1769 obj_type 

1770 The type of the object to find. 

1771 

1772 Returns 

1773 ------- 

1774 BaseObject 

1775 The object with the given ID and type. 

1776 

1777 Raises 

1778 ------ 

1779 KeyError 

1780 If the object with the given ID is not found in the schema. 

1781 TypeError 

1782 If the object that is found does not have the right type. 

1783 

1784 Notes 

1785 ----- 

1786 The actual return type is the user-specified argument ``T``, which is 

1787 expected to be a subclass of `BaseObject`. 

1788 """ 

1789 obj = self[id] 

1790 if not isinstance(obj, obj_type): 

1791 raise TypeError(f"Object with ID '{id}' is not of type '{obj_type.__name__}'") 

1792 return obj 

1793 

1794 def get_table_by_column(self, column: Column) -> Table: 

1795 """Find the table that contains a column. 

1796 

1797 Parameters 

1798 ---------- 

1799 column 

1800 The column to find. 

1801 

1802 Returns 

1803 ------- 

1804 `Table` 

1805 The table that contains the column. 

1806 

1807 Raises 

1808 ------ 

1809 ValueError 

1810 If the column is not found in any table. 

1811 """ 

1812 for table in self.tables: 

1813 if column in table.columns: 

1814 return table 

1815 raise ValueError(f"Column '{column.name}' not found in any table") 

1816 

1817 @classmethod 

1818 def from_uri(cls, resource_path: ResourcePathExpression, context: dict[str, Any] = {}) -> Schema: 

1819 """Load a `Schema` from a string representing a ``ResourcePath``. 

1820 

1821 Parameters 

1822 ---------- 

1823 resource_path 

1824 The ``ResourcePath`` pointing to a YAML file. 

1825 context 

1826 Pydantic context to be used in validation. 

1827 

1828 Returns 

1829 ------- 

1830 `str` 

1831 The ID of the object. 

1832 

1833 Raises 

1834 ------ 

1835 yaml.YAMLError 

1836 Raised if there is an error loading the YAML data. 

1837 ValueError 

1838 Raised if there is an error reading the resource. 

1839 pydantic.ValidationError 

1840 Raised if the schema fails validation. 

1841 """ 

1842 logger.debug(f"Loading schema from: '{resource_path}'") 

1843 try: 

1844 rp_stream = ResourcePath(resource_path).read() 

1845 except Exception as e: 

1846 raise ValueError(f"Error reading resource from '{resource_path}' : {e}") from e 

1847 yaml_data = yaml.safe_load(rp_stream) 

1848 return Schema.model_validate(yaml_data, context=context) 

1849 

1850 @classmethod 

1851 def from_stream(cls, source: IO[str], context: dict[str, Any] = {}) -> Schema: 

1852 """Load a `Schema` from a file stream which should contain YAML data. 

1853 

1854 Parameters 

1855 ---------- 

1856 source 

1857 The file stream to read from. 

1858 context 

1859 Pydantic context to be used in validation. 

1860 

1861 Returns 

1862 ------- 

1863 `Schema` 

1864 The Felis schema loaded from the stream. 

1865 

1866 Raises 

1867 ------ 

1868 yaml.YAMLError 

1869 Raised if there is an error loading the YAML file. 

1870 pydantic.ValidationError 

1871 Raised if the schema fails validation. 

1872 """ 

1873 logger.debug("Loading schema from: '%s'", source) 

1874 yaml_data = yaml.safe_load(source) 

1875 return Schema.model_validate(yaml_data, context=context) 

1876 

1877 def _model_dump(self, strip_ids: bool = False) -> dict[str, Any]: 

1878 """Dump the schema as a dictionary with some default arguments 

1879 applied. 

1880 

1881 Parameters 

1882 ---------- 

1883 strip_ids 

1884 Whether to strip the IDs from the dumped data. Defaults to `False`. 

1885 

1886 Returns 

1887 ------- 

1888 `dict` [ `str`, `Any` ] 

1889 The dumped schema data as a dictionary. 

1890 """ 

1891 data = self.model_dump(by_alias=True, exclude_none=True, exclude_defaults=True) 

1892 if strip_ids: 

1893 data = _strip_ids(data) 

1894 return data 

1895 

1896 def dump_yaml(self, stream: IO[str] = sys.stdout, strip_ids: bool = False) -> None: 

1897 """Pretty print the schema as YAML. 

1898 

1899 Parameters 

1900 ---------- 

1901 stream 

1902 The stream to write the YAML data to. 

1903 strip_ids 

1904 Whether to strip the IDs from the dumped data. Defaults to `False`. 

1905 """ 

1906 data = self._model_dump(strip_ids=strip_ids) 

1907 yaml.safe_dump( 

1908 data, 

1909 stream, 

1910 default_flow_style=False, 

1911 sort_keys=False, 

1912 ) 

1913 

1914 def dump_json(self, stream: IO[str] = sys.stdout, strip_ids: bool = False) -> None: 

1915 """Pretty print the schema as JSON. 

1916 

1917 Parameters 

1918 ---------- 

1919 stream 

1920 The stream to write the JSON data to. 

1921 strip_ids 

1922 Whether to strip the IDs from the dumped data. Defaults to `False`. 

1923 """ 

1924 data = self._model_dump(strip_ids=strip_ids) 

1925 json.dump( 

1926 data, 

1927 stream, 

1928 indent=4, 

1929 sort_keys=False, 

1930 )