redis-om-python/redis_developer/orm/model.py

import abc
import dataclasses
import decimal
import logging
import operator
import re
from copy import deepcopy
from enum import Enum
from functools import reduce
from typing import (
    AbstractSet,
    Any,
    Callable,
    Dict,
    Mapping,
    Optional,
    Set,
    Tuple,
    TypeVar,
    Union,
    Sequence,
    no_type_check,
    Protocol,
    List,
    Type,
    Pattern, get_origin, get_args
)
import uuid

import redis
from pydantic import BaseModel, validator
from pydantic.fields import FieldInfo as PydanticFieldInfo
from pydantic.fields import ModelField, Undefined, UndefinedType
from pydantic.main import ModelMetaclass
from pydantic.typing import NoArgAnyCallable, resolve_annotations
from pydantic.utils import Representation

from .encoders import jsonable_encoder

model_registry = {}
_T = TypeVar("_T")
log = logging.getLogger(__name__)


class TokenEscaper:
    """
    Escape punctuation within an input string.
    """

    # Characters that RediSearch requires us to escape during queries.
    # Source: https://oss.redis.com/redisearch/Escaping/#the_rules_of_text_field_tokenization
    DEFAULT_ESCAPED_CHARS = r"[,.<>{}\[\]\\\"\':;!@#$%^&*()\-+=~\ ]"

    def __init__(self, escape_chars_re: Optional[Pattern] = None):
        if escape_chars_re:
            self.escaped_chars_re = escape_chars_re
        else:
            self.escaped_chars_re = re.compile(self.DEFAULT_ESCAPED_CHARS)

    def escape(self, string):
        def escape_symbol(match):
            value = match.group(0)
            return f"\\{value}"

        return self.escaped_chars_re.sub(escape_symbol, string)


escaper = TokenEscaper()


class RedisModelError(Exception):
    pass


class NotFoundError(Exception):
    """A query found no results."""


class Operators(Enum):
    EQ = 1
    NE = 2
    LT = 3
    LE = 4
    GT = 5
    GE = 6
    OR = 7
    AND = 8
    NOT = 9
    IN = 10
    NOT_IN = 11
    LIKE = 12
    ALL = 13


@dataclasses.dataclass
class NegatedExpression:
    expression: 'Expression'

    def __invert__(self):
        return self.expression

    def __and__(self, other):
        return Expression(left=self, op=Operators.AND, right=other)

    def __or__(self, other):
        return Expression(left=self, op=Operators.OR, right=other)


@dataclasses.dataclass
class Expression:
    op: Operators
    left: Any
    right: Any

    def __invert__(self):
        return NegatedExpression(self)

    def __and__(self, other):
        return Expression(left=self, op=Operators.AND, right=other)

    def __or__(self, other):
        return Expression(left=self, op=Operators.OR, right=other)


ExpressionOrNegated = Union[Expression, NegatedExpression]


class ExpressionProxy:
    def __init__(self, field: ModelField):
        self.field = field

    def __eq__(self, other: Any) -> Expression:
        return Expression(left=self.field, op=Operators.EQ, right=other)

    def __ne__(self, other: Any) -> Expression:
        return Expression(left=self.field, op=Operators.NE, right=other)

    def __lt__(self, other: Any) -> Expression:
        return Expression(left=self.field, op=Operators.LT, right=other)

    def __le__(self, other: Any) -> Expression:
        return Expression(left=self.field, op=Operators.LE, right=other)

    def __gt__(self, other: Any) -> Expression:
        return Expression(left=self.field, op=Operators.GT, right=other)

    def __ge__(self, other: Any) -> Expression:
        return Expression(left=self.field, op=Operators.GE, right=other)


class QueryNotSupportedError(Exception):
    """The attempted query is not supported."""


class RediSearchFieldTypes(Enum):
    TEXT = 'TEXT'
    TAG = 'TAG'
    NUMERIC = 'NUMERIC'
    GEO = 'GEO'


# TODO: How to handle Geo fields?
NUMERIC_TYPES = (float, int, decimal.Decimal)
DEFAULT_PAGE_SIZE = 10


