Allow to save dictionaries, lists, tuples, etc, as document ids
[mojodb.git] / cursor.py
1 # -*- coding: utf-8 -*-
2 ##############################################################################
3 #
4 #    mojo, a Python library for implementing document based databases
5 #    Copyright (C) 2013-2014 by Javier Sancho Fernandez <jsf at jsancho dot org>
6 #
7 #    This program is free software: you can redistribute it and/or modify
8 #    it under the terms of the GNU General Public License as published by
9 #    the Free Software Foundation, either version 3 of the License, or
10 #    (at your option) any later version.
11 #
12 #    This program is distributed in the hope that it will be useful,
13 #    but WITHOUT ANY WARRANTY; without even the implied warranty of
14 #    MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
15 #    GNU General Public License for more details.
16 #
17 #    You should have received a copy of the GNU General Public License
18 #    along with this program.  If not, see <http://www.gnu.org/licenses/>.
19 #
20 ##############################################################################
21
22 import cPickle
23
24 class Cursor(object):
25     def __init__(self, collection, spec=None, fields=None, **kwargs):
26         if spec and not type(spec) is dict:
27             raise Exception("spec must be an instance of dict")
28
29         self.collection = collection
30         self.spec = spec
31         if self.collection.exists():
32             self.fields = self._get_fields(fields)
33             self.cursor = self._get_cursor()
34         else:
35             self.fields = None
36             self.cursor = None
37
38     def __iter__(self):
39         return self
40
41     def _get_fields(self, fields):
42         set_all_fields = set(self.collection._get_fields())
43         if fields is None:
44             res_fields = list(set_all_fields)
45         elif type(fields) is dict:
46             fields_without_id = filter(lambda x: x[0] != '_id', fields.iteritems())
47             if fields_without_id[0][1]:
48                 first = True
49                 res_fields = set()
50             else:
51                 first = False
52                 res_fields = set(set_all_fields)
53             for f in fields_without_id:
54                 if f[1] and f[0] in set_all_fields:
55                     if first:
56                         res_fields.add(f[0])
57                     else:
58                         raise Exception("You cannot currently mix including and excluding fields. Contact us if this is an issue.")
59                 elif not f[1]:
60                     if not first:
61                         res_fields.discard(f[0])
62                     else:
63                         raise Exception("You cannot currently mix including and excluding fields. Contact us if this is an issue.")
64             if '_id' in fields and not fields['_id']:
65                 res_fields.discard('_id')
66             else:
67                 res_fields.add('_id')
68             res_fields = list(res_fields)
69         else:
70             set_fields = set(list(fields))
71             set_fields.add('_id')
72             res_fields = list(set_all_fields.intersection(set_fields))
73
74         return res_fields
75
76     def _get_cursor(self):
77         query = {}
78         table_id = '%s$_id' % self.collection.table_name
79
80         query['select'] = [(table_id, 'id')]
81         for f in filter(lambda x: x != '_id', self.fields):
82             table_f = '%s$%s' % (self.collection.table_name, f)
83             q = self._get_cursor_field(table_id, table_f)
84             query['select'].append(q)
85
86         query['from'] = [table_id]
87
88         if self.spec:
89             query['where'] = []
90             for k, v in self.spec.iteritems():
91                 table_f = '%s$%s' % (self.collection.table_name, k)
92                 if k == u'_id':
93                     query['where'].append(((table_id, 'id'), '=', cPickle.dumps(v)))
94                 elif type(v) in (int, float):
95                     field_q = self._get_cursor_field(table_id, table_f, field_name='number')
96                     query['where'].append((field_q, '=', v))
97                 else:
98                     field_q = self._get_cursor_field(table_id, table_f)
99                     query['where'].append((field_q, '=', cPickle.dumps(v)))
100
101         return self.collection.database.connection._get_cursor(self.collection.database.db_name, query)
102
103     def _get_cursor_field(self, table_id, table_field, field_name='value'):
104         return {
105             'select': [(table_field, field_name)],
106             'from': [table_field],
107             'where': [((table_field, 'id'), '=', (table_id, 'id'))],
108             }
109
110     def next(self):
111         if self.cursor is None:
112             raise StopIteration
113
114         if self.cursor:
115             res = self.collection.database.connection._next(self.cursor)
116             if res is None:
117                 raise StopIteration
118             else:
119                 document = {}
120                 if u'_id' in self.fields:
121                     document[u'_id'] = cPickle.loads(res[0])
122                 fields_without_id = filter(lambda x: x != u'_id', self.fields)
123                 for i in xrange(len(fields_without_id)):
124                     if not res[i + 1] is None:
125                         document[fields_without_id[i]] = cPickle.loads(res[i + 1])
126                 return document
127         else:
128             return None