quay/data/fields.py

import base64
import resumablehashlib
import json

from peewee import TextField, CharField, Clause
from data.text import prefix_search


class _ResumableSHAField(TextField):
  def _create_sha(self):
    raise NotImplementedError

  def db_value(self, value):
    if value is None:
      return None

    sha_state = value.state()

    # One of the fields is a byte string, let's base64 encode it to make sure
    # we can store and fetch it regardless of default collocation.
    sha_state[3] = base64.b64encode(sha_state[3])

    return json.dumps(sha_state)

  def python_value(self, value):
    if value is None:
      return None

    sha_state = json.loads(value)

    # We need to base64 decode the data bytestring.
    sha_state[3] = base64.b64decode(sha_state[3])
    to_resume = self._create_sha()
    to_resume.set_state(sha_state)
    return to_resume


class ResumableSHA256Field(_ResumableSHAField):
  def _create_sha(self):
    return resumablehashlib.sha256()


class ResumableSHA1Field(_ResumableSHAField):
  def _create_sha(self):
    return resumablehashlib.sha1()


class JSONField(TextField):
  def db_value(self, value):
    return json.dumps(value)

  def python_value(self, value):
    if value is None or value == "":
      return {}
    return json.loads(value)


class Base64BinaryField(TextField):
  def db_value(self, value):
    if value is None:
      return None
    return base64.b64encode(value)

  def python_value(self, value):
    if value is None:
      return None
    return base64.b64decode(value)


def _add_fulltext(field_class):
  """ Adds support for full text indexing and lookup to the given field class. """
  class indexed_class(field_class):
    # Marker used by SQLAlchemy translation layer to add the proper index for full text searching.
    __fulltext__ = True

    def __init__(self, match_function, *args, **kwargs):
      field_class.__init__(self, *args, **kwargs)
      self.match_function = match_function

    def match(self, query):
      return self.match_function(self, query)

    def match_prefix(self, query):
      return prefix_search(self, query)

    def __mod__(self, _):
      raise Exception('Unsafe operation: Use `match` or `match_prefix`')

    def __pow__(self, _):
      raise Exception('Unsafe operation: Use `match` or `match_prefix`')

    def __contains__(self, _):
      raise Exception('Unsafe operation: Use `match` or `match_prefix`')

    def contains(self, _):
      raise Exception('Unsafe operation: Use `match` or `match_prefix`')

    def startswith(self, _):
      raise Exception('Unsafe operation: Use `match` or `match_prefix`')

    def endswith(self, _):
      raise Exception('Unsafe operation: Use `match` or `match_prefix`')

  return indexed_class


FullIndexedCharField = _add_fulltext(CharField)
FullIndexedTextField = _add_fulltext(TextField)