664961a040680aebadf96e4e3724e1e765713b9c
[mojodb.git] / collection.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 from cursor import Cursor
24 import uuid
25
26 class Collection(object):
27     def __init__(self, database, table_name):
28         self.database = database
29         self.table_name = unicode(table_name)
30
31     def __repr__(self):
32         return "Collection(%r, %r)" % (self.database, self.table_name)
33
34     def exists(self):
35         return (self.database.exists() and self.table_name in self.database.collection_names())
36
37     def _create_table(self):
38         fields = [
39             {'name': 'id', 'type': 'char', 'size': 32, 'primary': True},
40             ]
41         return self.database.connection._create_table(self.database.db_name, '%s$_id' % self.table_name, fields)
42
43     def _create_field(self, field_name):
44         fields = [
45             {'name': 'id', 'type': 'char', 'size': 32, 'primary': True},
46             {'name': 'value', 'type': 'text', 'null': False},
47             {'name': 'number', 'type': 'float'},
48             ]
49         return self.database.connection._create_table(self.database.db_name, '%s$%s' % (self.table_name, field_name), fields)
50
51     def _get_fields(self):
52         tables = self.database.connection._get_tables(self.database.db_name)
53         return [unicode(x[x.find('$')+1:]) for x in filter(lambda x: x.startswith('%s$' % self.table_name), tables)]
54
55     def count(self):
56         return self.database.connection._count(self.database.db_name, self.table_name)
57
58     def find(self, *args, **kwargs):
59         return Cursor(self, *args, **kwargs)
60
61     def insert(self, doc_or_docs):
62         if not self.database.db_name in self.database.connection.database_names():
63             self.database._create_database()
64         if not self.table_name in self.database.collection_names():
65             self._create_table()
66
67         if not type(doc_or_docs) in (list, tuple):
68             docs = [doc_or_docs]
69         else:
70             docs = doc_or_docs
71         for doc in docs:
72             if not '_id' in doc:
73                 doc['_id'] = uuid.uuid4().hex
74             self._insert_document(doc)
75
76         if type(doc_or_docs) in (list, tuple):
77             return [d['_id'] for d in docs]
78         else:
79             return docs[0]['_id']
80
81     def _insert_document(self, doc):
82         table_id = '%s$_id' % self.table_name
83         fields = self._get_fields()
84         self.database.connection._insert(self.database.db_name, table_id, {'id': doc['_id']})
85         for f in doc:
86             if f == '_id':
87                 continue
88             if not f in fields:
89                 self._create_field(f)
90             table_f = '%s$%s' % (self.table_name, f)
91             values = {
92                 'id': doc['_id'],
93                 'value': cPickle.dumps(doc[f]),
94                 }
95             if type(doc[f]) in (int, float):
96                 values['number'] = doc[f]
97             self.database.connection._insert(self.database.db_name, table_f, values)