def open(self):
log.info('Opening last visits websocket')
site = self.get_secure_cookie('_pystil_site')
site = site and site.decode('utf-8').split('|')[0]
if site is not None:
self.site = site
LastVisitsWebSocket.waiters.add(self)
else:
log.warn('Last visits websocket open without secure cookie')
self.close()
def open(self):
log.info('Opening query websocket')
site = self.get_secure_cookie('_pystil_site')
site = site and site.decode('utf-8').split('|')[0]
self.query = None
self.state = None
if site is None:
log.warn('Query websocket open without secure cookie')
self.close()
return
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))