testgres 1.16.4__tar.gz → 1.17.1__tar.gz

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.
Files changed (56) hide show
  1. {testgres-1.16.4/testgres.egg-info → testgres-1.17.1}/PKG-INFO +1 -1
  2. {testgres-1.16.4 → testgres-1.17.1}/src/__init__.py +3 -1
  3. {testgres-1.16.4 → testgres-1.17.1}/src/backup.py +56 -9
  4. testgres-1.17.1/src/connection.py +23 -0
  5. {testgres-1.16.4 → testgres-1.17.1}/src/exceptions.py +88 -1
  6. testgres-1.17.1/src/impl/connection_pg8000.py +197 -0
  7. testgres-1.17.1/src/impl/connection_psycopg2.py +203 -0
  8. testgres-1.17.1/src/impl/error_utils.py +113 -0
  9. {testgres-1.16.4 → testgres-1.17.1}/src/node.py +97 -37
  10. {testgres-1.16.4 → testgres-1.17.1}/src/pubsub.py +1 -1
  11. {testgres-1.16.4 → testgres-1.17.1/testgres.egg-info}/PKG-INFO +1 -1
  12. {testgres-1.16.4 → testgres-1.17.1}/testgres.egg-info/SOURCES.txt +3 -0
  13. {testgres-1.16.4 → testgres-1.17.1}/tests/test_testgres_common.py +238 -16
  14. testgres-1.16.4/src/connection.py +0 -121
  15. {testgres-1.16.4 → testgres-1.17.1}/LICENSE +0 -0
  16. {testgres-1.16.4 → testgres-1.17.1}/README.md +0 -0
  17. {testgres-1.16.4 → testgres-1.17.1}/pyproject.toml +0 -0
  18. {testgres-1.16.4 → testgres-1.17.1}/setup.cfg +0 -0
  19. {testgres-1.16.4 → testgres-1.17.1}/src/api.py +0 -0
  20. {testgres-1.16.4 → testgres-1.17.1}/src/cache.py +0 -0
  21. {testgres-1.16.4 → testgres-1.17.1}/src/config.py +0 -0
  22. {testgres-1.16.4 → testgres-1.17.1}/src/consts.py +0 -0
  23. {testgres-1.16.4 → testgres-1.17.1}/src/decorators.py +0 -0
  24. {testgres-1.16.4 → testgres-1.17.1}/src/defaults.py +0 -0
  25. {testgres-1.16.4 → testgres-1.17.1}/src/enums.py +0 -0
  26. {testgres-1.16.4 → testgres-1.17.1}/src/impl/file_line_reader.py +0 -0
  27. {testgres-1.16.4 → testgres-1.17.1}/src/impl/internal_node_utils.py +0 -0
  28. {testgres-1.16.4 → testgres-1.17.1}/src/impl/internal_utils.py +0 -0
  29. {testgres-1.16.4 → testgres-1.17.1}/src/impl/platforms/darwin/internal_platform_utils.py +0 -0
  30. {testgres-1.16.4 → testgres-1.17.1}/src/impl/platforms/internal_platform_utils.py +0 -0
  31. {testgres-1.16.4 → testgres-1.17.1}/src/impl/platforms/internal_platform_utils_factory.py +0 -0
  32. {testgres-1.16.4 → testgres-1.17.1}/src/impl/platforms/internal_platform_utils_provider.py +0 -0
  33. {testgres-1.16.4 → testgres-1.17.1}/src/impl/platforms/linux/internal_platform_utils.py +0 -0
  34. {testgres-1.16.4 → testgres-1.17.1}/src/impl/platforms/win32/internal_platform_utils.py +0 -0
  35. {testgres-1.16.4 → testgres-1.17.1}/src/impl/port_manager__generic.py +0 -0
  36. {testgres-1.16.4 → testgres-1.17.1}/src/impl/port_manager__generic2.py +0 -0
  37. {testgres-1.16.4 → testgres-1.17.1}/src/impl/port_manager__this_host.py +0 -0
  38. {testgres-1.16.4 → testgres-1.17.1}/src/logger.py +0 -0
  39. {testgres-1.16.4 → testgres-1.17.1}/src/node_app.py +0 -0
  40. {testgres-1.16.4 → testgres-1.17.1}/src/node_state.py +0 -0
  41. {testgres-1.16.4 → testgres-1.17.1}/src/port_manager.py +0 -0
  42. {testgres-1.16.4 → testgres-1.17.1}/src/raise_error.py +0 -0
  43. {testgres-1.16.4 → testgres-1.17.1}/src/standby.py +0 -0
  44. {testgres-1.16.4 → testgres-1.17.1}/src/utils.py +0 -0
  45. {testgres-1.16.4 → testgres-1.17.1}/testgres.egg-info/dependency_links.txt +0 -0
  46. {testgres-1.16.4 → testgres-1.17.1}/testgres.egg-info/requires.txt +0 -0
  47. {testgres-1.16.4 → testgres-1.17.1}/testgres.egg-info/top_level.txt +0 -0
  48. {testgres-1.16.4 → testgres-1.17.1}/tests/test_api.py +0 -0
  49. {testgres-1.16.4 → testgres-1.17.1}/tests/test_config.py +0 -0
  50. {testgres-1.16.4 → testgres-1.17.1}/tests/test_os_ops_common.py +0 -0
  51. {testgres-1.16.4 → testgres-1.17.1}/tests/test_os_ops_local.py +0 -0
  52. {testgres-1.16.4 → testgres-1.17.1}/tests/test_os_ops_remote.py +0 -0
  53. {testgres-1.16.4 → testgres-1.17.1}/tests/test_raise_error.py +0 -0
  54. {testgres-1.16.4 → testgres-1.17.1}/tests/test_testgres_local.py +0 -0
  55. {testgres-1.16.4 → testgres-1.17.1}/tests/test_testgres_remote.py +0 -0
  56. {testgres-1.16.4 → testgres-1.17.1}/tests/test_utils.py +0 -0
