forked from django/django
-
Notifications
You must be signed in to change notification settings - Fork 1
/
compiler.py
132 lines (119 loc) · 4.92 KB
/
compiler.py
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
# TODO: ...
class SQLCompiler(object):
def __init__(self, query, connection, using):
self.query = query
self.connection = connection
self.using = using
def get_filters(self, where, correct=False):
assert where.connector == "AND"
filters = {}
for child in where.children:
if isinstance(child, self.query.where_class):
child_filters = self.get_filters(child)
for k, v in child_filters.iteritems():
if k in filters:
v = {"$and": [filters[k], v]}
if where.negated:
v = {"$not": v}
filters[k] = v
else:
field, val = self.make_atom(*child, **{"negated": where.negated})
filters[field] = val
if correct:
self.correct_filters(filters)
return filters
def make_atom(self, lhs, lookup_type, value_annotation, params_or_value, negated):
assert lookup_type in ["exact", "isnull", "lt"], lookup_type
if hasattr(lhs, "process"):
lhs, params = lhs.process(lookup_type, params_or_value, self.connection)
else:
params = Field().get_db_prep_lookup(lookup_type, params_or_value,
connection=self.connection, prepared=True)
assert isinstance(lhs, (list, tuple))
table, column, _ = lhs
assert table == self.query.model._meta.db_table
if column == self.query.model._meta.pk.column:
column = "_id"
if lookup_type == "exact":
val = params[0]
if negated:
val = {"$ne": val}
return column, val
elif lookup_type == "isnull":
val = None
if value_annotation == negated:
val = {"$not": val}
return column, val
elif lookup_type == "lt":
return column, {"$lt": params[0]}
def correct_filters(self, filters):
for k, v in filters.items():
if isinstance(v, dict) and v.keys() == ["$not"]:
if isinstance(v["$not"], dict) and v["$not"].keys() == ["$and"]:
del filters[k]
or_vals = [self.negate(k, v) for v in v["$not"]["$and"]]
assert "$or" not in filters
filters["$or"] = or_vals
def negate(self, k, v):
if isinstance(v, dict):
if v.keys() == ["$not"]:
return {k: v["$not"]}
return {k: {"$not": v}}
return {k: {"$ne": v}}
def build_query(self, aggregates=False):
assert len([a for a in self.query.alias_map if self.query.alias_refcount[a]]) <= 1
if not aggregates:
assert self.query.default_cols
assert not self.query.distinct
assert not self.query.extra
assert not self.query.having
assert self.query.high_mark is None
assert not self.query.order_by
filters = self.get_filters(self.query.where, correct=True)
return self.connection.db[self.query.model._meta.db_table].find(filters)
def results_iter(self):
query = self.build_query()
for row in query:
yield tuple(
row[f.column if f is not self.query.model._meta.pk else "_id"]
for f in self.query.model._meta.fields
)
def has_results(self):
try:
self.build_query()[0]
except IndexError:
return False
else:
return True
def get_aggregates(self):
assert len(self.query.aggregates) == 1
agg = self.query.aggregates.values()[0]
assert (
isinstance(agg, self.query.aggregates_module.Count) and (
agg.col == "*" or
isinstance(agg.col, tuple) and agg.col == (self.query.model._meta.db_table, self.query.model._meta.pk.column)
)
)
return [self.build_query(aggregates=True).count()]
class SQLInsertCompiler(SQLCompiler):
def insert(self, return_id=False):
values = dict([
(c, v)
for c, v in zip(self.query.columns, self.query.params)
])
if self.query.model._meta.pk.column in values:
values["_id"] = values.pop(self.query.model._meta.pk.column)
if "_id" in values and not values["_id"]:
del values["_id"]
return self.connection.db[self.query.model._meta.db_table].insert(values)
class SQLUpdateCompiler(SQLCompiler):
def update(self, result_type):
# TODO: more asserts
filters = self.get_filters(self.query.where)
# TODO: Don't use set for everything, use INC and such where
# appropriate.
return self.connection.db[self.query.model._meta.db_table].update(
filters,
{"$set": dict((f.column, val) for f, o, val in self.query.values)},
multi=True
)