class FindQuery:
    def __init__(self,
                 expressions: Sequence[ExpressionOrNegated],
                 model: Type['RedisModel'],
                 offset: int = 0,
                 limit: int = DEFAULT_PAGE_SIZE,
                 page_size: int = DEFAULT_PAGE_SIZE,
                 sort_fields: Optional[List[str]] = None):
        self.expressions = expressions
        self.model = model
        self.offset = offset
        self.limit = limit
        self.page_size = page_size

        if sort_fields:
            self.sort_fields = self.validate_sort_fields(sort_fields)
        else:
            self.sort_fields = []

        self._expression = None
        self._query = None
        self._pagination = []
        self._model_cache = []

    @property
    def pagination(self):
        if self._pagination:
            return self._pagination
        self._pagination = self.resolve_redisearch_pagination()
        return self._pagination

    @property
    def expression(self):
        if self._expression:
            return self._expression
        if self.expressions:
            self._expression = reduce(operator.and_, self.expressions)
        else:
            self._expression = Expression(left=None, right=None, op=Operators.ALL)
        return self._expression

    @property
    def query(self):
        return self.resolve_redisearch_query(self.expression)

    def validate_sort_fields(self, sort_fields):
        for sort_field in sort_fields:
            field_name = sort_field.lstrip("-")
            if field_name not in self.model.__fields__:
                raise QueryNotSupportedError(f"You tried sort by {field_name}, but that field "
                                             f"does not exist on the model {self.model}")
            field_proxy = getattr(self.model, field_name)
            if not getattr(field_proxy.field.field_info, 'sortable', False):
                raise QueryNotSupportedError(f"You tried sort by {field_name}, but {self.model} does "
                                             "not define that field as sortable. See docs: XXX")
        return sort_fields

    @staticmethod
    def resolve_field_type(field: ModelField) -> RediSearchFieldTypes:
        if getattr(field.field_info, 'primary_key', None) is True:
            return RediSearchFieldTypes.TAG
        elif getattr(field.field_info, 'full_text_search', None) is True:
            return RediSearchFieldTypes.TEXT

        field_type = field.outer_type_

        # TODO: GEO
        if any(issubclass(field_type, t) for t in NUMERIC_TYPES):
            return RediSearchFieldTypes.NUMERIC
        else:
            # TAG fields are the default field type.
            # TODO: A ListField or ArrayField that supports multiple values
            #  and contains logic.
            return RediSearchFieldTypes.TAG

    @staticmethod
    def expand_tag_value(value):
        err = RedisModelError(f"Using the IN operator requires passing a sequence of "
                              "possible values. You passed: {value}")
        if isinstance(str, value):
            raise err
        try:
            expanded_value = "|".join([escaper.escape(v) for v in value])
        except TypeError:
            raise err
        return expanded_value

    @classmethod
    def resolve_value(cls, field_name: str, field_type: RediSearchFieldTypes,
                      op: Operators, value: Any) -> str:
        result = ""
        if field_type is RediSearchFieldTypes.TEXT:
            result = f"@{field_name}:"
            if op is Operators.EQ:
                result += f'"{value}"'
            elif op is Operators.NE:
                result = f'-({result}"{value}")'
            elif op is Operators.LIKE:
                result += value
            else:
                raise QueryNotSupportedError("Only equals (=), not-equals (!=), and like() "
                                             "comparisons are supported for TEXT fields. See "
                                             "docs: TODO.")
        elif field_type is RediSearchFieldTypes.NUMERIC:
            if op is Operators.EQ:
                result += f"@{field_name}:[{value} {value}]"
            elif op is Operators.NE:
                # TODO: Is this enough or do we also need a clause for all values
                #  ([-inf +inf]) from which we then subtract the undesirable value?
                result += f"-(@{field_name}:[{value} {value}])"
            elif op is Operators.GT:
                result += f"@{field_name}:[({value} +inf]"
            elif op is Operators.LT:
                result += f"@{field_name}:[-inf ({value}]"
            elif op is Operators.GE:
                result += f"@{field_name}:[{value} +inf]"
            elif op is Operators.LE:
                result += f"@{field_name}:[-inf {value}]"
        elif field_type is RediSearchFieldTypes.TAG:
            if op is Operators.EQ:
                value = escaper.escape(value)
                result += f"@{field_name}:{{{value}}}"
            elif op is Operators.NE:
                value = escaper.escape(value)
                result += f"-(@{field_name}:{{{value}}})"
            elif op is Operators.IN:
                expanded_value = cls.expand_tag_value(value)
                result += f"(@{field_name}:{{{expanded_value}}})"
            elif op is Operators.NOT_IN:
                expanded_value = cls.expand_tag_value(value)
                result += f"-(@{field_name}:{{{expanded_value}}})"

        return result

    def resolve_redisearch_pagination(self):
        """Resolve pagination options for a query."""
        return ["LIMIT", self.offset, self.limit]

    def resolve_redisearch_sort_fields(self):
        """Resolve sort options for a query."""
        if not self.sort_fields:
            return
        fields = []
        for f in self.sort_fields:
            direction = "desc" if f.startswith('-') else 'asc'
            fields.extend([f.lstrip('-'), direction])
        if self.sort_fields:
            return ["SORTBY", *fields]

    @classmethod
    def resolve_redisearch_query(cls, expression: ExpressionOrNegated):
        """Resolve an expression to a string RediSearch query."""
        field_type = None
        field_name = None
        encompassing_expression_is_negated = False
        result = ""

        if isinstance(expression, NegatedExpression):
            encompassing_expression_is_negated = True
            expression = expression.expression

        if expression.op is Operators.ALL:
            if encompassing_expression_is_negated:
                # TODO: Is there a use case for this, perhaps for dynamic
                # scoring purposes?
                raise QueryNotSupportedError("You cannot negate a query for all results.")
            return "*"

        if isinstance(expression.left, Expression) or \
                isinstance(expression.left, NegatedExpression):
            result += f"({cls.resolve_redisearch_query(expression.left)})"
        elif isinstance(expression.left, ModelField):
            field_type = cls.resolve_field_type(expression.left)
            field_name = expression.left.name
        else:
            raise QueryNotSupportedError(f"A query expression should start with either a field "
                                         f"or an expression enclosed in parenthesis. See docs: "
                                         f"TODO")

        right = expression.right

        if isinstance(right, Expression) or isinstance(right, NegatedExpression):
            if expression.op == Operators.AND:
                result += " "
            elif expression.op == Operators.OR:
                result += "| "
            else:
                raise QueryNotSupportedError("You can only combine two query expressions with"
                                             "AND (&) or OR (|). See docs: TODO")

            if isinstance(right, NegatedExpression):
                result += "-"
                # We're handling the RediSearch operator in this call ("-"), so resolve the
                # inner expression instead of the NegatedExpression.
                right = right.expression

            result += f"({cls.resolve_redisearch_query(right)})"
        else:
            if isinstance(right, ModelField):
                raise QueryNotSupportedError("Comparing fields is not supported. See docs: TODO")
            else:
                # TODO: Optionals causing IDE errors here
                result += cls.resolve_value(field_name, field_type, expression.op, right)

        if encompassing_expression_is_negated:
            result = f"-({result})"

        return result

    def execute(self, exhaust_results=True):
        args = ["ft.search", self.model.Meta.index_name, self.query, *self.pagination]
        if self.sort_fields:
            args += self.resolve_redisearch_sort_fields()

        # Reset the cache if we're executing from offset 0.
        if self.offset == 0:
            self._model_cache.clear()

        # If the offset is greater than 0, we're paginating through a result set,
        # so append the new results to results already in the cache.
        raw_result = self.model.db().execute_command(*args)
        count = raw_result[0]
        results = self.model.from_redis(raw_result)
        self._model_cache += results

        if not exhaust_results:
            return self._model_cache

        # The query returned all results, so we have no more work to do.
        if count <= len(results):
            return self._model_cache

        # Transparently (to the user) make subsequent requests to paginate
        # through the results and finally return them all.
        query = self
        while True:
            # Make a query for each pass of the loop, with a new offset equal to the
            # current offset plus `page_size`, until we stop getting results back.
            query = FindQuery(expressions=query.expressions,
                              model=query.model,
                              offset=query.offset + query.page_size,
                              page_size=query.page_size,
                              limit=query.limit)
            _results = query.execute(exhaust_results=False)
            if not _results:
                break
            self._model_cache += _results
        return self._model_cache

    def first(self):
        query = FindQuery(expressions=self.expressions, model=self.model,
                          offset=0, limit=1, sort_fields=self.sort_fields)
        return query.execute()[0]

    def all(self, batch_size=10):
        if batch_size != self.page_size:
            # TODO: There's probably a copy-with-change mechanism in Pydantic,
            #  or can we use one from dataclasses?
            query = FindQuery(expressions=self.expressions,
                              model=self.model,
                              offset=self.offset,
                              page_size=batch_size,
                              limit=batch_size,
                              sort_fields=self.sort_fields)
            return query.execute()
        return self.execute()

    def sort_by(self, *fields: str):
        if not fields:
            return self
        return FindQuery(expressions=self.expressions,
                         model=self.model,
                         offset=self.offset,
                         page_size=self.page_size,
                         limit=self.limit,
                         sort_fields=list(fields))

    def update(self, **kwargs):
        """Update all matching records in this query."""
        # TODO

    def delete(cls, **field_values):
        """Delete all matching records in this query."""
        for field_name, value in field_values:
            valid_attr = hasattr(cls.model, field_name)
            if not valid_attr:
                raise RedisModelError(f"Can't update field {field_name} because "
                                      f"the field does not exist on the model {cls}")

        return cls

    def __iter__(self):
        if self._model_cache:
            for m in self._model_cache:
                yield m
        else:
            for m in self.execute():
                yield m

    def __getitem__(self, item: int):
        """
        Given this code:
            Model.find()[1000]

        We should return only the 1000th result.

            1. If the result is loaded in the query cache for this query,
               we can return it directly from the cache.

            2. If the query cache does not have enough elements to return
               that result, then we should clone the current query and
               give it a new offset and limit: offset=n, limit=1.
        """
        if self._model_cache and len(self._model_cache) >= item:
            return self._model_cache[item]

        query = FindQuery(expressions=self.expressions,
                          model=self.model,
                          offset=item,
                          sort_fields=self.sort_fields,
                          limit=1)

        return query.execute()[0]


