Coverage for /pythoncovmergedfiles/medio/medio/usr/local/lib/python3.8/site-packages/marshmallow_sqlalchemy/schema.py: 32%

1from marshmallow.fields import Field 

2from marshmallow.schema import Schema, SchemaMeta, SchemaOpts 

3import sqlalchemy as sa 

4from sqlalchemy.ext.declarative import DeclarativeMeta 

5 

6from .convert import ModelConverter 

7from .exceptions import IncorrectSchemaTypeError 

8from .load_instance_mixin import LoadInstanceMixin 

9 

10 

11# This isn't really a field; it's a placeholder for the metaclass. 

12# This should be considered private API. 

13class SQLAlchemyAutoField(Field): 

14 def __init__(self, *, column_name=None, model=None, table=None, field_kwargs): 

15 super().__init__() 

16 

17 if model and table: 

18 raise ValueError("Cannot pass both `model` and `table` options.") 

19 

20 self.column_name = column_name 

21 self.model = model 

22 self.table = table 

23 self.field_kwargs = field_kwargs 

24 

25 def create_field(self, schema_opts, column_name, converter): 

26 model = self.model or schema_opts.model 

27 if model: 

28 return converter.field_for(model, column_name, **self.field_kwargs) 

29 else: 

30 table = self.table if self.table is not None else schema_opts.table 

31 column = getattr(table.columns, column_name) 

32 return converter.column2field(column, **self.field_kwargs) 

33 

34 # This field should never be bound to a schema. 

35 # If this method is called, it's probably because the schema is not a SQLAlchemySchema. 

36 def _bind_to_schema(self, field_name, schema): 

37 raise IncorrectSchemaTypeError( 

38 f"Cannot bind SQLAlchemyAutoField. Make sure that {schema} is a SQLAlchemySchema or SQLAlchemyAutoSchema." 

39 ) 

40 

41 

42class SQLAlchemySchemaOpts(LoadInstanceMixin.Opts, SchemaOpts): 

43 """Options class for `SQLAlchemySchema`. 

44 Adds the following options: 

45 

46 - ``model``: The SQLAlchemy model to generate the `Schema` from (mutually exclusive with ``table``). 

47 - ``table``: The SQLAlchemy table to generate the `Schema` from (mutually exclusive with ``model``). 

48 - ``load_instance``: Whether to load model instances. 

49 - ``sqla_session``: SQLAlchemy session to be used for deserialization. 

50 This is only needed when ``load_instance`` is `True`. You can also pass a session to the Schema's `load` method. 

51 - ``transient``: Whether to load model instances in a transient state (effectively ignoring the session). 

52 Only relevant when ``load_instance`` is `True`. 

53 - ``model_converter``: `ModelConverter` class to use for converting the SQLAlchemy model to marshmallow fields. 

54 """ 

55 

56 def __init__(self, meta, *args, **kwargs): 

57 super().__init__(meta, *args, **kwargs) 

58 

59 self.model = getattr(meta, "model", None) 

60 self.table = getattr(meta, "table", None) 

61 if self.model is not None and self.table is not None: 

62 raise ValueError("Cannot set both `model` and `table` options.") 

63 self.model_converter = getattr(meta, "model_converter", ModelConverter) 

64 

65 

66class SQLAlchemyAutoSchemaOpts(SQLAlchemySchemaOpts): 

67 """Options class for `SQLAlchemyAutoSchema`. 

68 Has the same options as `SQLAlchemySchemaOpts`, with the addition of: 

69 

70 - ``include_fk``: Whether to include foreign fields; defaults to `False`. 

71 - ``include_relationships``: Whether to include relationships; defaults to `False`. 

72 """ 

73 

74 def __init__(self, meta, *args, **kwargs): 

75 super().__init__(meta, *args, **kwargs) 

76 self.include_fk = getattr(meta, "include_fk", False) 

77 self.include_relationships = getattr(meta, "include_relationships", False) 

78 if self.table is not None and self.include_relationships: 

79 raise ValueError("Cannot set `table` and `include_relationships = True`.") 

80 

81 

82class SQLAlchemySchemaMeta(SchemaMeta): 

83 @classmethod 

84 def get_declared_fields(mcs, klass, cls_fields, inherited_fields, dict_cls): 

85 opts = klass.opts 

86 Converter = opts.model_converter 

