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
133
134
135
|
# Copyright (c) 2012 Canonical Ltd
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE. See the
# GNU Affero General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public
# License along with this program. If not, see
# <http://www.gnu.org/licenses/>.
import datetime
import dateutil.tz
import email.parser
from email.utils import parsedate_tz
import logging
import uuid
import pycassa
from pycassa.system_manager import (
LEXICAL_UUID_TYPE,
SystemManager,
TIME_UUID_TYPE,
)
from grackle.server.cassandra import workaround_1779
def create_schema(host, keyspace, clobber=False):
mgr = SystemManager(host)
if clobber:
for cf in mgr.get_keyspace_column_families(keyspace):
mgr.drop_column_family(keyspace, cf)
try:
workaround_1779(
mgr.create_column_family, keyspace, 'message',
key_validation_class=LEXICAL_UUID_TYPE)
workaround_1779(
mgr.create_column_family, keyspace, 'archive_message',
comparator_type=TIME_UUID_TYPE,
default_validation_class=LEXICAL_UUID_TYPE)
pass
finally:
mgr.close()
class CassandraConnection(object):
def __init__(self, keyspace, host):
self._keyspace = keyspace
self._host = host
self._connection = self._connect()
self.messages = self._column_family('message')
self.archive_messages = self._column_family('archive_message')
def _connect(self):
return pycassa.connect(self._keyspace, self._host)
def _column_family(self, name):
return pycassa.ColumnFamily(self._connection, name)
def _parse_message(self, message):
"""Get a date and dict of an RFC822 message."""
parsed = email.parser.Parser().parsestr(message)
message_dict = {}
message_dict['content'] = message
for key in ('from', 'to', 'subject', 'message-id'):
value = parsed.get(key, None)
if value is not None:
message_dict[key] = value
date = parsed.get('date')
if date is not None:
try:
pdate = parsedate_tz(date)
date = datetime.datetime(
*pdate[:6],
tzinfo=dateutil.tz.tzoffset('', pdate[9]))
except ValueError:
pass
message_dict['date'] = date.isoformat() if date is not None else None
return date, message_dict
def add_message(self, archive_uuid, message):
message_uuid = uuid.uuid4()
message_date, message_dict = self._parse_message(message)
message_dict['date_created'] = (
datetime.datetime.utcnow().isoformat() + 'Z')
self.messages.insert(message_uuid, message_dict)
self.archive_messages.insert(
archive_uuid,
{message_date.astimezone(dateutil.tz.tzutc()): message_uuid})
logging.debug(
'Imported %s into %s'
% (message_dict.get('message-id', None), archive_uuid))
return message_uuid
def _format_message(self, message):
return {
'date': message.get('date'),
'from': message.get('from'),
'subject': message.get('subject'),
}
def get_messages(self, archive_uuid, order, count, start):
if order in ("date", "-date"):
reversed = order[0] == '-'
else:
raise AssertionError("Unsupported order.")
pairs = self.archive_messages.get(
archive_uuid, column_count=count + 1,
column_start=start, column_reversed=reversed).items()
ids = [v for k, v in pairs]
messages = self.messages.multiget(
ids, columns=['date', 'from', 'subject'])
actual_count = len(pairs)
if len(pairs) > count:
assert len(pairs) == count + 1
actual_count -= 1
next_memo = str(pairs[count][0])
else:
next_memo = None
return (
[self._format_message(messages[id]) for id in ids[:actual_count]],
next_memo,
)
|