class PrimaryKeyCreator(Protocol):
    def create_pk(self, *args, **kwargs) -> str:
        """Create a new primary key"""


class Uuid4PrimaryKey:
    def create_pk(self, *args, **kwargs) -> str:
        return str(uuid.uuid4())


def __dataclass_transform__(
    *,
    eq_default: bool = True,
    order_default: bool = False,
    kw_only_default: bool = False,
    field_descriptors: Tuple[Union[type, Callable[..., Any]], ...] = (()),
) -> Callable[[_T], _T]:
    return lambda a: a


class FieldInfo(PydanticFieldInfo):
    def __init__(self, default: Any = Undefined, **kwargs: Any) -> None:
        primary_key = kwargs.pop("primary_key", False)
        sortable = kwargs.pop("sortable", Undefined)
        index = kwargs.pop("index", Undefined)
        full_text_search = kwargs.pop("full_text_search", Undefined)
        super().__init__(default=default, **kwargs)
        self.primary_key = primary_key
        self.sortable = sortable
        self.index = index
        self.full_text_search = full_text_search


class RelationshipInfo(Representation):
    def __init__(
        self,
        *,
        back_populates: Optional[str] = None,
        link_model: Optional[Any] = None,
    ) -> None:
        self.back_populates = back_populates
        self.link_model = link_model


