@@ -105,6 +105,7 @@ class ResultSet(list, ColumnGuesserMixin):
105105
106106 def __init__ (self , sqlaproxy , sql , config ):
107107 self .keys = sqlaproxy .keys ()
108+ self .status = getattr (sqlaproxy , 'statusmessage' , None )
108109 self .sql = sql
109110 self .config = config
110111 self .limit = config .autolimit
@@ -128,6 +129,20 @@ def _repr_html_(self):
128129 self .pretty .add_rows (self )
129130 result = self .pretty .get_html_string ()
130131 result = _cell_with_spaces_pattern .sub (_nonbreaking_spaces , result )
132+ if self .status :
133+ status_tables = []
134+ status_table = None
135+ for line in self .status .splitlines ():
136+ if not line .startswith (' ' ):
137+ if status_table :
138+ status_tables .append (status_table )
139+ status_table = PrettyTable ([line ])
140+ else :
141+ status_table .add_row ([line .strip ()])
142+ status_tables .append (status_table )
143+ result = '\n ' .join (
144+ [result ] +
145+ [st .get_html_string () for st in status_tables ])
131146 if self .config .displaylimit and len (
132147 self ) > self .config .displaylimit :
133148 result = '%s\n <span style="font-style:italic;text-align:center;">%d rows, truncated to displaylimit of %d</span>' % (
@@ -296,13 +311,41 @@ class FakeResultProxy(object):
296311 SqlAlchemy.
297312 """
298313
299- def __init__ (self , cursor , headers ):
300- self .fetchall = cursor . fetchall
301- self .fetchmany = cursor . fetchmany
302- self .rowcount = cursor . rowcount
314+ def __init__ (self , cur , headers , status ):
315+ self ._cursor = cur
316+ self ._cursor_index = 0
317+ self ._status = status
303318 self .keys = lambda : headers
304319 self .returns_rows = True
305320
321+ def fetchall (self ):
322+ if isinstance (self ._cursor , list ):
323+ self ._cursor_index = self .rowcount
324+ return self ._cursor
325+ return self ._cursor .fetchall ()
326+
327+
328+ def fetchmany (self , size ):
329+ if isinstance (self ._cursor , list ):
330+ prev = self ._cursor_index
331+ next = prev + size
332+ self ._cursor_index = next
333+ return self ._cursor [prev :next ]
334+ return self ._cursor .fetchmany (size )
335+
336+ @property
337+ def rowcount (self ):
338+ if isinstance (self ._cursor , list ):
339+ return len (self ._cursor )
340+ return self ._cursor .rowcount
341+
342+ @property
343+ def statusmessage (self ):
344+ if isinstance (self ._cursor , list ):
345+ return self ._status
346+ return self ._cursor .statusmessage
347+
348+
306349# some dialects have autocommit
307350# specific dialects break when commit is used:
308351_COMMIT_BLACKLIST_DIALECTS = ('mssql' , 'clickhouse' , 'teradata' )
@@ -332,9 +375,9 @@ def run(conn, sql, config, user_namespace):
332375 if not PGSpecial :
333376 raise ImportError ('pgspecial not installed' )
334377 pgspecial = PGSpecial ()
335- _ , cur , headers , _ = pgspecial .execute (
378+ _ , cur , headers , status = pgspecial .execute (
336379 conn .session .connection .cursor (), statement )[0 ]
337- result = FakeResultProxy (cur , headers )
380+ result = FakeResultProxy (cur , headers , status )
338381 else :
339382 txt = sqlalchemy .sql .text (statement )
340383 result = conn .session .execute (txt , user_namespace )
0 commit comments