87 converter = Converter(schema_cls=klass) 

88 fields = super().get_declared_fields( 

89 klass, cls_fields, inherited_fields, dict_cls 

90 ) 

91 fields.update(mcs.get_declared_sqla_fields(fields, converter, opts, dict_cls)) 

92 fields.update(mcs.get_auto_fields(fields, converter, opts, dict_cls)) 

93 return fields 

94 

95 @classmethod 

96 def get_declared_sqla_fields(mcs, base_fields, converter, opts, dict_cls): 

97 return {} 

98 

99 @classmethod 

100 def get_auto_fields(mcs, fields, converter, opts, dict_cls): 

101 return dict_cls( 

102 { 

103 field_name: field.create_field( 

104 opts, field.column_name or field_name, converter 

105 ) 

106 for field_name, field in fields.items() 

107 if isinstance(field, SQLAlchemyAutoField) 

108 and field_name not in opts.exclude 

109 } 

110 ) 

111 

112 

113class SQLAlchemyAutoSchemaMeta(SQLAlchemySchemaMeta): 

114 @classmethod 

115 def get_declared_sqla_fields(cls, base_fields, converter, opts, dict_cls): 

116 fields = dict_cls() 

117 if opts.table is not None: 

118 fields.update( 

119 converter.fields_for_table( 

120 opts.table, 

121 fields=opts.fields, 

122 exclude=opts.exclude, 

123 include_fk=opts.include_fk, 

124 base_fields=base_fields, 

125 dict_cls=dict_cls, 

126 ) 

127 ) 

128 elif opts.model is not None: 

129 fields.update( 

130 converter.fields_for_model( 

131 opts.model, 

132 fields=opts.fields, 

133 exclude=opts.exclude, 

134 include_fk=opts.include_fk, 

135 include_relationships=opts.include_relationships, 

136 base_fields=base_fields, 

137 dict_cls=dict_cls, 

138 ) 

139 ) 

140 return fields 

141 

142 

143class SQLAlchemySchema( 

144 LoadInstanceMixin.Schema, Schema, metaclass=SQLAlchemySchemaMeta 

145): 

146 """Schema for a SQLAlchemy model or table. 

147 Use together with `auto_field` to generate fields from columns. 

148 

149 Example: :: 

150 

151 from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field 

152 

153 from mymodels import User 

154 

155 class UserSchema(SQLAlchemySchema): 

156 class Meta: 

157 model = User 

158 

159 id = auto_field() 

160 created_at = auto_field(dump_only=True) 

161 name = auto_field() 

162 """ 

163 

164 OPTIONS_CLASS = SQLAlchemySchemaOpts 

165 

166 

167class SQLAlchemyAutoSchema(SQLAlchemySchema, metaclass=SQLAlchemyAutoSchemaMeta): 

168 """Schema that automatically generates fields from the columns of 

169 a SQLAlchemy model or table. 

170 

171 Example: :: 

172 

173 from marshmallow_sqlalchemy import SQLAlchemyAutoSchema, auto_field 

174 

175 from mymodels import User 

176 

177 class UserSchema(SQLAlchemyAutoSchema): 

178 class Meta: 

179 model = User 

180 # OR 

181 # table = User.__table__ 

182 

183 created_at = auto_field(dump_only=True) 

184 """ 

185 

186 OPTIONS_CLASS = SQLAlchemyAutoSchemaOpts 

187 

188 

189def auto_field( 

190 column_name: str = None, 

191 *, 

192 model: DeclarativeMeta = None, 

193 table: sa.Table = None, 

194 **kwargs, 

195): 

196 """Mark a field to autogenerate from a model or table. 

197 

198 :param column_name: Name of the column to generate the field from. 

199 If ``None``, matches the field name. If ``attribute`` is unspecified, 

200 ``attribute`` will be set to the same value as ``column_name``. 

201 :param model: Model to generate the field from. 

202 If ``None``, uses ``model`` specified on ``class Meta``. 

203 :param table: Table to generate the field from. 

204 If ``None``, uses ``table`` specified on ``class Meta``. 

205 :param kwargs: Field argument overrides. 

206 """ 

207 if column_name is not None: 

208 kwargs.setdefault("attribute", column_name) 

209 return SQLAlchemyAutoField( 

210 column_name=column_name, model=model, table=table, field_kwargs=kwargs 

211 )