def Field(
    default: Any = Undefined,
    *,
    default_factory: Optional[NoArgAnyCallable] = None,
    alias: str = None,
    title: str = None,
    description: str = None,
    exclude: Union[
        AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
    ] = None,
    include: Union[
        AbstractSet[Union[int, str]], Mapping[Union[int, str], Any], Any
    ] = None,
    const: bool = None,
    gt: float = None,
    ge: float = None,
    lt: float = None,
    le: float = None,
    multiple_of: float = None,
    min_items: int = None,
    max_items: int = None,
    min_length: int = None,
    max_length: int = None,
    allow_mutation: bool = True,
    regex: str = None,
    primary_key: bool = False,
    sortable: Union[bool, UndefinedType] = Undefined,
    index: Union[bool, UndefinedType] = Undefined,
    full_text_search: Union[bool, UndefinedType] = Undefined,
    schema_extra: Optional[Dict[str, Any]] = None,
) -> Any:
    current_schema_extra = schema_extra or {}
    field_info = FieldInfo(
        default,
        default_factory=default_factory,
        alias=alias,
        title=title,
        description=description,
        exclude=exclude,
        include=include,
        const=const,
        gt=gt,
        ge=ge,
        lt=lt,
        le=le,
        multiple_of=multiple_of,
        min_items=min_items,
        max_items=max_items,
        min_length=min_length,
        max_length=max_length,
        allow_mutation=allow_mutation,
        regex=regex,
        primary_key=primary_key,
        sortable=sortable,
        index=index,
        full_text_search=full_text_search,
        **current_schema_extra,
    )
    field_info._validate()
    return field_info


