Skip to content

Commit 6cb3a80

Browse files
demonolockdemonolock
authored
Add work with remote host (#78)
1 parent 09e9f01 commit 6cb3a80

19 files changed

+2321
-244
lines changed

‎README.md

Lines changed: 27 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -173,6 +173,33 @@ with testgres.get_new_node().init() as master:
173173
Note that `default_conf()` is called by `init()` function; both of them overwrite
174174
the configuration file, which means that they should be called before `append_conf()`.
175175

176+
### Remote mode
177+
Testgres supports the creation of PostgreSQL nodes on a remote host. This is useful when you want to run distributed tests involving multiple nodes spread across different machines.
178+
179+
To use this feature, you need to use the RemoteOperations class.
180+
Here is an example of how you might set this up:
181+
182+
```python
183+
from testgres import ConnectionParams, RemoteOperations, TestgresConfig, get_remote_node
184+
185+
# Set up connection params
186+
conn_params = ConnectionParams(
187+
host='your_host', # replace with your host
188+
username='user_name', # replace with your username
189+
ssh_key='path_to_ssh_key' # replace with your SSH key path
190+
)
191+
os_ops = RemoteOperations(conn_params)
192+
193+
# Add remote testgres config before test
194+
TestgresConfig.set_os_ops(os_ops=os_ops)
195+
196+
# Proceed with your test
197+
def test_basic_query(self):
198+
with get_remote_node(conn_params=conn_params) as node:
199+
node.init().start()
200+
res = node.execute('SELECT 1')
201+
self.assertEqual(res, [(1,)])
202+
```
176203

177204
## Authors
178205

‎setup.py

Lines changed: 5 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -12,6 +12,9 @@
1212
"six>=1.9.0",
1313
"psutil",
1414
"packaging",
15+
"paramiko",
16+
"fabric",
17+
"sshtunnel"
1518
]
1619

1720
# Add compatibility enum class
@@ -27,9 +30,9 @@
2730
readme = f.read()
2831

