Beispiel #1
0
    def run(self):
        from pystil.websocket import broadcast
        from pystil.utils import visit_to_table_line
        while True:
            try:
                self.log.debug('Blocking for messages')
                message = MESSAGE_QUEUE.get(True)
                self.log.debug('Message got %r' % message)
                try:
                    visit_or_uuid, opening = message.process(self.db)

                    if opening:
                        visit = Visit(**visit_or_uuid)
                        broadcast(
                            'VISIT|' + visit.host + '|' +
                            visit_to_table_line(visit))
                    else:
                        visit_or_uuid and broadcast('EXIT|%s' % visit_or_uuid)
                except:
                    self.log.exception('Error processing visit')
            except:
                self.log.exception('Exception in loop')
Beispiel #2
0
    def io_callback(self, fd=None, events=None):
        try:
            state = self.connection.poll()
        except psycopg2.extensions.QueryCanceledError:
            log.info('Canceling request %r' % self, exc_info=True)
            self.cursor.execute('ROLLBACK -- CANCEL')
            self.state = 'terminated'
        except (psycopg2.Warning, psycopg2.Error):
            log.exception('Poll error')
            self.momoko_connection.ioloop.remove_handler(
                self.momoko_connection.fileno)
            raise
        else:
            if state == POLL_OK:
                if self.state == 'terminated':
                    self.momoko_connection.ioloop.remove_handler(
                        self.momoko_connection.fileno)
                    return

                rows = self.cursor.fetchmany()

                if not rows:
                    self.state = 'terminated'
                    self.cursor.execute('CLOSE visit_cur; ROLLBACK; -- NOROWS')
                    try:
                        self.write_message('END|Done found %d visit%s' % (
                            self.count, 's' if self.count > 1 else ''))
                    except:
                        pass
                else:
                    try:
                        for row in rows:
                            if row.id:
                                self.count += 1
                                self.write_message(
                                    'VISIT|' + visit_to_table_line(row))
                    except Exception as e:
                        log.warn('During write', exc_info=True)
                        self.state = 'terminated'
                        self.cursor.execute(
                            'CLOSE visit_cur; ROLLBACK; -- WSERROR')
                        try:
                            self.write_message(
                                'END|%s: %s' % (type(e), str(e)))
                        except:
                            pass
                    else:
                        if self.count < self.stop:
                            self.cursor.execute(
                                'FETCH FORWARD 1 FROM visit_cur;')
                        else:
                            self.state = 'paused'
                            try:
                                self.write_message(
                                    'PAUSE|Paused on %d visits' % self.count)
                            except:
                                pass
            elif state == POLL_READ:
                self.momoko_connection.ioloop.update_handler(
                    self.momoko_connection.fileno, IOLoop.READ)
            elif state == POLL_WRITE:
                self.momoko_connection.ioloop.update_handler(
                    self.momoko_connection.fileno, IOLoop.WRITE)
            else:
                raise psycopg2.OperationalError(
                    'poll() returned {0}'.format(state))