@dataclasses.dataclass
class PrimaryKey:
    name: str
    field: ModelField


class DefaultMeta:
    # TODO: Should this really be optional here?
    global_key_prefix: Optional[str] = None
    model_key_prefix: Optional[str] = None
    primary_key_pattern: Optional[str] = None
    database: Optional[redis.Redis] = None
    primary_key: Optional[PrimaryKey] = None
    primary_key_creator_cls: Optional[Type[PrimaryKeyCreator]] = None
    index_name: Optional[str] = None
    abstract: Optional[bool] = False


class ModelMeta(ModelMetaclass):
    def __new__(cls, name, bases, attrs, **kwargs):  # noqa C901
        meta = attrs.pop('Meta', None)
        new_class = super().__new__(cls, name, bases, attrs, **kwargs)
        meta = meta or getattr(new_class, 'Meta', None)
        base_meta = getattr(new_class, '_meta', None)

        if meta and meta != DefaultMeta and meta != base_meta:
            new_class.Meta = meta
            new_class._meta = meta
        elif base_meta:
            new_class._meta = deepcopy(base_meta)
            new_class.Meta = new_class._meta
            # Unset inherited values we don't want to reuse (typically based on the model name).
            new_class._meta.abstract = False
            new_class._meta.model_key_prefix = None
            new_class._meta.index_name = None
        else:
            new_class._meta = deepcopy(DefaultMeta)
            new_class.Meta = new_class._meta

        # Create proxies for each model field so that we can use the field
        # in queries, like Model.get(Model.field_name == 1)
        for field_name, field in new_class.__fields__.items():
            setattr(new_class, field_name, ExpressionProxy(field))
            # Check if this is our FieldInfo version with extended ORM metadata.
            if isinstance(field.field_info, FieldInfo):
                if field.field_info.primary_key:
                    new_class._meta.primary_key = PrimaryKey(name=field_name, field=field)

        if not getattr(new_class._meta, 'global_key_prefix', None):
            new_class._meta.global_key_prefix = getattr(base_meta, "global_key_prefix", "")
        if not getattr(new_class._meta, 'model_key_prefix', None):
            # Don't look at the base class for this.
            new_class._meta.model_key_prefix = f"{new_class.__module__}.{new_class.__name__}"
        if not getattr(new_class._meta, 'primary_key_pattern', None):
            new_class._meta.primary_key_pattern = getattr(base_meta, "primary_key_pattern",
                                                          "{pk}")
        if not getattr(new_class._meta, 'database', None):
            new_class._meta.database = getattr(base_meta, "database",
                                               redis.Redis(decode_responses=True))
        if not getattr(new_class._meta, 'primary_key_creator_cls', None):
            new_class._meta.primary_key_creator_cls = getattr(base_meta, "primary_key_creator_cls",
                                                              Uuid4PrimaryKey)
        if not getattr(new_class._meta, 'index_name', None):
            new_class._meta.index_name = f"{new_class._meta.global_key_prefix}:" \
                                         f"{new_class._meta.model_key_prefix}:index"

        return new_class


