Coverage for snekql/_schema_plan.py: 96%

1"""Shared model-derived schema planning for backend schema startup.""" 

2 

3from __future__ import annotations 

4 

5from collections.abc import Sequence 

6from dataclasses import dataclass 

7from typing import Any 

8 

9from snekql.errors import SchemaError 

10from snekql.indexes import NormalizedIndex 

11from snekql.model import Table, require_model_columns, require_model_table_name 

12from snekql.storage import Attr, SchemaPolicy 

13 

14 

15@dataclass(frozen=True) 

16class PlannedColumn: 

17 """One model column resolved for schema startup.""" 

18 

19 column: Attr[Any, Any, Any, Any, Any] 

20 name: str 

21 

22 

23@dataclass(frozen=True) 

24class PlannedModel: 

25 """One table model's backend-neutral schema startup inputs.""" 

26 

27 columns: tuple[PlannedColumn, ...] 

28 indexes: tuple[NormalizedIndex, ...] 

29 model: type[Table[Any]] 

30 table_name: str 

31 

32 

33@dataclass(frozen=True) 

34class SchemaPlan: 

35 """Ordered schema startup plan shared by backend adapters.""" 

36 

37 models: tuple[PlannedModel, ...] 

38 

39 

40def _column_unique_indexes( 

41 table_name: str, 

42 columns: tuple[PlannedColumn, ...], 

43) -> tuple[NormalizedIndex, ...]: 

44 return tuple( 

45 NormalizedIndex( 

46 column_names=(planned_column.name,), 

47 name=f"ux_{table_name}_{planned_column.name}", 

48 unique=True, 

49 ) 

50 for planned_column in columns 

51 if planned_column.column.unique 

52 ) 

53 

54 

55def _model_indexes( 

56 model: type[Table[Any]], 

57 table_name: str, 

58 columns: tuple[PlannedColumn, ...], 

59) -> tuple[NormalizedIndex, ...]: 

60 table_indexes = getattr(model, "__snekql_indexes__", ()) 

61 return (*_column_unique_indexes(table_name, columns), *table_indexes) 

62 

63 

64def _plan_model(model: type[Table[Any]]) -> PlannedModel: 

65 table_name = require_model_table_name(model) 

66 columns = tuple( 

67 PlannedColumn(column=column, name=name) 

68 for name, column in require_model_columns(model).items() 

69 ) 

70 return PlannedModel( 

71 columns=columns, 

72 indexes=_model_indexes(model, table_name, columns), 

73 model=model, 

74 table_name=table_name, 

75 ) 

76 

77 

78def build_schema_plan(models: Sequence[type[Table[Any]]]) -> SchemaPlan: 

79 """Resolve and validate model-derived schema startup metadata once.""" 

80 

81 planned_models = tuple(_plan_model(model) for model in models) 

82 table_names: set[str] = set() 

83 index_names: set[str] = set() 

84 for planned_model in planned_models: 

85 if planned_model.table_name in table_names: 

86 msg = f"duplicate table name: {planned_model.table_name!r}" 

87 raise SchemaError(msg) 

88 table_names.add(planned_model.table_name) 

89 for index in planned_model.indexes: 

90 if index.name in index_names: 

91 msg = f"duplicate index name: {index.name!r}" 

92 raise SchemaError(msg) 

93 index_names.add(index.name) 

94 return SchemaPlan(models=planned_models) 

95 

96 

97def validate_schema_policy(schema_policy: SchemaPolicy) -> None: 

98 """Reject unsupported schema policy values.""" 

99 

100 if schema_policy not in {"strict", "warn"}: 

101 msg = "schema_policy must be 'strict' or 'warn'" 

102 raise SchemaError(msg)