quay/data/fields.py

import base64
import resumablehashlib
import json

from peewee import TextField


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

  def db_value(self, value):
    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):
    to_resume = self._create_sha()
    if value is None:
      return to_resume

    sha_state = json.loads(value)

    # We need to base64 decode the data bytestring.
    sha_state[3] = base64.b64decode(sha_state[3])
    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)