class RedisModel(BaseModel, abc.ABC, metaclass=ModelMeta):
    pk: Optional[str] = Field(default=None, primary_key=True)

    Meta = DefaultMeta
    # TODO: Missing _meta here is causing IDE warnings.

    class Config:
        orm_mode = True
        arbitrary_types_allowed = True
        extra = 'allow'

    def __init__(__pydantic_self__, **data: Any) -> None:
        super().__init__(**data)
        __pydantic_self__.validate_primary_key()

    def __lt__(self, other):
        """Default sort: compare all shared model fields."""
        my_keys = set(self.__fields__.keys())
        other_keys = set(other.__fields__.keys())
        shared_keys = list(my_keys & other_keys)
        lt = [getattr(self, k) < getattr(other, k) for k in shared_keys]
        return len(lt) > len(shared_keys) / 2

    @validator("pk", always=True)
    def validate_pk(cls, v):
        if not v:
            v = cls._meta.primary_key_creator_cls().create_pk()
        return v

    @classmethod
    def validate_primary_key(cls):
        """Check for a primary key. We need one (and only one)."""
        primary_keys = 0
        for name, field in cls.__fields__.items():
            if getattr(field.field_info, 'primary_key', None):
                primary_keys += 1
        if primary_keys == 0:
            raise RedisModelError("You must define a primary key for the model")
        elif primary_keys > 1:
            raise RedisModelError("You must define only one primary key for a model")

    @classmethod
    def make_key(cls, part: str):
        global_prefix = getattr(cls._meta, 'global_key_prefix', '').strip(":")
        model_prefix = getattr(cls._meta, 'model_key_prefix', '').strip(":")
        return f"{global_prefix}:{model_prefix}:{part}"

    @classmethod
    def make_primary_key(cls, pk: Any):
        """Return the Redis key for this model."""
        return cls.make_key(cls._meta.primary_key_pattern.format(pk=pk))

    def key(self):
        """Return the Redis key for this model."""
        pk = getattr(self, self._meta.primary_key.field.name)
        return self.make_primary_key(pk)

    @classmethod
    def db(cls):
        return cls._meta.database

    @classmethod
    def find(cls, *expressions: Union[Any, Expression]):  # TODO: How to type annotate this?
        return FindQuery(expressions=expressions, model=cls)

    @classmethod
    def from_redis(cls, res: Any):
        # TODO: Parsing logic borrowed from redisearch-py. Evaluate.
        import six
        from six.moves import xrange, zip as izip

        def to_string(s):
            if isinstance(s, six.string_types):
                return s
            elif isinstance(s, six.binary_type):
                return s.decode('utf-8', 'ignore')
            else:
                return s  # Not a string we care about

        docs = []
        step = 2  # Because the result has content
        offset = 1  # The first item is the count of total matches.

        for i in xrange(1, len(res), step):
            fields_offset = offset

            fields = dict(
                dict(izip(map(to_string, res[i + fields_offset][::2]),
                          map(to_string, res[i + fields_offset][1::2])))
            )

            try:
                del fields['id']
            except KeyError:
                pass

            doc = cls(**fields)
            docs.append(doc)
        return docs


    @classmethod
    def add(cls, models: Sequence['RedisModel']) -> Sequence['RedisModel']:
        return [model.save() for model in models]

    @classmethod
    def update(cls, **field_values):
        """Update this model instance."""
        return cls

    @classmethod
    def values(cls):
        """Return raw values from Redis instead of model instances."""
        return cls

    def delete(self):
        return self.db().delete(self.key())

    def save(self, *args, **kwargs) -> 'RedisModel':
        raise NotImplementedError

    @classmethod
    def redisearch_schema(cls):
        raise NotImplementedError


