1
# Copyright (c) 2012 Canonical Ltd
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.
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.
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/>.
20
from email.utils import parsedate_tz
25
from pycassa.system_manager import (
31
from grackle.server.cassandra import workaround_1779
34
def create_schema(host, keyspace, clobber=False):
35
mgr = SystemManager(host)
38
for cf in mgr.get_keyspace_column_families(keyspace):
39
mgr.drop_column_family(keyspace, cf)
43
mgr.create_column_family, keyspace, 'message',
44
key_validation_class=LEXICAL_UUID_TYPE)
46
mgr.create_column_family, keyspace, 'archive_message',
47
comparator_type=TIME_UUID_TYPE,
48
default_validation_class=LEXICAL_UUID_TYPE)
54
class CassandraConnection(object):
56
def __init__(self, keyspace, host):
57
self._keyspace = keyspace
59
self._connection = self._connect()
60
self.messages = self._column_family('message')
61
self.archive_messages = self._column_family('archive_message')
64
return pycassa.connect(self._keyspace, self._host)
66
def _column_family(self, name):
67
return pycassa.ColumnFamily(self._connection, name)
69
def add_message(self, archive_uuid, message):
70
message_uuid = uuid.uuid4()
71
parsed = email.parser.Parser().parsestr(message)
72
date = parsed.get('date')
75
pdate = parsedate_tz(date)
76
date = datetime.datetime(
78
tzinfo=dateutil.tz.tzoffset('', pdate[9]))
83
{'date_created': datetime.datetime.utcnow().isoformat() + 'Z',
85
'from': parsed.get('From'),
86
'to': parsed.get('To'),
87
'subject': parsed.get('Subject'),
88
'date': date.isoformat() if date is not None else None,
89
'message_id': parsed.get('Message-ID'),
91
self.archive_messages.insert(
93
{date.astimezone(dateutil.tz.tzutc()): message_uuid})
95
'Imported %s into %s' % (parsed.get('Message-ID'), archive_uuid))
98
def _format_message(self, message):
100
'date': message['date'],
101
'from': message['from'],
102
'subject': message['subject'],
105
def get_messages(self, archive_uuid, order, count, start):
106
if order in ("date", "-date"):
107
reversed = order[0] == '-'
109
raise AssertionError("Unsupported order.")
110
pairs = self.archive_messages.get(
111
archive_uuid, column_count=count + 1,
112
column_start=start, column_reversed=reversed).items()
113
ids = [v for k, v in pairs]
114
messages = self.messages.multiget(
115
ids, columns=['date', 'from', 'subject'])
116
actual_count = len(pairs)
117
if len(pairs) > count:
118
assert len(pairs) == count + 1
120
next_memo = str(pairs[count][0])
124
[self._format_message(messages[id]) for id in ids[:actual_count]],