# -*- coding: utf-8 -*-
# vi:si:et:sw=4:sts=4:ts=4
##
## Copyright (C) 2012-2013 Async Open Source <http://www.async.com.br>
## All rights reserved
##
## This program is free software; you can redistribute it and/or modify
## it under the terms of the GNU General Public License as published by
## the Free Software Foundation; either version 2 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 General Public License for more details.
##
## You should have received a copy of the GNU General Public License
## along with this program; if not, write to the Free Software
## Foundation, Inc., or visit: http://www.gnu.org/.
##
## Author(s): Stoq Team <stoq-devel@async.com.br>
##
import datetime
import os
import sys
import platform
import struct
import psycopg2
from storm.tracer import BaseStatementTracer, install_tracer
try:
from sqlparse import engine, filters, sql
has_sqlparse = True
except ImportError:
has_sqlparse = False
# http://stackoverflow.com/questions/566746/how-to-get-console-window-width-in-python
[docs]def getTerminalSize():
if platform.system() != 'Linux':
return 80, 20
import fcntl
import termios
env = os.environ
def ioctl_GWINSZ(fd):
try:
cr = struct.unpack('hh', fcntl.ioctl(fd, termios.TIOCGWINSZ, '1234'))
except:
return 80, 20
return cr
cr = ioctl_GWINSZ(0) or ioctl_GWINSZ(1) or ioctl_GWINSZ(2)
if not cr:
try:
fd = os.open(os.ctermid(), os.O_RDONLY)
cr = ioctl_GWINSZ(fd)
os.close(fd)
except:
pass
if not cr:
try:
cr = (env['LINES'], env['COLUMNS'])
except:
cr = (25, 80)
return int(cr[1]), int(cr[0])
if has_sqlparse:
[docs] class MyReindentFilter(filters.ReindentFilter):
def __init__(self, max_width):
self.max_width = max_width
filters.ReindentFilter.__init__(self)
def _process_identifierlist(self, tlist):
identifiers = list(tlist.get_identifiers())
if len(identifiers) > 1 and not tlist.within(sql.Function):
# This is not working in some cases
# first = list(identifiers[0].flatten())[0]
# num_offset = self._get_offset(first) - len(first.value)
num_offset = 7
self.offset += num_offset
width = self.offset
for token in identifiers:
width += len(str(token)) + 2
if width > self.max_width:
tlist.insert_before(token, self.nl())
width = self.offset + len(str(token))
self.offset -= num_offset
return True
[docs]class StoqlibDebugTracer(BaseStatementTracer):
ATTRIBUTES = dict(bold=1, dark=2, underline=4, blink=5,
reverse=7, concealed=8)
COLORS = dict(grey=30, red=31, green=32, yellow=33, blue=34,
magenta=35, cyan=36, white=37)
RESET = '\033[0m' # pylint: disable=W1401
def __init__(self, stream=None):
# This colors will be used to highlight the transaction
self._available_colors = ['blue', 'green', 'magenta', 'yellow', 'cyan',
'red']
self._current_color = 0
# Mapping pid > color
self._transactions = {}
# Mapping pid > query count
self._transactions_count = {}
if stream is None:
stream = sys.stderr
self._stream = stream
def _colored(self, text, color=None, attrs=None):
if os.getenv('ANSI_COLORS_DISABLED') is None:
fmt_str = '\033[%dm%s' # pylint: disable=W1401
if color is not None:
text = fmt_str % (self.COLORS[color], text)
if attrs is not None:
for attr in attrs:
text = fmt_str % (self.ATTRIBUTES[attr], text)
text += self.RESET
return text
def _format_statement(self, statement, header_size):
if has_sqlparse:
statement = format_sql(statement, header_size)
replaces = []
if statement.startswith('SELECT'):
color = 'blue'
replaces = ['SELECT', 'FROM', 'WHERE', 'GROUP BY', 'JOIN', 'LEFT',
'AND', 'OR', 'ORDER BY']
elif statement.startswith('UPDATE'):
color = 'yellow'
replaces = ['UPDATE', 'SET', 'WHERE']
elif statement.startswith('INSERT INTO transaction_entry'):
# transaction entry inserting is quite common and always the same query.
# Make it less prominent
statement = self._colored(statement, 'white')
elif statement.startswith('INSERT'):
color = 'green'
replaces = ['INSERT', 'INTO', 'VALUES']
elif statement.startswith('DELETE'):
color = 'red'
replaces = ['DELETE', 'FROM', 'WHERE']
for i in replaces:
statement = statement.replace(i + ' ', self._colored(i, color) + ' ')
statement = statement.replace(i + '\n', self._colored(i, color) + '\n')
return statement
[docs] def write(self, msg):
self._stream.write(msg)
self._stream.flush()
def _expanded_raw_execute(self, transaction, raw_cursor, statement):
pid = raw_cursor.connection.get_backend_pid()
self._transactions_count.setdefault(pid, 0)
self._transactions_count[pid] += 1
count = self._transactions_count[pid]
header_size = 9 + len(str(pid))
color = self._get_transaction_color(pid)
pid = self._colored(pid, color)
self._start_time = datetime.datetime.now()
self.statement = self._format_statement(statement, header_size)
# Dont write new line, so we can print the time at the end
self.header(pid, color, count, tail=' ')
self.write(self.statement + '\n')
[docs] def connection_raw_execute_success(self, transaction, raw_cursor, statement,
params):
pid = raw_cursor.connection.get_backend_pid()
header_size = 9 + len(str(pid))
now = datetime.datetime.now()
duration = now - self._start_time
seconds = duration.seconds + float(duration.microseconds) / 10 ** 6
rows = raw_cursor.rowcount
text = '%s%s seconds | %s rows | %s' % (
' ' * header_size,
self._colored(seconds, attrs=['bold']),
self._colored(rows, attrs=['bold']),
self._colored(now.strftime('%F %H:%M:%S.%f'), attrs=['bold']))
if statement.startswith('INSERT') and rows == 1:
try:
rowid = raw_cursor.fetchone()[0]
raw_cursor.scroll(-1)
text += ' | id: ' + self._colored(repr(rowid), attrs=['bold'])
except psycopg2.ProgrammingError:
text = ''
self.write(text + '\n')
def _get_transaction_color(self, pid):
if pid not in self._transactions:
self._transactions[pid] = self._available_colors[self._current_color]
self._current_color += 1
if self._current_color == len(self._available_colors):
self._current_color = 0
return self._transactions[pid]
[docs] def transaction_create(self, store):
pid = store._connection._raw_connection.get_backend_pid()
color = self._get_transaction_color(pid)
self.header(pid, color, 'BEGIN')
[docs] def transaction_commit(self, store):
pid = store._connection._raw_connection.get_backend_pid()
color = self._get_transaction_color(pid)
self.header(pid, color, 'COMIT')
[docs] def transaction_rollback(self, store, xid=None):
pid = store._connection._raw_connection.get_backend_pid()
color = self._get_transaction_color(pid)
self.header(pid, color, 'ROLLB')
[docs] def transaction_close(self, store):
pid = store._connection._raw_connection.get_backend_pid()
color = self._get_transaction_color(pid)
del self._transactions[pid]
self.header(pid, color, 'CLOSE')
[docs]def enable():
install_tracer(StoqlibDebugTracer())