Source

djsqlalchemy / djsqlalchemy / mapper.py

Full commit
import logging

from django.db.models import ForeignKey

from djsqlalchemy.alchemy import get_engine


__all__ = ['fetchall', 'get_list', 'get_one']


def map_related(obj, values, prefix=''):
    # special case for master obj
    local_prefix = prefix or obj._meta.db_table

    fks = []
    for field in obj._meta.local_fields:
        if isinstance(field, ForeignKey):
            fks.append(field)

            label = local_prefix + '_' + field.name + '_' + field.rel.field_name
            attr = field.name + '_' + field.rel.field_name

        else:
            label = local_prefix + '_' + field.name
            attr = field.name

        if label in values:
            setattr(obj, attr, values[label])
            del values[label]

    if obj.pk:
        for fk in fks:
            rel_obj = fk.rel.to()
            map_related(rel_obj, values, prefix=prefix + fk.name)
            if rel_obj.pk:
                cache_name = '_%s_cache' % fk.name
                setattr(obj, cache_name, rel_obj)


def map_values(model, values):
    obj = model()

    # TODO: replace with values = dict(values)
    d = {}
    for i, key in enumerate(values.keys()):
        d[key] = values._row[i]
    values = d

    map_related(obj, values)

    # set other values on master obj as is
    for label, value in values.items():
        setattr(obj, label, value)
    return obj


def fetchall(s, model=None):
    engine = get_engine()

    result = engine.execute(s)
    values_list = result.fetchall()

    for values in values_list:
        if model:
            yield map_values(model, values)
        else:
            yield values


def get_list(s, model=None):
    return list(fetchall(s, model))


def get_one(s, model=None):
    objs = list(fetchall(s, model))

    if not objs:
        if model:
            raise model.DoesNotExist()
        else:
            raise ValueError('Query returned empty result')
    elif len(objs) > 1:
        if model:
            raise model.MultipleObjectsReturned()
        else:
            raise ValueError('Query returned multiple rows')

    return objs[0]