from marshmallow.fields import Field from marshmallow.schema import Schema, SchemaMeta, SchemaOpts import sqlalchemy as sa from sqlalchemy.ext.declarative import DeclarativeMeta from .convert import ModelConverter from .exceptions import IncorrectSchemaTypeError from .load_instance_mixin import LoadInstanceMixin # This isn't really a field; it's a placeholder for the metaclass. # This should be considered private API. class SQLAlchemyAutoField(Field): def __init__(self, *, column_name=None, model=None, table=None, field_kwargs): super().__init__() if model and table: raise ValueError("Cannot pass both `model` and `table` options.") self.column_name = column_name self.model = model self.table = table self.field_kwargs = field_kwargs def create_field(self, schema_opts, column_name, converter): model = self.model or schema_opts.model if model: return converter.field_for(model, column_name, **self.field_kwargs) else: table = self.table if self.table is not None else schema_opts.table column = getattr(table.columns, column_name) return converter.column2field(column, **self.field_kwargs) # This field should never be bound to a schema. # If this method is called, it's probably because the schema is not a SQLAlchemySchema. def _bind_to_schema(self, field_name, schema): raise IncorrectSchemaTypeError( f"Cannot bind SQLAlchemyAutoField. Make sure that {schema} is a SQLAlchemySchema or SQLAlchemyAutoSchema." ) class SQLAlchemySchemaOpts(LoadInstanceMixin.Opts, SchemaOpts): """Options class for `SQLAlchemySchema`. Adds the following options: - ``model``: The SQLAlchemy model to generate the `Schema` from (mutually exclusive with ``table``). - ``table``: The SQLAlchemy table to generate the `Schema` from (mutually exclusive with ``model``). - ``load_instance``: Whether to load model instances. - ``sqla_session``: SQLAlchemy session to be used for deserialization. This is only needed when ``load_instance`` is `True`. You can also pass a session to the Schema's `load` method. - ``transient``: Whether to load model instances in a transient state (effectively ignoring the session). Only relevant when ``load_instance`` is `True`. - ``model_converter``: `ModelConverter` class to use for converting the SQLAlchemy model to marshmallow fields. """ def __init__(self, meta, *args, **kwargs): super().__init__(meta, *args, **kwargs) self.model = getattr(meta, "model", None) self.table = getattr(meta, "table", None) if self.model is not None and self.table is not None: raise ValueError("Cannot set both `model` and `table` options.") self.model_converter = getattr(meta, "model_converter", ModelConverter) class SQLAlchemyAutoSchemaOpts(SQLAlchemySchemaOpts): """Options class for `SQLAlchemyAutoSchema`. Has the same options as `SQLAlchemySchemaOpts`, with the addition of: - ``include_fk``: Whether to include foreign fields; defaults to `False`. - ``include_relationships``: Whether to include relationships; defaults to `False`. """ def __init__(self, meta, *args, **kwargs): super().__init__(meta, *args, **kwargs) self.include_fk = getattr(meta, "include_fk", False) self.include_relationships = getattr(meta, "include_relationships", False) if self.table is not None and self.include_relationships: raise ValueError("Cannot set `table` and `include_relationships = True`.") class SQLAlchemySchemaMeta(SchemaMeta): @classmethod def get_declared_fields(mcs, klass, cls_fields, inherited_fields, dict_cls): opts = klass.opts Converter = opts.model_converter converter = Converter(schema_cls=klass) fields = super().get_declared_fields( klass, cls_fields, inherited_fields, dict_cls ) fields.update(mcs.get_declared_sqla_fields(fields, converter, opts, dict_cls)) fields.update(mcs.get_auto_fields(fields, converter, opts, dict_cls)) return fields @classmethod def get_declared_sqla_fields(mcs, base_fields, converter, opts, dict_cls): return {} @classmethod def get_auto_fields(mcs, fields, converter, opts, dict_cls): return dict_cls( { field_name: field.create_field( opts, field.column_name or field_name, converter ) for field_name, field in fields.items() if isinstance(field, SQLAlchemyAutoField) and field_name not in opts.exclude } ) class SQLAlchemyAutoSchemaMeta(SQLAlchemySchemaMeta): @classmethod def get_declared_sqla_fields(cls, base_fields, converter, opts, dict_cls): fields = dict_cls() if opts.table is not None: fields.update( converter.fields_for_table( opts.table, fields=opts.fields, exclude=opts.exclude, include_fk=opts.include_fk, base_fields=base_fields, dict_cls=dict_cls, ) ) elif opts.model is not None: fields.update( converter.fields_for_model( opts.model, fields=opts.fields, exclude=opts.exclude, include_fk=opts.include_fk, include_relationships=opts.include_relationships, base_fields=base_fields, dict_cls=dict_cls, ) ) return fields class SQLAlchemySchema( LoadInstanceMixin.Schema, Schema, metaclass=SQLAlchemySchemaMeta ): """Schema for a SQLAlchemy model or table. Use together with `auto_field` to generate fields from columns. Example: :: from marshmallow_sqlalchemy import SQLAlchemySchema, auto_field from mymodels import User class UserSchema(SQLAlchemySchema): class Meta: model = User id = auto_field() created_at = auto_field(dump_only=True) name = auto_field() """ OPTIONS_CLASS = SQLAlchemySchemaOpts class SQLAlchemyAutoSchema(SQLAlchemySchema, metaclass=SQLAlchemyAutoSchemaMeta): """Schema that automatically generates fields from the columns of a SQLAlchemy model or table. Example: :: from marshmallow_sqlalchemy import SQLAlchemyAutoSchema, auto_field from mymodels import User class UserSchema(SQLAlchemyAutoSchema): class Meta: model = User # OR # table = User.__table__ created_at = auto_field(dump_only=True) """ OPTIONS_CLASS = SQLAlchemyAutoSchemaOpts def auto_field( column_name: str = None, *, model: DeclarativeMeta = None, table: sa.Table = None, **kwargs, ): """Mark a field to autogenerate from a model or table. :param column_name: Name of the column to generate the field from. If ``None``, matches the field name. If ``attribute`` is unspecified, ``attribute`` will be set to the same value as ``column_name``. :param model: Model to generate the field from. If ``None``, uses ``model`` specified on ``class Meta``. :param table: Table to generate the field from. If ``None``, uses ``table`` specified on ``class Meta``. :param kwargs: Field argument overrides. """ if column_name is not None: kwargs.setdefault("attribute", column_name) return SQLAlchemyAutoField( column_name=column_name, model=model, table=table, field_kwargs=kwargs )