@@ -1,6 +1,6 @@
1
1
  Metadata-Version: 2.4
2
2
  Name: testgres
3
- Version: 1.16.4
3
+ Version: 1.17.1
4
4
  Summary: Testing utility for PostgreSQL and its extensions
5
5
  Author-email: Postgres Professional <testgres@postgrespro.ru>
6
6
  License: PostgreSQL
@@ -24,6 +24,7 @@ from .exceptions import \
24
24
  TimeoutException, \
25
25
  CatchUpException, \
26
26
  StartNodeException, \
27
+ StopNodeException, \
27
28
  InitNodeException, \
28
29
  BackupException, \
29
30
  InvalidOperationException
@@ -60,7 +61,7 @@ from testgres.operations.os_ops import OsProcessController
60
61
  from testgres.operations.local_ops import LocalOperations
61
62
  from testgres.operations.remote_ops import RemoteOperations
62
63
 
63
- __version__ = "1.16.4"
64
+ __version__ = "1.17.1"
64
65
 
65
66
  __all__ = [
66
67
  "get_new_node",
@@ -73,6 +74,7 @@ __all__ = [
73
74
  "QueryTimeoutException",
74
75
  "TimeoutException", "CatchUpException", "StartNodeException", "InitNodeException", "BackupException", "InvalidOperationException",
75
76
  "XLogMethod", "IsolationLevel", "NodeStatus", "ProcessType", "DumpFormat",
77
+ "StopNodeException",
76
78
  "NodeApp",
77
79
  "PostgresNode",
78
80
  "PortManager",
@@ -35,7 +35,8 @@ class NodeBackup:
35
35
  base_dir=None,
36
36
  username=None,
37
37
  xlog_method=XLogMethod.fetch,
38
- options=None):
38
+ options=None,
39
+ os_ops=None):
39
40
  """
40
41
  Create a new backup.
41
42
 
@@ -44,13 +45,18 @@ class NodeBackup:
44
45
  base_dir: where should we store it?
45
46
  username: database user name.
46
47
  xlog_method: none | fetch | stream (see docs)
48
+ os_ops: execution context where pg_basebackup is run and the
49
+ backup is stored. If omitted, use the source node's os_ops.
50
+ The source node must be reachable from this context using its
51
+ configured host and port.
47
52
  """