class HashModel(RedisModel, abc.ABC):
    def __init_subclass__(cls, **kwargs):
        super().__init_subclass__(**kwargs)

        for name, field in cls.__fields__.items():
            if issubclass(field.outer_type_, RedisModel):
                raise RedisModelError(f"HashModels cannot have embedded model "
                                      f"fields. Field: {name}")

            for typ in (Set, Mapping, List):
                if issubclass(field.outer_type_, typ):
                    raise RedisModelError(f"HashModels cannot have set, list,"
                                          f" or mapping fields. Field: {name}")

    def save(self, *args, **kwargs) -> 'HashModel':
        document = jsonable_encoder(self.dict())
        success = self.db().hset(self.key(), mapping=document)

        return success

    @classmethod
    def get(cls, pk: Any) -> 'HashModel':
        document = cls.db().hgetall(cls.make_primary_key(pk))
        if not document:
            raise NotFoundError
        return cls.parse_obj(document)

    @classmethod
    @no_type_check
    def _get_value(cls, *args, **kwargs) -> Any:
        """
        Always send None as an empty string.

        TODO: We do this because redis-py's hset() method requires non-null
        values. Is there a better way?
        """
        val = super()._get_value(*args, **kwargs)
        if val is None:
            return ""
        return val

    @classmethod
    def redisearch_schema(cls):
        hash_prefix = cls.make_key(cls._meta.primary_key_pattern.format(pk=""))
        schema_prefix = f"ON HASH PREFIX 1 {hash_prefix} SCHEMA"
        schema_parts = [schema_prefix] + cls.schema_for_fields()
        return " ".join(schema_parts)

    @classmethod
    def schema_for_fields(cls):
        schema_parts = []

        for name, field in cls.__fields__.items():
            # TODO: Merge this code with schema_for_type()
            _type = field.outer_type_
            if getattr(field.field_info, 'primary_key', None):
                if issubclass(_type, str):
                    redisearch_field = f"{name} TAG"
                else:
                    redisearch_field = cls.schema_for_type(name, _type, field.field_info)
                schema_parts.append(redisearch_field)
            elif getattr(field.field_info, 'index', None) is True:
                schema_parts.append(cls.schema_for_type(name, _type, field.field_info))
                # TODO: Raise error if user embeds a model field or list and makes it
                #  sortable. Instead, the embedded model should mark individual fields
                #  as sortable.
                if getattr(field.field_info, 'sortable', False) is True:
                    schema_parts.append("SORTABLE")
            elif get_origin(_type) == list:
                embedded_cls = get_args(_type)
                if not embedded_cls:
                    # TODO: Test if this can really happen.
                    log.warning("Model %s defined an empty list field: %s", cls, name)
                    continue
                embedded_cls = embedded_cls[0]
                schema_parts.append(cls.schema_for_type(name, embedded_cls,
                                                        field.field_info))
            elif issubclass(_type, RedisModel):
                schema_parts.append(cls.schema_for_type(name, _type, field.field_info))
        return schema_parts

    @classmethod
    def schema_for_type(cls, name, typ: Type, field_info: FieldInfo):
        if get_origin(typ) == list:
            embedded_cls = get_args(typ)
            if not embedded_cls:
                # TODO: Test if this can really happen.
                log.warning("Model %s defined an empty list field: %s", cls, name)
                return ""
            embedded_cls = embedded_cls[0]
            return cls.schema_for_type(name, embedded_cls, field_info)
        elif any(issubclass(typ, t) for t in NUMERIC_TYPES):
            return f"{name} NUMERIC"
        elif issubclass(typ, str):
            if getattr(field_info, 'full_text_search', False) is True:
                return f"{name} TAG {name}_fts TEXT"
            else:
                return f"{name} TAG"
        elif issubclass(typ, RedisModel):
            sub_fields = []
            for embedded_name, field in typ.__fields__.items():
                sub_fields.append(cls.schema_for_type(f"{name}_{embedded_name}", field.outer_type_,
                                                      field.field_info))
            return " ".join(sub_fields)
        else:
            return f"{name} TAG"


