8
by William Grant
A little bit of server. |
1 |
# Copyright (c) 2012 Canonical Ltd
|
2 |
#
|
|
3 |
# This program is free software: you can redistribute it and/or modify
|
|
4 |
# it under the terms of the GNU Affero General Public License as published by
|
|
5 |
# the Free Software Foundation, either version 3 of the License, or
|
|
6 |
# (at your option) any later version.
|
|
7 |
#
|
|
8 |
# This program is distributed in the hope that it will be useful,
|
|
9 |
# but WITHOUT ANY WARRANTY; without even the implied warranty of
|
|
10 |
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
|
|
11 |
# GNU Affero General Public License for more details.
|
|
12 |
#
|
|
13 |
# You should have received a copy of the GNU Affero General Public
|
|
14 |
# License along with this program. If not, see
|
|
15 |
# <http://www.gnu.org/licenses/>.
|
|
16 |
||
17 |
import datetime |
|
18 |
import uuid |
|
19 |
||
20 |
import pycassa |
|
21 |
from pycassa.system_manager import ( |
|
22 |
LEXICAL_UUID_TYPE, |
|
23 |
SystemManager, |
|
24 |
TIME_UUID_TYPE, |
|
25 |
)
|
|
26 |
||
27 |
from grackle.server.cassandra import workaround_1779 |
|
28 |
||
29 |
||
30 |
def create_schema(host, keyspace, clobber=False): |
|
31 |
mgr = SystemManager(host) |
|
32 |
||
33 |
if clobber: |
|
34 |
for cf in mgr.get_keyspace_column_families(keyspace): |
|
35 |
mgr.drop_column_family(keyspace, cf) |
|
36 |
||
37 |
try: |
|
38 |
workaround_1779( |
|
39 |
mgr.create_column_family, keyspace, 'message', |
|
40 |
key_validation_class=LEXICAL_UUID_TYPE) |
|
41 |
workaround_1779( |
|
42 |
mgr.create_column_family, keyspace, 'archive_message', |
|
43 |
comparator_type=TIME_UUID_TYPE, |
|
44 |
default_validation_class=LEXICAL_UUID_TYPE) |
|
45 |
pass
|
|
46 |
finally: |
|
47 |
mgr.close() |
|
48 |
||
49 |
||
50 |
class CassandraConnection(object): |
|
51 |
||
52 |
def __init__(self, keyspace, host): |
|
53 |
self._keyspace = keyspace |
|
54 |
self._host = host |
|
55 |
self._connection = self._connect() |
|
56 |
self.messages = self._column_family('message') |
|
57 |
self.archive_messages = self._column_family('archive_message') |
|
58 |
||
59 |
def _connect(self): |
|
60 |
return pycassa.connect(self._keyspace, self._host) |
|
61 |
||
62 |
def _column_family(self, name): |
|
63 |
return pycassa.ColumnFamily(self._connection, name) |
|
64 |
||
65 |
def add_message(self, archive_uuid, message): |
|
66 |
message_uuid = uuid.uuid4() |
|
67 |
self.messages.insert( |
|
68 |
message_uuid, |
|
69 |
{'date_created': datetime.datetime.utcnow().isoformat() + 'Z', |
|
70 |
'content': message, |
|
71 |
})
|
|
72 |
self.archive_messages.insert( |
|
73 |
archive_uuid, {datetime.datetime.now(): message_uuid}) |
|
74 |
return message_uuid |
|
75 |
||
76 |
def get_messages(self, archive_uuid): |
|
77 |
ids = self.archive_messages.get( |
|
78 |
archive_uuid, column_count=10, column_reversed=True).values() |
|
79 |
messages = self.messages.multiget(ids) |
|
80 |
return [messages[id]['content'] for id in ids] |