48
53
  assert node.os_ops is not None
49
54
  assert isinstance(node.os_ops, OsOperations)
55
+ assert os_ops is None or isinstance(os_ops, OsOperations)
50
56
 
51
57
  if not options:
52
58
  options = []
53
- self.os_ops = node.os_ops
59
+ self.os_ops = os_ops if os_ops is not None else node.os_ops
54
60
  if not node.status():
55
61
  raise BackupException('Node must be running')
56
62
 
@@ -63,7 +69,9 @@ class NodeBackup:
63
69
  raise BackupException(msg)
64
70
 
65
71
  # Set default arguments
66
- username = username or self.os_ops.get_user()
72
+ # The username belongs to the source PostgreSQL instance, not to the
73
+ # environment where pg_basebackup happens to run.
74
+ username = username or node.os_ops.get_user()
67
75
  base_dir = base_dir or self.os_ops.mkdtemp(prefix=TMP_BACKUP)
68
76
 
69
77
  # public
@@ -139,13 +147,24 @@ class NodeBackup:
139
147
  # Return path to new node
140
148
  return dest_base_dir
141
149
 
142
- def spawn_primary(self, name=None, destroy=True):
150
+ def spawn_primary(
151
+ self,
152
+ name=None,
153
+ destroy=True,
154
+ target_host=None,
155
+ target_port=None,
156
+ ):
143
157
  """
144
158
  Create a primary node from a backup.
145
159
 
146
160
  Args:
147
161
  name: primary's application name.
148
162
  destroy: should we convert this backup into a node?
163
+ target_host: host/address for the new node. If omitted, preserve
164
+ the source host when the backup uses the source os_ops,
165
+ otherwise use the backup os_ops host.
166
+ target_port: explicit port for the new node. If omitted, reserve a
167
+ port in the backup's execution context.
149
168
 
150
169
  Returns:
151
170
  New instance of :class:`.PostgresNode`.
@@ -157,19 +176,37 @@ class NodeBackup:
157
176
  # Build a new PostgresNode
158
177
  assert self.original_node is not None
159
178
 
160
- node = self.original_node.clone_with_new_name_and_base_dir(name=name, base_dir=base_dir)
179
+ node = self.original_node.clone_with_new_name_and_base_dir(
180
+ name=name,
181
+ base_dir=base_dir,
182
+ target_os_ops=self.os_ops,
183
+ target_host=target_host,
184
+ target_port=target_port,
185
+ )
161
186
 
162
187
  assert node is not None
163
188
  assert type(node) is self.original_node.__class__
164
189
 
165
190
  with clean_on_error(node) as node:
166
- # Set a new port
191
+ # The backup contains the source node's connection settings.
192
+ # Override settings that belong to the spawned node.
167
193
  node.append_conf(filename=PG_CONF_FILE, line='\n')
168
- node.append_conf(filename=PG_CONF_FILE, port=node.port)
194
+ node.append_conf(
195
+ filename=PG_CONF_FILE,
196
+ listen_addresses=node.host,
197
+ port=node.port,
198
+ )
169
199
 
170
200
  return node
171
201
 
172
- def spawn_replica(self, name=None, destroy=True, slot=None):
202
+ def spawn_replica(
203
+ self,
204
+ name=None,
205
+ destroy=True,
206
+ slot=None,
207
+ target_host=None,
208
+ target_port=None,
209
+ ):
173
210
  """
