← Back to branch summary

~unity-2d-team/unity-2d/Shell-MultiMonitor

« back to all changes in this revision

Viewing changes to grackle/model.py

  • Committer: William Grant
  • Date: 2012-04-20 02:20:58 UTC
  • mfrom: (6.1.70 trunk)
  • Revision ID: william.grant@canonical.com-20120420022058-3nkracsmlg7akydu
Merge trunk.

Show diffs side-by-side

added added

removed removed

Lines of Context:
grackle/model.py
14
14
# License along with this program. If not, see
15
15
# <http://www.gnu.org/licenses/>.
16
16
 
 
17
import calendar
17
18
import datetime
18
 
import dateutil.tz
19
19
import email.parser
20
20
from email.utils import parsedate_tz
 
21
import functools
21
22
import logging
22
 
import time
23
23
import uuid
24
24
 
25
 
import pycassa
 
25
import dateutil.tz
 
26
import pycassa.pool
26
27
from pycassa.system_manager import (
27
28
    LEXICAL_UUID_TYPE,
28
29
    SystemManager,
85
86
 
86
87
 
87
88
def _utc_timestamp(dt):
88
 
    return time.mktime(_utc_datetime(dt).timetuple()) - time.timezone
 
89
    return calendar.timegm(_utc_datetime(dt).timetuple())
89
90
 
90
91
 
91
92
def _utc_timeuuid(dt, lowest_val=True):
121
122
    return memo, start, finish
122
123
 
123
124
 
124
 
def _format_message(message):
125
 
    return {
126
 
        'date': message.get('date'),
127
 
        'from': message.get('from'),
128
 
        'subject': message.get('subject'),
129
 
        'message-id': message.get('message-id'),
130
 
        }
 
125
LEGAL_HEADERS = set([
 
126
    'date', 'from', 'subject', 'message-id',
 
127
    ])
 
128
 
 
129
 
 
130
def _format_message(message, headers=[], include_raw=False):
 
131
    data = {}
 
132
 
 
133
    if headers:
 
134
        assert not set(headers).difference(LEGAL_HEADERS)
 
135
        hdict = {}
 
136
        for header in headers:
 
137
            hdict[header] = message.get(header)
 
138
        data['headers'] = hdict
 
139
 
 
140
    if include_raw:
 
141
        data['raw'] = message['raw']
 
142
 
 
143
    return data
131
144
 
132
145
 
133
146
class CassandraConnection(object):
134
147
 
135
 
    def __init__(self, keyspace, host):
 
148
    def __init__(self, keyspace, hosts):
136
149
        self._keyspace = keyspace
137
 
        self._host = host
 
150
        self._hosts = hosts
138
151
        self._connection = self._connect()
 
152
        self._pool = self._connect()
139
153
        self.messages = self._column_family('message')
140
154
        self.archive_messages = self._column_family('archive_message')
141
155
 
142
156
    def _connect(self):
143
 
        return pycassa.connect(self._keyspace, self._host)
 
157
        return pycassa.pool.ConnectionPool(self._keyspace, self._hosts)
144
158
 
145
159
    def _column_family(self, name):
146
 
        return pycassa.ColumnFamily(self._connection, name)
 
160
        return pycassa.ColumnFamily(self._pool, name)
147
161
 
148
162
    def add_message(self, archive_uuid, message):
149
163
        message_uuid = uuid.uuid4()
150
164
        message_date, message_dict = _parse_message(message)
151
 
        message_dict['content'] = message
 
165
        message_dict['raw'] = message
152
166
        message_dict['date_created'] = (
153
167
            datetime.datetime.utcnow().isoformat() + 'Z')
154
168
        self.messages.insert(message_uuid, message_dict)
172
186
            return sequence[1:]
173
187
 
174
188
    def get_messages(self, archive_uuid, order, count, memo, backward=False,
175
 
                     start_date=None, finish_date=None):
 
189
                     start_date=None, finish_date=None, format='all',
 
190
                     headers=['from', 'date', 'subject', 'message-id']):
176
191
        if order in ("date", "-date"):
177
192
            reversed = order[0] == '-'
178
193
        else:
204
219
 
205
220
        # We've narrowed down the message references. Fetch the messages.
206
221
        ids = [v for k, v in pairs]
207
 
        messages = self.messages.multiget(
208
 
            ids, columns=['date', 'from', 'subject', 'message-id'])
 
222
        formatter = functools.partial(
 
223
            _format_message, headers=headers, include_raw=True)
 
224
        # XXX: No need to get all columns. Restrict based on format.
 
225
        messages = self.messages.multiget(ids)
209
226
 
210
227
        return (
211
228
            str(pairs[0][0]),
212
 
            [_format_message(messages[id]) for id in ids],
 
229
            [formatter(messages[id]) for id in ids],
213
230
            str(pairs[-1][0]),
214
231
            )