|
| 1 | +from singledispatch import singledispatch |
| 2 | +from sqlalchemy import types |
| 3 | +from sqlalchemy.dialects import postgresql |
| 4 | +from sqlalchemy.orm import interfaces |
| 5 | + |
| 6 | +from graphene import (ID, Boolean, Dynamic, Enum, Field, Float, Int, List, |
| 7 | + String) |
| 8 | +from graphene.types.json import JSONString |
| 9 | + |
| 10 | +from .enums import enum_for_sa_enum |
| 11 | +from .registry import get_global_registry |
| 12 | + |
| 13 | +try: |
| 14 | + from sqlalchemy_utils import ChoiceType, JSONType, ScalarListType, TSVectorType |
| 15 | +except ImportError: |
| 16 | + ChoiceType = JSONType = ScalarListType = TSVectorType = object |
| 17 | + |
| 18 | + |
| 19 | +def get_column_doc(column): |
| 20 | + return getattr(column, "doc", None) |
| 21 | + |
| 22 | + |
| 23 | +def is_column_nullable(column): |
| 24 | + return bool(getattr(column, "nullable", True)) |
| 25 | + |
| 26 | + |
| 27 | +def convert_sqlalchemy_relationship(relationship, registry, connection_field_factory): |
| 28 | + direction = relationship.direction |
| 29 | + model = relationship.mapper.entity |
| 30 | + |
| 31 | + def dynamic_type(): |
| 32 | + _type = registry.get_type_for_model(model) |
| 33 | + if not _type: |
| 34 | + return None |
| 35 | + if direction == interfaces.MANYTOONE or not relationship.uselist: |
| 36 | + return Field(_type) |
| 37 | + elif direction in (interfaces.ONETOMANY, interfaces.MANYTOMANY): |
| 38 | + if _type._meta.connection: |
| 39 | + return connection_field_factory(relationship, registry) |
| 40 | + return Field(List(_type)) |
| 41 | + |
| 42 | + return Dynamic(dynamic_type) |
| 43 | + |
| 44 | + |
| 45 | +def convert_sqlalchemy_hybrid_method(hybrid_item): |
| 46 | + return String(description=getattr(hybrid_item, "__doc__", None), required=False) |
| 47 | + |
| 48 | + |
| 49 | +def convert_sqlalchemy_composite(composite, registry): |
| 50 | + converter = registry.get_converter_for_composite(composite.composite_class) |
| 51 | + if not converter: |
| 52 | + try: |
| 53 | + raise Exception( |
| 54 | + "Don't know how to convert the composite field %s (%s)" |
| 55 | + % (composite, composite.composite_class) |
| 56 | + ) |
| 57 | + except AttributeError: |
| 58 | + # handle fields that are not attached to a class yet (don't have a parent) |
| 59 | + raise Exception( |
| 60 | + "Don't know how to convert the composite field %r (%s)" |
| 61 | + % (composite, composite.composite_class) |
| 62 | + ) |
| 63 | + return converter(composite, registry) |
| 64 | + |
| 65 | + |
| 66 | +def _register_composite_class(cls, registry=None): |
| 67 | + if registry is None: |
| 68 | + from .registry import get_global_registry |
| 69 | + |
| 70 | + registry = get_global_registry() |
| 71 | + |
| 72 | + def inner(fn): |
| 73 | + registry.register_composite_converter(cls, fn) |
| 74 | + |
| 75 | + return inner |
| 76 | + |
| 77 | + |
| 78 | +convert_sqlalchemy_composite.register = _register_composite_class |
| 79 | + |
| 80 | + |
| 81 | +def convert_sqlalchemy_column(column, registry=None): |
| 82 | + return convert_sqlalchemy_type(getattr(column, "type", None), column, registry) |
| 83 | + |
| 84 | + |
| 85 | +@singledispatch |
| 86 | +def convert_sqlalchemy_type(type, column, registry=None): |
| 87 | + raise Exception( |
| 88 | + "Don't know how to convert the SQLAlchemy field %s (%s)" |
| 89 | + % (column, column.__class__) |
| 90 | + ) |
| 91 | + |
| 92 | + |
| 93 | +@convert_sqlalchemy_type.register(types.Date) |
| 94 | +@convert_sqlalchemy_type.register(types.Time) |
| 95 | +@convert_sqlalchemy_type.register(types.String) |
| 96 | +@convert_sqlalchemy_type.register(types.Text) |
| 97 | +@convert_sqlalchemy_type.register(types.Unicode) |
| 98 | +@convert_sqlalchemy_type.register(types.UnicodeText) |
| 99 | +@convert_sqlalchemy_type.register(postgresql.UUID) |
| 100 | +@convert_sqlalchemy_type.register(postgresql.INET) |
| 101 | +@convert_sqlalchemy_type.register(postgresql.CIDR) |
| 102 | +@convert_sqlalchemy_type.register(TSVectorType) |
| 103 | +def convert_column_to_string(type, column, registry=None): |
| 104 | + return String( |
| 105 | + description=get_column_doc(column), required=not (is_column_nullable(column)) |
| 106 | + ) |
| 107 | + |
| 108 | + |
| 109 | +@convert_sqlalchemy_type.register(types.DateTime) |
| 110 | +def convert_column_to_datetime(type, column, registry=None): |
| 111 | + from graphene.types.datetime import DateTime |
| 112 | + |
| 113 | + return DateTime( |
| 114 | + description=get_column_doc(column), required=not (is_column_nullable(column)) |
| 115 | + ) |
| 116 | + |
| 117 | + |
| 118 | +@convert_sqlalchemy_type.register(types.SmallInteger) |
| 119 | +@convert_sqlalchemy_type.register(types.Integer) |
| 120 | +def convert_column_to_int_or_id(type, column, registry=None): |
| 121 | + if column.primary_key: |
| 122 | + return ID( |
| 123 | + description=get_column_doc(column), |
| 124 | + required=not (is_column_nullable(column)), |
| 125 | + ) |
| 126 | + else: |
| 127 | + return Int( |
| 128 | + description=get_column_doc(column), |
| 129 | + required=not (is_column_nullable(column)), |
| 130 | + ) |
| 131 | + |
| 132 | + |
| 133 | +@convert_sqlalchemy_type.register(types.Boolean) |
| 134 | +def convert_column_to_boolean(type, column, registry=None): |
| 135 | + return Boolean( |
| 136 | + description=get_column_doc(column), required=not (is_column_nullable(column)) |
| 137 | + ) |
| 138 | + |
| 139 | + |
| 140 | +@convert_sqlalchemy_type.register(types.Float) |
| 141 | +@convert_sqlalchemy_type.register(types.Numeric) |
| 142 | +@convert_sqlalchemy_type.register(types.BigInteger) |
| 143 | +def convert_column_to_float(type, column, registry=None): |
| 144 | + return Float( |
| 145 | + description=get_column_doc(column), required=not (is_column_nullable(column)) |
| 146 | + ) |
| 147 | + |
| 148 | + |
| 149 | +@convert_sqlalchemy_type.register(types.Enum) |
| 150 | +def convert_enum_to_enum(type, column, registry=None): |
| 151 | + return Field( |
| 152 | + lambda: enum_for_sa_enum(type, registry or get_global_registry()), |
| 153 | + description=get_column_doc(column), |
| 154 | + required=not (is_column_nullable(column)), |
| 155 | + ) |
| 156 | + |
| 157 | + |
| 158 | +@convert_sqlalchemy_type.register(ChoiceType) |
| 159 | +def convert_choice_to_enum(type, column, registry=None): |
| 160 | + name = "{}_{}".format(column.table.name, column.name).upper() |
| 161 | + return Enum(name, type.choices, description=get_column_doc(column)) |
| 162 | + |
| 163 | + |
| 164 | +@convert_sqlalchemy_type.register(ScalarListType) |
| 165 | +def convert_scalar_list_to_list(type, column, registry=None): |
| 166 | + return List(String, description=get_column_doc(column)) |
| 167 | + |
| 168 | + |
| 169 | +@convert_sqlalchemy_type.register(postgresql.ARRAY) |
| 170 | +def convert_postgres_array_to_list(_type, column, registry=None): |
| 171 | + graphene_type = convert_sqlalchemy_type(column.type.item_type, column) |
| 172 | + inner_type = type(graphene_type) |
| 173 | + return List( |
| 174 | + inner_type, |
| 175 | + description=get_column_doc(column), |
| 176 | + required=not (is_column_nullable(column)), |
| 177 | + ) |
| 178 | + |
| 179 | + |
| 180 | +@convert_sqlalchemy_type.register(postgresql.HSTORE) |
| 181 | +@convert_sqlalchemy_type.register(postgresql.JSON) |
| 182 | +@convert_sqlalchemy_type.register(postgresql.JSONB) |
| 183 | +def convert_json_to_string(type, column, registry=None): |
| 184 | + return JSONString( |
| 185 | + description=get_column_doc(column), required=not (is_column_nullable(column)) |
| 186 | + ) |
| 187 | + |
| 188 | + |
| 189 | +@convert_sqlalchemy_type.register(JSONType) |
| 190 | +def convert_json_type_to_string(type, column, registry=None): |
| 191 | + return JSONString( |
| 192 | + description=get_column_doc(column), required=not (is_column_nullable(column)) |
| 193 | + ) |
0 commit comments