174
211
  Create a replica of the original node from a backup.
175
212
 
@@ -177,13 +214,23 @@ class NodeBackup:
177
214
  name: replica's application name.
178
215
  slot: create a replication slot with the specified name.
179
216
  destroy: should we convert this backup into a node?
217
+ target_host: host/address for the replica. If omitted, preserve
218
+ the source host when the backup uses the source os_ops,
219
+ otherwise use the backup os_ops host.
220
+ target_port: explicit port for the replica. If omitted, reserve a
221
+ port in the backup's execution context.
180
222
 
181
223
  Returns:
182
224
  New instance of :class:`.PostgresNode`.
183
225
  """
184
226
 
185
227
  # Build a new PostgresNode
186
- node = self.spawn_primary(name=name, destroy=destroy)
228
+ node = self.spawn_primary(
229
+ name=name,
230
+ destroy=destroy,
231
+ target_host=target_host,
232
+ target_port=target_port,
233
+ )
187
234
  assert node is not None
188
235
 
189
236
  try:
@@ -0,0 +1,23 @@
1
+ # coding: utf-8
2
+ # we support both pg8000 and psycopg2
3
+
4
+ # to backward compatibility
5
+ from .exceptions import NodeConnectionDatabaseException as DatabaseError
6
+ from .exceptions import NodeConnectionOperationalException as OperationalError
7
+ from .exceptions import NodeConnectionInternalException as InternalError
8
+ from .exceptions import NodeConnectionProgrammingException as ProgrammingError
9
+
10
+ try:
11
+ from .impl.connection_psycopg2 import NodeConnection, pglib
12
+ except ImportError:
13
+ try:
14
+ from .impl.connection_pg8000 import NodeConnection, pglib
15
+ except ImportError:
16
+ raise ImportError("You must have psycopg2 or pg8000 modules installed")
17
+
18
+ assert DatabaseError.__name__ == "NodeConnectionDatabaseException"
19
+ assert OperationalError.__name__ == "NodeConnectionOperationalException"
20
+ assert InternalError.__name__ == "NodeConnectionInternalException"
21
+ assert ProgrammingError.__name__ == "NodeConnectionProgrammingException"
22
+ assert NodeConnection.__name__ == "NodeConnection"
23
+ assert pglib is not None
@@ -160,7 +160,7 @@ class CatchUpException(TestgresException):
160
160
  return result
161
161
 
162
162
 
163
- class StartNodeException(TestgresException):
163
+ class NodeException(TestgresException):
164
164
  _description: typing.Optional[str]
165
165
  _files: typing.Optional[typing.Iterable]
166
166
 
@@ -224,6 +224,14 @@ class StartNodeException(TestgresException):
224
224
  return result
225
225
 
226
226
 
227
+ class StartNodeException(NodeException):
228
+ pass
229
+
230
+
231
+ class StopNodeException(NodeException):
232
+ pass
233
+
234
+
227
235
  class InitNodeException(TestgresException):
228
236
  _message: typing.Optional[str]
229
237
 
@@ -294,6 +302,85 @@ class BackupException(TestgresException):
294
302
  return result
295
303
 
296
304
 
305
+ class NodeConnectionException(TestgresException):
306
+ _message: typing.Optional[str]
307
+
308
+ def __init__(
309
+ self,
310
+ message: typing.Optional[str] = None,
311
+ ):
312
+ assert message is None or type(message) is str
313
+ super().__init__(message)
314
+ self._message = message
315
+ return
316
+
317
+ @property
318
+ def message(self) -> str:
319
+ assert self._message is None or type(self._message) is str
320
+ if self._message is None:
321
+ return ""
322
+ return self._message
323
+
324
+ def __repr__(self) -> str:
325
+ args = []
326
+
327
+ if self._message is not None:
328
+ args.append(("message", self._message))
329
+
330
+ result = "{}(".format(type(self).__name__)
331
+ sep = ""
332
+ for a in args:
333
+ result += sep + a[0] + "=" + repr(a[1])
334
+ sep = ", "
335
+ continue
336
+ result += ")"
337
+ return result
338
+
339
+
340
+ class NodeConnectionInterfaceException(NodeConnectionException):
341
+ pass
342
+
343
+
344
+ class NodeConnectionDatabaseException(NodeConnectionException):
345
+ pass
346
+
347
+
348
+ class NodeConnectionDataException(NodeConnectionDatabaseException):
349
+ pass
350
+
351
+
352
+ class NodeConnectionOperationalException(NodeConnectionDatabaseException):
353
+ pass
354
+
355
+
356
+ class NodeConnectionIntegrityException(NodeConnectionDatabaseException):
357
+ pass
358
+
359
+
360
+ class NodeConnectionInternalException(NodeConnectionDatabaseException):
361
+ pass
362
+
363
+
364
+ class NodeConnectionProgrammingException(NodeConnectionDatabaseException):
365
+ pass
366
+
367
+
368
+ class NodeConnectionNotSupportedException(NodeConnectionDatabaseException):
369
+ pass
370
+
371
+
372
+ class NodeConnectionQueryCanceledException(NodeConnectionOperationalException):
373
+ pass
374
+
375
+
376
+ class NodeConnectionTransactionRollbackException(NodeConnectionOperationalException):
377
+ pass
378
+
379
+
380
+ class NodeConnectionArrayContentNotSupportedException(NodeConnectionNotSupportedException):
381
+ pass
382
+
383
+
297
384
  assert ExecUtilException.__name__ == "ExecUtilException"
298
385
  assert InvalidOperationException.__name__ == "InvalidOperationException"
299
386
  assert ExecTimeoutException.__name__ == "ExecTimeoutException"
@@ -0,0 +1,197 @@
1
+ # coding: utf-8
2
+ # we support both pg8000 and psycopg2
3
+
4
+ from ..enums import IsolationLevel
5
+
6
+ from ..defaults import \
7
+ default_dbname, \
8
+ default_username2
9
+
10
+ from ..exceptions import QueryException
11
+ from ..exceptions import NodeConnectionException
12
+ from ..exceptions import NodeConnectionInterfaceException
13
+ from ..exceptions import NodeConnectionDatabaseException
14
+ from ..exceptions import NodeConnectionDataException
15
+ from ..exceptions import NodeConnectionOperationalException
16
+ from ..exceptions import NodeConnectionIntegrityException
17
+ from ..exceptions import NodeConnectionInternalException
18
+ from ..exceptions import NodeConnectionProgrammingException
19
+ from ..exceptions import NodeConnectionNotSupportedException
20
+ from ..exceptions import NodeConnectionArrayContentNotSupportedException
21
+
22
+ import typing
23
+
24
+ import pg8000 as pglib
25
+ import pg8000.dbapi as pglib_dbapi
26
+
27
+
28
+ class NodeConnection:
29
+ """
30
+ Transaction wrapper returned by Node
31
+ """
32
+ def __init__(
33
+ self,
34
+ node,
35
+ dbname=None,
36
+ username=None,
37
+ password=None,
38
+ autocommit=False,
39
+ ):
40
+ # Set default arguments
41
+ dbname = dbname or default_dbname()
42
+ username = username or default_username2(node.os_ops)
43
+
44
+ self._node = node
45
+
46
+ c = __class__._exec_dblib(
47
+ pglib.connect,
48
+ database=dbname,
49
+ user=username,
50
+ password=password,
51
+ host=node.host,
52
+ port=node.port,
53
+ )
54
+ assert isinstance(c, pglib.Connection)
55
+
56
+ self._connection = c
57
+
58
+ self._connection.autocommit = autocommit
59
+
60
+ c = self.connection.cursor()
61
+ assert isinstance(c, pglib.Cursor)
62
+ self._cursor = c
63
+
64
+ @property
65
+ def node(self):
66
+ return self._node
67
+
68
+ @property
69
+ def connection(self):
70
+ return self._connection
71
+
72
+ @property
73
+ def pid(self):
74
+ r = self.execute("select pg_catalog.pg_backend_pid()")
75
+ assert r is not None
76
+ return r[0][0]
77
+
78
+ @property
79
+ def cursor(self):
80
+ return self._cursor
81
+
82
+ def __enter__(self):
83
+ return self
84
+
85
+ def __exit__(self, type, value, traceback):
86
+ self.close()
87
+ return
88
+
89
+ def begin(self, isolation_level=IsolationLevel.ReadCommitted):
90
+ # Check if level isn't an IsolationLevel
91
+ if not isinstance(isolation_level, IsolationLevel):
92
+ # Get name of isolation level
93
+ level_str = str(isolation_level).lower()
94
+
95
+ # Validate level string
96
+ try:
97
+ isolation_level = IsolationLevel(level_str)
98
+ except ValueError:
99
+ error = 'Invalid isolation level "{}"'
100
+ raise QueryException(error.format(level_str))
101
+
102
+ # Set isolation level
103
+ cmd_f = 'SET TRANSACTION ISOLATION LEVEL {}'
104
+ cmd = cmd_f.format(isolation_level.value)
105
+
106
+ __class__._exec_dblib(
107
+ self.cursor.execute,
108
+ cmd,
109
+ )
110
+ return self
111
+
112
+ def commit(self):
113
+ __class__._exec_dblib(
114
+ self.connection.commit,
115
+ )
116
+ return self
117
+
118
+ def rollback(self):
119
+ __class__._exec_dblib(
120
+ self.connection.rollback,
121
+ )
122
+
123
+ return self
124
+
125
+ def execute(self, query, *args):
126
+ __class__._exec_dblib(
127
+ self.cursor.execute,
128
+ query,
129
+ args,
130
+ )
131
+
132
+ if self.cursor.description is None:
133
+ return None
134
+
135
+ # pg8000 might return tuples
136
+ r = __class__._exec_dblib(
137
+ self.cursor.fetchall,
138
+ )
139
+ assert type(r) is tuple
140
+ res = [tuple(t) for t in r]
141
+ return res
142
+
143
+ def close(self):
144
+ __class__._exec_dblib(
145
+ self.cursor.close,
146
+ )
147
+ __class__._exec_dblib(
148
+ self.connection.close,
149
+ )
150
+ return
151
+
152
+ @staticmethod
153
+ def _exec_dblib(
154
+ func: typing.Callable,
155
+ *args,
156
+ **kwargs,
157
+ ) -> typing.Any:
158
+ try:
159
+ r = func(*args, **kwargs)
160
+ except Exception as e:
161
+ __class__._map_pglib_exception(e)
162
+
163
+ return r
164
+
165
+ @staticmethod
166
+ def _map_pglib_exception(exc: Exception) -> typing.NoReturn:
167
+ if isinstance(exc, pglib_dbapi.ArrayContentNotSupportedError):
168
+ raise NodeConnectionArrayContentNotSupportedException(str(exc)) from exc
169
+
170
+ if isinstance(exc, pglib.NotSupportedError):
171
+ raise NodeConnectionNotSupportedException(str(exc)) from exc
172
+
173
+ if isinstance(exc, pglib.InternalError):
174
+ raise NodeConnectionInternalException(str(exc)) from exc
175
+
176
+ if isinstance(exc, pglib.ProgrammingError):
177
+ raise NodeConnectionProgrammingException(str(exc)) from exc
178
+
179
+ if isinstance(exc, pglib.IntegrityError):
180
+ raise NodeConnectionIntegrityException(str(exc)) from exc
181
+
182
+ if isinstance(exc, pglib.OperationalError):
183
+ raise NodeConnectionOperationalException(str(exc)) from exc
184
+
185
+ if isinstance(exc, pglib.DataError):
186
+ raise NodeConnectionDataException(str(exc)) from exc
187
+
188
+ if isinstance(exc, pglib.DatabaseError):
189
+ raise NodeConnectionDatabaseException(str(exc)) from exc
190
+
191
+ if isinstance(exc, pglib.InterfaceError):
192
+ raise NodeConnectionInterfaceException(str(exc)) from exc
193
+
194
+ if isinstance(exc, pglib.Error):
195
+ raise NodeConnectionException(str(exc)) from exc
196
+
197
+ raise exc
@@ -0,0 +1,203 @@
1
+ # coding: utf-8
2
+ # we support both pg8000 and psycopg2
3
+
4
+ from ..enums import IsolationLevel
5
+
6
+ from ..defaults import \
7
+ default_dbname, \
8
+ default_username2
9
+
10
+ from ..exceptions import QueryException
11
+ from ..exceptions import NodeConnectionException
12
+ from ..exceptions import NodeConnectionInterfaceException
13
+ from ..exceptions import NodeConnectionDatabaseException
14
+ from ..exceptions import NodeConnectionDataException
15
+ from ..exceptions import NodeConnectionOperationalException
16
+ from ..exceptions import NodeConnectionIntegrityException
17
+ from ..exceptions import NodeConnectionInternalException
18
+ from ..exceptions import NodeConnectionProgrammingException
19
+ from ..exceptions import NodeConnectionNotSupportedException
20
+ from ..exceptions import NodeConnectionTransactionRollbackException
21
+ from ..exceptions import NodeConnectionQueryCanceledException
22
+
23
+ import typing
24
+
25
+ import psycopg2 as pglib
26
+ from psycopg2._psycopg import connection as PSYCOPG2__Connection
27
+ from psycopg2._psycopg import cursor as PSYCOPG2__Cursor
28
+ from psycopg2._psycopg import QueryCanceledError as PSYCOPG2__QueryCanceledError
29
+ from psycopg2._psycopg import TransactionRollbackError as PSYCOPG2__TransactionRollbackError
30
+
31
+
32
+ class NodeConnection:
33
+ """
34
+ Transaction wrapper returned by Node
35
+ """
36
+ def __init__(
37
+ self,
38
+ node,
39
+ dbname=None,
40
+ username=None,
41
+ password=None,
42
+ autocommit=False,
43
+ ):
44
+ # Set default arguments
45
+ dbname = dbname or default_dbname()
46
+ username = username or default_username2(node.os_ops)
47
+
48
+ self._node = node
49
+
50
+ c = __class__._exec_dblib(
51
+ pglib.connect,
52
+ database=dbname,
53
+ user=username,
54
+ password=password,
55
+ host=node.host,
56
+ port=node.port,
57
+ )
58
+ assert isinstance(c, PSYCOPG2__Connection)
59
+
60
+ self._connection = c
61
+
62
+ self._connection.autocommit = autocommit
63
+
64
+ c = self.connection.cursor()
65
+ assert isinstance(c, PSYCOPG2__Cursor)
66
+ self._cursor = c
67
+
68
+ @property
69
+ def node(self):
70
+ return self._node
71
+
72
+ @property
73
+ def connection(self):
74
+ return self._connection
75
+
76
+ @property
77
+ def pid(self):
78
+ r = self.execute("select pg_catalog.pg_backend_pid()")
79
+ assert r is not None
80
+ return r[0][0]
81
+
82
+ @property
83
+ def cursor(self):
84
+ return self._cursor
85
+
86
+ def __enter__(self):
87
+ return self
88
+
89
+ def __exit__(self, type, value, traceback):
90
+ self.close()
91
+ return
92
+
93
+ def begin(self, isolation_level=IsolationLevel.ReadCommitted):
94
+ # Check if level isn't an IsolationLevel
95
+ if not isinstance(isolation_level, IsolationLevel):
96
+ # Get name of isolation level
97
+ level_str = str(isolation_level).lower()
98
+
99
+ # Validate level string
100
+ try:
101
+ isolation_level = IsolationLevel(level_str)
102
+ except ValueError:
103
+ error = 'Invalid isolation level "{}"'
104
+ raise QueryException(error.format(level_str))
105
+
106
+ # Set isolation level
107
+ cmd_f = 'SET TRANSACTION ISOLATION LEVEL {}'
108
+ cmd = cmd_f.format(isolation_level.value)
109
+
110
+ __class__._exec_dblib(
111
+ self.cursor.execute,
112
+ cmd,
113
+ )
114
+ return self
115
+
116
+ def commit(self):
117
+ __class__._exec_dblib(
118
+ self.connection.commit,
119
+ )
120
+ return self
121
+
122
+ def rollback(self):
123
+ __class__._exec_dblib(
124
+ self.connection.rollback,
125
+ )
126
+
127
+ return self
128
+
129
+ def execute(self, query, *args):
130
+ __class__._exec_dblib(
131
+ self.cursor.execute,
132
+ query,
133
+ args,
134
+ )
135
+
136
+ if self.cursor.description is None:
137
+ return None
138
+
139
+ # note: pg8000 might return tuples
140
+ res = __class__._exec_dblib(
141
+ self.cursor.fetchall,
142
+ )
143
+ assert type(res) is list
144
+ return res
145
+
146
+ def close(self):
147
+ __class__._exec_dblib(
148
+ self.cursor.close,
149
+ )
150
+ __class__._exec_dblib(
151
+ self.connection.close,
152
+ )
153
+ return
154
+
155
+ @staticmethod
156
+ def _exec_dblib(
157
+ func: typing.Callable,
158
+ *args,
159
+ **kwargs,
160
+ ) -> typing.Any:
161
+ try:
162
+ r = func(*args, **kwargs)
163
+ except Exception as e:
164
+ __class__._map_pglib_exception(e)
165
+
166
+ return r
167
+
168
+ @staticmethod
169
+ def _map_pglib_exception(exc: Exception) -> typing.NoReturn:
170
+ if isinstance(exc, PSYCOPG2__QueryCanceledError):
171
+ raise NodeConnectionQueryCanceledException(str(exc)) from exc
172
+
173
+ if isinstance(exc, PSYCOPG2__TransactionRollbackError):
174
+ raise NodeConnectionTransactionRollbackException(str(exc)) from exc
175
+
176
+ if isinstance(exc, pglib.NotSupportedError):
177
+ raise NodeConnectionNotSupportedException(str(exc)) from exc
178
+
179
+ if isinstance(exc, pglib.InternalError):
180
+ raise NodeConnectionInternalException(str(exc)) from exc
181
+
182
+ if isinstance(exc, pglib.ProgrammingError):
183
+ raise NodeConnectionProgrammingException(str(exc)) from exc
184
+
185
+ if isinstance(exc, pglib.IntegrityError):
186
+ raise NodeConnectionIntegrityException(str(exc)) from exc
187
+
188
+ if isinstance(exc, pglib.OperationalError):
189
+ raise NodeConnectionOperationalException(str(exc)) from exc
190
+
191
+ if isinstance(exc, pglib.DataError):
192
+ raise NodeConnectionDataException(str(exc)) from exc
193
+
194
+ if isinstance(exc, pglib.DatabaseError):
195
+ raise NodeConnectionDatabaseException(str(exc)) from exc
196
+
197
+ if isinstance(exc, pglib.InterfaceError):
198
+ raise NodeConnectionInterfaceException(str(exc)) from exc
199
+
200
+ if isinstance(exc, pglib.Error):
201
+ raise NodeConnectionException(str(exc)) from exc
202
+
203
+ raise exc