2932
setup(
30-
version='1.8.9',
33+
version='1.9.0',
3134
name='testgres',
32-
packages=['testgres'],
35+
packages=['testgres', 'testgres.operations'],
3336
description='Testing utility for PostgreSQL and its extensions',
3437
url='https://github.com/postgrespro/testgres',
3538
long_description=readme,

‎testgres/__init__.py

Lines changed: 7 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,4 +1,4 @@
1-
from .api import get_new_node
1+
from .api import get_new_node, get_remote_node
22
from .backup import NodeBackup
33

44
from .config import \
@@ -46,8 +46,13 @@
4646
First, \
4747
Any
4848

49+
from .operations.os_ops import OsOperations, ConnectionParams
50+
from .operations.local_ops import LocalOperations
51+
from .operations.remote_ops import RemoteOperations
52+
4953
__all__ = [
5054
"get_new_node",
55+
"get_remote_node",
5156
"NodeBackup",
5257
"TestgresConfig", "configure_testgres", "scoped_config", "push_config", "pop_config",
5358
"NodeConnection", "DatabaseError", "InternalError", "ProgrammingError", "OperationalError",
@@ -56,4 +61,5 @@
5661
"PostgresNode", "NodeApp",
5762
"reserve_port", "release_port", "bound_ports", "get_bin_path", "get_pg_config", "get_pg_version",
5863
"First", "Any",
64+
"OsOperations", "LocalOperations", "RemoteOperations", "ConnectionParams"
5965
]

‎testgres/api.py

Lines changed: 12 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -40,3 +40,15 @@ def get_new_node(name=None, base_dir=None, **kwargs):
4040
"""
4141
# NOTE: leave explicit 'name' and 'base_dir' for compatibility
4242
return PostgresNode(name=name, base_dir=base_dir, **kwargs)
43+
44+
45+
def get_remote_node(name=None, conn_params=None):
46+
"""
47+
Simply a wrapper around :class:`.PostgresNode` constructor for remote node.
48+
See :meth:`.PostgresNode.__init__` for details.
49+
For remote connection you can add the next parameter:
50+
conn_params = ConnectionParams(host='127.0.0.1',
51+
ssh_key=None,
52+
username=default_username())
53+
"""
54+
return get_new_node(name=name, conn_params=conn_params)

‎testgres/backup.py

Lines changed: 7 additions & 11 deletions
Original file line numberDiff line numberDiff line change
@@ -2,9 +2,7 @@
22

33
import os
44

5-
from shutil import rmtree, copytree
65
from six import raise_from
7-
from tempfile import mkdtemp
86

97
from .enums import XLogMethod
108

@@ -15,8 +13,6 @@
1513
PG_CONF_FILE, \
1614
BACKUP_LOG_FILE
1715

18-
from .defaults import default_username
19-
2016
from .exceptions import BackupException
2117

2218
from .utils import \
@@ -47,7 +43,7 @@ def __init__(self,
4743
username: database user name.
4844
xlog_method: none | fetch | stream (see docs)
4945
"""
50-
46+
self.os_ops = node.os_ops
5147
if not node.status():
5248
raise BackupException('Node must be running')
5349

@@ -60,8 +56,8 @@ def __init__(self,
6056
raise BackupException(msg)
6157

6258
# Set default arguments
63-
username = username or default_username()
64-
base_dir = base_dir or mkdtemp(prefix=TMP_BACKUP)
59+
username = username or self.os_ops.get_user()
60+
base_dir = base_dir or self.os_ops.mkdtemp(prefix=TMP_BACKUP)
6561

6662
# public
6763
self.original_node = node
@@ -107,14 +103,14 @@ def _prepare_dir(self, destroy):
107103
available = not destroy
108104

109105
if available:
110-
dest_base_dir = mkdtemp(prefix=TMP_NODE)
106+
dest_base_dir = self.os_ops.mkdtemp(prefix=TMP_NODE)
111107

112108
data1 = os.path.join(self.base_dir, DATA_DIR)
113109
data2 = os.path.join(dest_base_dir, DATA_DIR)
114110

115111
try:
116112
# Copy backup to new data dir
117-
copytree(data1, data2)
113+
self.os_ops.copytree(data1, data2)
118114
except Exception as e:
119115
raise_from(BackupException('Failed to copy files'), e)
120116
else:
@@ -143,7 +139,7 @@ def spawn_primary(self, name=None, destroy=True):
143139

144140
# Build a new PostgresNode
145141
NodeClass = self.original_node.__class__
146-
with clean_on_error(NodeClass(name=name, base_dir=base_dir)) as node:
142+
with clean_on_error(NodeClass(name=name, base_dir=base_dir, conn_params=self.original_node.os_ops.conn_params)) as node:
147143

148144
# New nodes should always remove dir tree
149145
node._should_rm_dirs = True
@@ -185,4 +181,4 @@ def cleanup(self):
185181

186182
if self._available:
187183
self._available = False
188-
rmtree(self.base_dir, ignore_errors=True)
184+
self.os_ops.rmdirs(self.base_dir, ignore_errors=True)

‎testgres/cache.py

Lines changed: 12 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,7 @@
11
# coding: utf-8
22

3-
import io
43
import os
54

6-
from shutil import copytree
75
from six import raise_from
86

97
from .config import testgres_config
@@ -20,12 +18,16 @@
2018
get_bin_path, \
2119
execute_utility
2220

21+
from .operations.local_ops import LocalOperations
22+
from .operations.os_ops import OsOperations
2323

24-
def cached_initdb(data_dir, logfile=None, params=None):
24+
25+
def cached_initdb(data_dir, logfile=None, params=None, os_ops: OsOperations = LocalOperations()):
2526
"""
2627
Perform initdb or use cached node files.
2728
"""
28-
def call_initdb(initdb_dir, log=None):
29+
30+
def call_initdb(initdb_dir, log=logfile):
2931
try:
3032
_params = [get_bin_path("initdb"), "-D", initdb_dir, "-N"]
3133
execute_utility(_params + (params or []), log)
@@ -39,22 +41,23 @@ def call_initdb(initdb_dir, log=None):
3941
cached_data_dir = testgres_config.cached_initdb_dir
4042

4143
# Initialize cached initdb
42-
if not os.path.exists(cached_data_dir) or \
43-
not os.listdir(cached_data_dir):
44+
45+
if not os_ops.path_exists(cached_data_dir) or \
46+
not os_ops.listdir(cached_data_dir):
4447
call_initdb(cached_data_dir)
4548

4649
try:
4750
# Copy cached initdb to current data dir
48-
copytree(cached_data_dir, data_dir)
51+
os_ops.copytree(cached_data_dir, data_dir)
4952

5053
# Assign this node a unique system id if asked to
5154
if testgres_config.cached_initdb_unique:
5255
# XXX: write new unique system id to control file
5356
# Some users might rely upon unique system ids, but
5457
# our initdb caching mechanism breaks this contract.
5558
pg_control = os.path.join(data_dir, XLOG_CONTROL_FILE)
56-
with io.open(pg_control, "r+b") as f:
57-
f.write(generate_system_id()) # overwrite id
59+
system_id = generate_system_id()
60+
os_ops.write(pg_control, system_id, truncate=True, binary=True, read_and_write=True)
5861

5962
# XXX: build new WAL segment with our system id
6063
_params = [get_bin_path("pg_resetwal"), "-D", data_dir, "-f"]

‎testgres/config.py

Lines changed: 13 additions & 4 deletions
Original file line numberDiff line numberDiff line change
@@ -5,10 +5,10 @@
55
import tempfile
66

77
from contextlib import contextmanager
8-
from shutil import rmtree
9-
from tempfile import mkdtemp
108

119
from .consts import TMP_CACHE
10+
from .operations.os_ops import OsOperations
11+
from .operations.local_ops import LocalOperations
1212

1313

1414
class GlobalConfig(object):
@@ -43,6 +43,9 @@ class GlobalConfig(object):
4343

4444
_cached_initdb_dir = None
4545
""" underlying class attribute for cached_initdb_dir property """
46+
47+
os_ops = LocalOperations()
48+
""" OsOperation object that allows work on remote host """
4649
@property
4750
def cached_initdb_dir(self):
4851
""" path to a temp directory for cached initdb. """
@@ -54,6 +57,7 @@ def cached_initdb_dir(self, value):
5457

5558
if value:
5659
cached_initdb_dirs.add(value)
60+
return testgres_config.cached_initdb_dir
5761

5862
@property
5963
def temp_dir(self):
@@ -118,6 +122,11 @@ def copy(self):
118122

119123
return copy.copy(self)
120124

125+
@staticmethod
126+
def set_os_ops(os_ops: OsOperations):
127+
testgres_config.os_ops = os_ops
128+
testgres_config.cached_initdb_dir = os_ops.mkdtemp(prefix=TMP_CACHE)
129+
121130

122131
# cached dirs to be removed
123132
cached_initdb_dirs = set()
@@ -135,7 +144,7 @@ def copy(self):
135144
@atexit.register
136145
def _rm_cached_initdb_dirs():
137146
for d in cached_initdb_dirs:
138-
rmtree(d, ignore_errors=True)
147+
testgres_config.os_ops.rmdirs(d, ignore_errors=True)
139148

140149

141150
def push_config(**options):
@@ -198,4 +207,4 @@ def configure_testgres(**options):
198207

199208

200209
# NOTE: assign initial cached dir for initdb
201-
testgres_config.cached_initdb_dir = mkdtemp(prefix=TMP_CACHE)
210+
testgres_config.cached_initdb_dir = testgres_config.os_ops.mkdtemp(prefix=TMP_CACHE)

‎testgres/connection.py

Lines changed: 7 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -41,11 +41,11 @@ def __init__(self,
4141

4242
self._node = node
4343

44-
self._connection = pglib.connect(database=dbname,
45-
user=username,
46-
password=password,
47-
host=node.host,
48-
port=node.port)
44+
self._connection = node.os_ops.db_connect(dbname=dbname,
45+
user=username,
46+
password=password,
47+
host=node.host,
48+
port=node.port)
4949

5050
self._connection.autocommit = autocommit
5151
self._cursor = self.connection.cursor()
@@ -103,16 +103,15 @@ def rollback(self):
103103

104104
def execute(self, query, *args):
105105
self.cursor.execute(query, args)
106-
107106
try:
108107
res = self.cursor.fetchall()
109-
110108
# pg8000 might return tuples
111109
if isinstance(res, tuple):
112110
res = [tuple(t) for t in res]
113111

114112
return res
115-
except Exception:
113+
except Exception as e:
114+
print("Error executing query: {}".format(e))
116115
return None
117116

118117
def close(self):

‎testgres/defaults.py

Lines changed: 4 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -1,9 +1,9 @@
11
import datetime
2-
import getpass
3-
import os
42
import struct
53
import uuid
64

5+
from .config import testgres_config as tconf
6+
77

88
def default_dbname():
99
"""
@@ -17,8 +17,7 @@ def default_username():
1717
"""
1818
Return default username (current user).
1919
"""
20-
21-
return getpass.getuser()
20+
return tconf.os_ops.get_user()
2221

2322

2423
def generate_app_name():
@@ -44,7 +43,7 @@ def generate_system_id():
4443
system_id = 0
4544
system_id |= (secs << 32)
4645
system_id |= (usecs << 12)
47-
system_id |= (os.getpid() & 0xFFF)
46+
system_id |= (tconf.os_ops.get_pid() & 0xFFF)
4847

4948
# pack ULL in native byte order
5049
return struct.pack('=Q', system_id)

0 commit comments

Comments
 (0)
pFad - Phonifier reborn

Pfad - The Proxy pFad of © 2024 Garber Painting. All rights reserved.

Note: This service is not intended for secure transactions such as banking, social media, email, or purchasing. Use at your own risk. We assume no liability whatsoever for broken pages.


Alternative Proxies:

Alternative Proxy

pFad Proxy

pFad v3 Proxy

pFad v4 Proxy