summaryrefslogtreecommitdiffstats
diff options
context:
space:
mode:
authorKarl-Aksel Puulmann <oxymaccy@gmail.com>2015-01-10 11:55:35 +0200
committerKarl-Aksel Puulmann <oxymaccy@gmail.com>2015-01-10 11:55:35 +0200
commit89547a171e79f00c0b3e708436ba051b7626e90e (patch)
tree4a00ba388f98086e4c8afe5d4247173e35595ccd
parenta939a9a4915d64d5f616ef77c3a54bbf98126efb (diff)
Support having multiple queries on the same line.
Note that if any of the queries raises an error, nothing besides the error is output. This is different behavior from psql which outputs separately for each query: _test_db=# select 1; invalid syntax; select 2; ?column? ---------- 1 (1 row) ERROR: syntax error at or near "invalid" LINE 1: invalid syntax; ^ ?column? ---------- 2 (1 row)
-rw-r--r--pgcli/pgexecute.py34
-rw-r--r--tests/test_pgexecute.py39
2 files changed, 53 insertions, 20 deletions
diff --git a/pgcli/pgexecute.py b/pgcli/pgexecute.py
index 27af3085..0f97d4f3 100644
--- a/pgcli/pgexecute.py
+++ b/pgcli/pgexecute.py
@@ -2,6 +2,7 @@ import logging
import psycopg2
import psycopg2.extras
import psycopg2.extensions
+import sqlparse
from collections import defaultdict
from .packages import pgspecial
@@ -89,11 +90,12 @@ class PGExecute(object):
are a list of tuples. Each tuple has 3 values (rows, headers, status).
"""
+ # Remove spaces and EOL
+ sql = sql.strip()
if not sql: # Empty string
return [(None, None, None)]
# Remove spaces, eol and semi-colons.
- sql = sql.strip()
sql = sql.rstrip(';')
# Check if the command is a \c or 'use'. This is a special exception
@@ -106,31 +108,35 @@ class PGExecute(object):
except:
_logger.debug('Database name missing.')
raise RuntimeError('Database name missing.')
- self.conn = psycopg2.connect(database=dbname,
- user=self.user, password=self.password, host=self.host,
- port=self.port)
self.dbname = dbname
- self.conn.autocommit = True
+ self.connect()
_logger.debug('Successfully switched to DB: %r', dbname)
return [(None, None, 'You are now connected to database "%s" as '
'user "%s"' % (self.dbname, self.user))]
- with self.conn.cursor() as cur:
- try:
- _logger.debug('Trying a pgspecial command. sql: %r', sql)
+ # Special command
+ try:
+ _logger.debug('Trying a pgspecial command. sql: %r', sql)
+ with self.conn.cursor() as cur:
return pgspecial.execute(cur, sql)
- except KeyError:
- _logger.debug('Regular sql statement. sql: %r', sql)
- cur.execute(sql)
-
+ except KeyError:
+ # 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.
+ 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)]
+ return (cur.fetchall(), headers, cur.statusmessage)
else:
_logger.debug('No rows in result.')
- return [(None, None, cur.statusmessage)]
+ return (None, None, cur.statusmessage)
def tables(self):
""" Returns tuple (sorted_tables, columns). Columns is a dictionary of
diff --git a/tests/test_pgexecute.py b/tests/test_pgexecute.py
index a918ec75..6888067c 100644
--- a/tests/test_pgexecute.py
+++ b/tests/test_pgexecute.py
@@ -69,16 +69,43 @@ def test_database_list(executor):
databases = executor.databases()
assert '_test_db' in databases
-@pytest.fixture(params=[True, False])
-def expanded(request):
- return request.param
+@dbtest
+def test_invalid_syntax(executor):
+ with pytest.raises(psycopg2.ProgrammingError) as excinfo:
+ run(executor, 'invalid syntax!')
+ assert 'syntax error at or near "invalid"' in str(excinfo.value)
@dbtest
-def test_unicode_support_in_output(executor, expanded):
- if expanded:
+def test_invalid_column_name(executor):
+ with pytest.raises(psycopg2.ProgrammingError) as excinfo:
+ run(executor, 'select invalid command')
+ assert 'column "invalid" does not exist' in str(excinfo.value)
+
+@pytest.yield_fixture(params=[True, False])
+def expanded(request, executor):
+ if request.param:
+ run(executor, '\\x')
+ yield request.param
+ if request.param:
run(executor, '\\x')
+
+@dbtest
+def test_unicode_support_in_output(executor, expanded):
run(executor, "create table unicodechars(t text)")
run(executor, "insert into unicodechars (t) values ('é')")
# See issue #24, this raises an exception without proper handling
- assert u'é' in run(executor, "select * from unicodechars", join=True) \ No newline at end of file
+ assert u'é' in run(executor, "select * from unicodechars", join=True)
+
+@dbtest
+def test_multiple_queries_same_line(executor):
+ result = run(executor, "select 'foo'; select 'bar'")
+ assert len(result) == 4 # 2 * (output+status)
+ assert "foo" in result[0]
+ assert "bar" in result[2]
+
+@dbtest
+def test_multiple_queries_same_line_syntaxerror(executor):
+ with pytest.raises(psycopg2.ProgrammingError) as excinfo:
+ run(executor, "select 'foo'; invalid syntax")
+ assert 'syntax error at or near "invalid"' in str(excinfo.value)