class JsonModel(RedisModel, abc.ABC):
    def save(self, *args, **kwargs) -> 'JsonModel':
        success = self.db().execute_command('JSON.SET', self.key(), ".", self.json())
        return success

    @classmethod
    def get(cls, pk: Any) -> 'JsonModel':
        document = cls.db().execute_command("JSON.GET", cls.make_primary_key(pk))
        if not document:
            raise NotFoundError
        return cls.parse_raw(document)

    @classmethod
    def redisearch_schema(cls):
        key_prefix = cls.make_key(cls._meta.primary_key_pattern.format(pk=""))
        schema_prefix = f"ON JSON PREFIX 1 {key_prefix} SCHEMA"
        schema_parts = [schema_prefix] + cls.schema_for_fields()
        return " ".join(schema_parts)

    @classmethod
    def schema_for_fields(cls):
        schema_parts = []
        json_path = "$"

        for name, field in cls.__fields__.items():
            # TODO: Merge this code with schema_for_type()?
            _type = field.outer_type_
            if getattr(field.field_info, 'primary_key', None):
                if issubclass(_type, str):
                    redisearch_field = f"{json_path}.{name} AS {name} TAG"
                else:
                    redisearch_field = cls.schema_for_type(f"{json_path}.{name}", name, "", _type, field.field_info)
                schema_parts.append(redisearch_field)
            elif getattr(field.field_info, 'index', None) is True:
                schema_parts.append(cls.schema_for_type(f"{json_path}.{name}", name, "", _type, field.field_info))
                # TODO: Raise error if user embeds a model field or list and makes it
                #  sortable. Instead, the embedded model should mark individual fields
                #  as sortable.
                if getattr(field.field_info, 'sortable', False) is True:
                    schema_parts.append("SORTABLE")
            elif get_origin(_type) == list:
                embedded_cls = get_args(_type)
                if not embedded_cls:
                    # TODO: Test if this can really happen.
                    log.warning("Model %s defined an empty list field: %s", cls, name)
                    continue
                embedded_cls = embedded_cls[0]
                schema_parts.append(cls.schema_for_type(f"{json_path}.{name}[]", name, f"{name}",
                                                        embedded_cls, field.field_info))
            elif issubclass(_type, RedisModel):
                schema_parts.append(cls.schema_for_type(f"{json_path}.{name}", name, f"{name}", _type,
                                                        field.field_info))
        return schema_parts

    @classmethod
    # TODO: We need both the "name" of the field (address_line_1) as we'll
    #  find it in the JSON document, AND the name of the field as it should
    #  be in the redisearch schema (address_address_line_1). Maybe both "name"
    #  and "name_prefix"?
    def schema_for_type(cls, json_path: str, name: str, name_prefix: str, typ: Type,
                        field_info: FieldInfo) -> str:
        index_field_name = f"{name_prefix}{name}"
        should_index = getattr(field_info, 'index', False)

        if get_origin(typ) == list:
            embedded_cls = get_args(typ)
            if not embedded_cls:
                # TODO: Test if this can really happen.
                log.warning("Model %s defined an empty list field: %s", cls, name)
                return ""
            embedded_cls = embedded_cls[0]
            # TODO: We need to pass the "JSON Path so far" which should include the
            #  correct syntax for an array.
            return cls.schema_for_type(f"{json_path}[]", name, f"{name_prefix}{name}", embedded_cls, field_info)
        elif issubclass(typ, RedisModel):
            sub_fields = []
            for embedded_name, field in typ.__fields__.items():
                sub_fields.append(cls.schema_for_type(f"{json_path}.{embedded_name}",
                                                      embedded_name,
                                                      f"{name_prefix}_",
                                                      field.outer_type_,
                                                      field.field_info))
            return " ".join(filter(None, sub_fields))
        elif should_index:
            if any(issubclass(typ, t) for t in NUMERIC_TYPES):
                return f"{json_path} AS {index_field_name} NUMERIC"
            elif issubclass(typ, str):
                if getattr(field_info, 'full_text_search', False) is True:
                    return f"{json_path} AS {index_field_name} TAG " \
                           f"{json_path} AS {index_field_name}_fts TEXT"
                else:
                    return f"{json_path} AS {index_field_name} TAG"
            else:
                return f"{json_path} AS {index_field_name} TAG"

        return ""