import six
from sqlalchemy.orm import ColumnProperty
from sqlalchemy.types import Enum as SQLAlchemyEnumType
from graphene import Argument, Enum, List
from .utils import EnumValue, to_enum_value_name, to_type_name
def _convert_sa_to_graphene_enum(sa_enum, fallback_name=None):
"""Convert the given SQLAlchemy Enum type to a Graphene Enum type.
The name of the Graphene Enum will be determined as follows:
If the SQLAlchemy Enum is based on a Python Enum, use the name
of the Python Enum. Otherwise, if the SQLAlchemy Enum is named,
use the SQL name after conversion to a type name. Otherwise, use
the given fallback_name or raise an error if it is empty.
The Enum value names are converted to upper case if necessary.
"""
if not isinstance(sa_enum, SQLAlchemyEnumType):
raise TypeError(
"Expected sqlalchemy.types.Enum, but got: {!r}".format(sa_enum)
)
enum_class = sa_enum.enum_class
if enum_class:
if all(to_enum_value_name(key) == key for key in enum_class.__members__):
return Enum.from_enum(enum_class)
name = enum_class.__name__
members = [
(to_enum_value_name(key), value.value)
for key, value in enum_class.__members__.items()
]
else:
sql_enum_name = sa_enum.name
if sql_enum_name:
name = to_type_name(sql_enum_name)
elif fallback_name:
name = fallback_name
else:
raise TypeError("No type name specified for {!r}".format(sa_enum))
members = [(to_enum_value_name(key), key) for key in sa_enum.enums]
return Enum(name, members)
def enum_for_sa_enum(sa_enum, registry):
"""Return the Graphene Enum type for the specified SQLAlchemy Enum type."""
if not isinstance(sa_enum, SQLAlchemyEnumType):
raise TypeError(
"Expected sqlalchemy.types.Enum, but got: {!r}".format(sa_enum)
)
enum = registry.get_graphene_enum_for_sa_enum(sa_enum)
if not enum:
enum = _convert_sa_to_graphene_enum(sa_enum)
registry.register_enum(sa_enum, enum)
return enum
def enum_for_field(obj_type, field_name):
"""Return the Graphene Enum type for the specified Graphene field."""
from .types import SQLAlchemyObjectType
if not isinstance(obj_type, type) or not issubclass(obj_type, SQLAlchemyObjectType):
raise TypeError(
"Expected SQLAlchemyObjectType, but got: {!r}".format(obj_type))
if not field_name or not isinstance(field_name, six.string_types):
raise TypeError(
"Expected a field name, but got: {!r}".format(field_name))
registry = obj_type._meta.registry
orm_field = registry.get_orm_field_for_graphene_field(obj_type, field_name)
if orm_field is None:
raise TypeError("Cannot get {}.{}".format(obj_type._meta.name, field_name))
if not isinstance(orm_field, ColumnProperty):
raise TypeError(
"{}.{} does not map to model column".format(obj_type._meta.name, field_name)
)
column = orm_field.columns[0]
sa_enum = column.type
if not isinstance(sa_enum, SQLAlchemyEnumType):
raise TypeError(
"{}.{} does not map to enum column".format(obj_type._meta.name, field_name)
)
enum = registry.get_graphene_enum_for_sa_enum(sa_enum)
if not enum:
fallback_name = obj_type._meta.name + to_type_name(field_name)
enum = _convert_sa_to_graphene_enum(sa_enum, fallback_name)
registry.register_enum(sa_enum, enum)
return enum
def _default_sort_enum_symbol_name(column_name, sort_asc=True):
return to_enum_value_name(column_name) + ("_ASC" if sort_asc else "_DESC")
def sort_enum_for_object_type(
obj_type, name=None, only_fields=None, only_indexed=None, get_symbol_name=None
):
"""Return Graphene Enum for sorting the given SQLAlchemyObjectType.
Parameters
- obj_type : SQLAlchemyObjectType
The object type for which the sort Enum shall be generated.
- name : str, optional, default None
Name to use for the sort Enum.
If not provided, it will be set to the object type name + 'SortEnum'
- only_fields : sequence, optional, default None
If this is set, only fields from this sequence will be considered.
- only_indexed : bool, optional, default False
If this is set, only indexed columns will be considered.
- get_symbol_name : function, optional, default None
Function which takes the column name and a boolean indicating
if the sort direction is ascending, and returns the symbol name
for the current column and sort direction. If no such function
is passed, a default function will be used that creates the symbols
'foo_asc' and 'foo_desc' for a column with the name 'foo'.
Returns
- Enum
The Graphene Enum type
"""
name = name or obj_type._meta.name + "SortEnum"
registry = obj_type._meta.registry
enum = registry.get_sort_enum_for_object_type(obj_type)
custom_options = dict(
only_fields=only_fields,
only_indexed=only_indexed,
get_symbol_name=get_symbol_name,
)
if enum:
if name != enum.__name__ or custom_options != enum.custom_options:
raise ValueError(
"Sort enum for {} has already been customized".format(obj_type)
)
else:
members = []
default = []
fields = obj_type._meta.fields
get_name = get_symbol_name or _default_sort_enum_symbol_name
for field_name in fields:
if only_fields and field_name not in only_fields:
continue
orm_field = registry.get_orm_field_for_graphene_field(obj_type, field_name)
if not isinstance(orm_field, ColumnProperty):
continue
column = orm_field.columns[0]
if only_indexed and not (column.primary_key or column.index):
continue
asc_name = get_name(column.name, True)
asc_value = EnumValue(asc_name, column.asc())
desc_name = get_name(column.name, False)
desc_value = EnumValue(desc_name, column.desc())
if column.primary_key:
default.append(asc_value)
members.extend(((asc_name, asc_value), (desc_name, desc_value)))
enum = Enum(name, members)
enum.default = default # store default as attribute
enum.custom_options = custom_options
registry.register_sort_enum(obj_type, enum)
return enum
def sort_argument_for_object_type(
obj_type,
enum_name=None,
only_fields=None,
only_indexed=None,
get_symbol_name=None,
has_default=True,
):
""""Returns Graphene Argument for sorting the given SQLAlchemyObjectType.
Parameters
- obj_type : SQLAlchemyObjectType
The object type for which the sort Argument shall be generated.
- enum_name : str, optional, default None
Name to use for the sort Enum.
If not provided, it will be set to the object type name + 'SortEnum'
- only_fields : sequence, optional, default None
If this is set, only fields from this sequence will be considered.
- only_indexed : bool, optional, default False
If this is set, only indexed columns will be considered.
- get_symbol_name : function, optional, default None
Function which takes the column name and a boolean indicating
if the sort direction is ascending, and returns the symbol name
for the current column and sort direction. If no such function
is passed, a default function will be used that creates the symbols
'foo_asc' and 'foo_desc' for a column with the name 'foo'.
- has_default : bool, optional, default True
If this is set to False, no sorting will happen when this argument is not
passed. Otherwise results will be sortied by the primary key(s) of the model.
Returns
- Enum
A Graphene Argument that accepts a list of sorting directions for the model.
"""
enum = sort_enum_for_object_type(
obj_type,
enum_name,
only_fields=only_fields,
only_indexed=only_indexed,
get_symbol_name=get_symbol_name,
)
if not has_default:
enum.default = None
return Argument(List(enum), default_value=enum.default)