diff --git a/README.md b/README.md index c018f591..0df8d94d 100644 --- a/README.md +++ b/README.md @@ -29,7 +29,7 @@ import testgres node = None try: node = testgres.get_new_node('test').init().start() - print(node.execute('postgres', 'select 1')) + print(node.execute('postgres','127.0.0.1', 'select 1')) except testgres.ClusterException as e: print(e) finally: @@ -66,7 +66,7 @@ Finally our temporary cluster is able to process queries. There are four ways to * `node.psql(database, query)` - runs query via `psql` command and returns tuple `(error code, stdout, stderr)` * `node.safe_psql(database, query)` - same as `psql()` except that it returns only `stdout`. If an error occures during the execution, an exception will be thrown. -* `node.execute(database, query)` - connects to postgresql server using `psycopg2` or `pg8000` library (depends on which is installed in your system) and returns two-dimensional array with data. +* `node.execute(database, host, query)` - connects to postgresql server using `psycopg2` or `pg8000` library (depends on which is installed in your system) and returns two-dimensional array with data. * `node.connect(database='postgres')` - returns connection wrapper (`NodeConnection`) capable of running several queries within a single transaction. The last one is the most powerful: you can use `begin(isolation_level)`, `commit()` and `rollback()`: diff --git a/testgres/testgres.py b/testgres/testgres.py index 91ce4cfa..78cd973f 100644 --- a/testgres/testgres.py +++ b/testgres/testgres.py @@ -70,14 +70,14 @@ class NodeConnection(object): Transaction wrapper returned by Node """ - def __init__(self, parent_node, dbname): + def __init__(self, parent_node, dbname, host='127.0.0.1'): self.parent_node = parent_node self.connection = pglib.connect( database=dbname, user=get_username(), port=parent_node.port, - host="127.0.0.1" + host=host ) self.cursor = self.connection.cursor() @@ -111,11 +111,9 @@ def begin(self, isolation_level=0): # Something is wrong, emit exception else: - raise QueryException('Invalid isolation level "{}"'.format( - isolation_level)) + raise QueryException('Invalid isolation level "{}"'.format(isolation_level)) - self.cursor.execute('SET TRANSACTION ISOLATION LEVEL {}'.format( - isolation_level)) + self.cursor.execute('SET TRANSACTION ISOLATION LEVEL {}'.format(isolation_level)) def commit(self): self.connection.commit() @@ -135,9 +133,9 @@ def close(self): class PostgresNode(object): - def __init__(self, name, port, base_dir=None): + def __init__(self, name, port, host='127.0.0.1', base_dir=None): self.name = name - self.host = '127.0.0.1' + self.host = host self.port = port if base_dir is None: self.base_dir = tempfile.mkdtemp() @@ -439,9 +437,15 @@ def poll_query_until(self, dbname, query): attemps += 1 raise QueryException("Timeout while waiting for query to return True") - def execute(self, dbname, query): - """Executes the query and returns all rows""" - with self.connect(dbname) as node_con: + def execute(self, dbname, host, query): + """Executes the query and returns all rows + + :param dbname: string, database name + :param host: string, hostname or ip address of database + :param query: string, query for database + :return : connection object + """ + with self.connect(dbname, host) as node_con: return node_con.execute(query) def backup(self, name): @@ -463,8 +467,8 @@ def backup(self, name): return backup_path - def connect(self, dbname='postgres'): - return NodeConnection(parent_node=self, dbname=dbname) + def connect(self, dbname='postgres', host='127.0.0.1' ): + return NodeConnection(parent_node=self, dbname=dbname, host=host) def get_username(): diff --git a/testgres/tests/test_simple.py b/testgres/tests/test_simple.py index ac33377e..79583375 100644 --- a/testgres/tests/test_simple.py +++ b/testgres/tests/test_simple.py @@ -15,7 +15,7 @@ def test_start_stop(self): node = get_new_node('test') node.init() node.start() - res = node.execute('postgres', 'select 1') + res = node.execute('postgres', '127.0.0.1', 'select 1') self.assertEqual(len(res), 1) self.assertEqual(res[0][0], 1) node.stop() @@ -32,7 +32,7 @@ def test_backup_and_replication(self): replica.init_from_backup(node, 'my_backup', has_streaming=True) replica.start() - res = replica.execute('postgres', 'select * from abc') + res = replica.execute('postgres', '127.0.0.1', 'select * from abc') self.assertEqual(len(res), 1) self.assertEqual(res[0], (1, 2)) @@ -46,7 +46,7 @@ def test_backup_and_replication(self): % replica.name) # time.sleep(0.5) # Check that this record was exported to replica - res = replica.execute('postgres', 'select * from abc') + res = replica.execute('postgres', '127.0.0.1', 'select * from abc') self.assertEqual(len(res), 2) self.assertEqual(res[1], (3, 4))