diff options
author | Karl-Aksel Puulmann <oxymaccy@gmail.com> | 2015-01-10 11:55:35 +0200 |
---|---|---|
committer | Karl-Aksel Puulmann <oxymaccy@gmail.com> | 2015-01-10 11:55:35 +0200 |
commit | 89547a171e79f00c0b3e708436ba051b7626e90e (patch) | |
tree | 4a00ba388f98086e4c8afe5d4247173e35595ccd | |
parent | a939a9a4915d64d5f616ef77c3a54bbf98126efb (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.py | 34 | ||||
-rw-r--r-- | tests/test_pgexecute.py | 39 |
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) |