quay/data/model/sqlalchemybridge.py

from sqlalchemy import (Table, MetaData, Column, ForeignKey, Integer, String, Boolean, Text,
                        DateTime, Date, BigInteger, Index)
from peewee import (PrimaryKeyField, CharField, BooleanField, DateTimeField, TextField,
                    ForeignKeyField, BigIntegerField, IntegerField, DateField)


OPTIONS_TO_COPY = [
  'null',
  'default',
  'primary_key',
]


OPTION_TRANSLATIONS = {
  'null': 'nullable',
}


def gen_sqlalchemy_metadata(peewee_model_list):
  metadata = MetaData(naming_convention={
    "ix": 'ix_%(column_0_label)s',
    "uq": "uq_%(table_name)s_%(column_0_name)s",
    "fk": "fk_%(table_name)s_%(column_0_name)s_%(referred_table_name)s",
    "pk": "pk_%(table_name)s"
  })

  for model in peewee_model_list:
    meta = model._meta

    all_indexes = set(meta.indexes)

    columns = []
    for field in meta.get_fields():
      alchemy_type = None
      col_args = []
      col_kwargs = {}
      if isinstance(field, PrimaryKeyField):
        alchemy_type = Integer
      elif isinstance(field, CharField):
        alchemy_type = String(field.max_length)
      elif isinstance(field, BooleanField):
        alchemy_type = Boolean
      elif isinstance(field, DateTimeField):
        alchemy_type = DateTime
      elif isinstance(field, DateField):
        alchemy_type = Date
      elif isinstance(field, TextField):
        alchemy_type = Text
      elif isinstance(field, ForeignKeyField):
        alchemy_type = Integer
        target_name = '%s.%s' % (field.to_field.model_class._meta.db_table,
                                 field.to_field.db_column)
        col_args.append(ForeignKey(target_name))
        all_indexes.add(((field.name, ), field.unique))
      elif isinstance(field, BigIntegerField):
        alchemy_type = BigInteger
      elif isinstance(field, IntegerField):
        alchemy_type = Integer
      else:
        raise RuntimeError('Unknown column type: %s' % field)

      for option_name in OPTIONS_TO_COPY:
        alchemy_option_name = (OPTION_TRANSLATIONS[option_name]
                               if option_name in OPTION_TRANSLATIONS else option_name)
        if alchemy_option_name not in col_kwargs:
          option_val = getattr(field, option_name)
          col_kwargs[alchemy_option_name] = option_val

      if field.unique or field.index:
        all_indexes.add(((field.name, ), field.unique))

      new_col = Column(field.db_column, alchemy_type, *col_args, **col_kwargs)
      columns.append(new_col)

    new_table = Table(meta.db_table, metadata, *columns)

    for col_prop_names, unique in all_indexes:
      col_names = [meta.fields[prop_name].db_column for prop_name in col_prop_names]
      index_name = '%s_%s' % (meta.db_table, '_'.join(col_names))
      col_refs = [getattr(new_table.c, col_name) for col_name in col_names]
      Index(index_name, *col_refs, unique=unique)

  return metadata