1
0
Fork 0

Consume the sql results as a generator.

This commit is contained in:
Amjith Ramanujam 2015-01-17 19:24:02 -08:00
parent 05865f4b39
commit bf8736a6c0
3 changed files with 31 additions and 23 deletions

View File

@ -182,17 +182,21 @@ class PGCli(object):
logger.debug('sql: %r', document.text)
successful = False
start = time()
# Initialized to None because res might never get
# initialized if an exception occurs in pgexecute.run().
# Which causes finally clause to fail.
res = None
res = pgexecute.run(document.text)
duration = time() - start
successful = True
output = []
total = 0
for rows, headers, status in res:
for cur, headers, status in res:
logger.debug("headers: %r", headers)
logger.debug("rows: %r", rows)
logger.debug("rows: %r", cur)
logger.debug("status: %r", status)
start = time()
output.extend(format_output(rows, headers, status))
output.extend(format_output(cur, headers, status))
end = time()
total += end - start
mutating = mutating or is_mutating(status)
@ -210,7 +214,11 @@ class PGCli(object):
if is_timing_enabled():
print('Command Time:', duration)
print('Render Time:', total)
print('Format Time:', total)
finally:
for cur, _, _ in res:
if hasattr(cur, 'close'):
cur.close()
# Refresh the table names and column names if necessary.
if need_completion_refresh(document.text):

View File

@ -58,7 +58,7 @@ def list_schemas(cur, pattern, verbose):
cur.execute(sql)
if cur.description:
headers = [x[0] for x in cur.description]
return [(cur.fetchall(), headers, cur.statusmessage)]
return [(cur, headers, cur.statusmessage)]
def describe_table_details(cur, pattern, verbose):
"""
@ -90,7 +90,7 @@ def describe_table_details(cur, pattern, verbose):
cur.execute(sql)
if cur.description:
headers = [x[0] for x in cur.description]
return [(cur.fetchall(), headers, cur.statusmessage)]
return [(cur, headers, cur.statusmessage)]
# This is a \d <tablename> command. A royal pain in the ass.
schema, relname = sql_name_pattern(pattern)
@ -860,7 +860,7 @@ def execute(cur, sql):
cur.execute(command_executor)
if cur.description:
headers = [x[0] for x in cur.description]
return [(cur.fetchall(), headers, cur.statusmessage)]
return [(cur, headers, cur.statusmessage)]
else:
return [(None, None, cur.statusmessage)]

View File

@ -117,29 +117,29 @@ class PGExecute(object):
return [(None, None, 'You are now connected to database "%s" as '
'user "%s"' % (self.dbname, self.user))]
# Special command
try:
try: # Special command
_logger.debug('Trying a pgspecial command. sql: %r', sql)
with self.conn.cursor() as cur:
return pgspecial.execute(cur, sql)
except KeyError:
cur = self.conn.cursor()
return pgspecial.execute(cur, sql)
except KeyError: # Regular SQL
# Split the sql into separate queries and run each one. If any
# single query fails, the rest are not run and no results are shown.
# single query fails, the rest of them are not run and no results
# are shown.
queries = sqlparse.split(sql)
return [self.execute_normal_sql(query) for query in queries]
def execute_normal_sql(self, split_sql):
_logger.debug('Regular sql statement. sql: %r', split_sql)
with self.conn.cursor() as cur:
cur.execute(split_sql)
# cur.description will be None for operations that do not return
# rows.
if cur.description:
headers = [x[0] for x in cur.description]
return (cur.fetchall(), headers, cur.statusmessage)
else:
_logger.debug('No rows in result.')
return (None, None, cur.statusmessage)
cur = self.conn.cursor()
cur.execute(split_sql)
# cur.description will be None for operations that do not return
# rows.
if cur.description:
headers = [x[0] for x in cur.description]
return (cur, headers, cur.statusmessage)
else:
_logger.debug('No rows in result.')
return (None, None, cur.statusmessage)
def tables(self):
""" Returns tuple (sorted_tables, columns). Columns is a dictionary of