Skip to content

Commit 236ad65

Browse files
New PostgresNode implementation (#453)
1) Separated implementations for psycopg2 and pg8000 2) NodeConnection raises own exceptions: - NodeConnectionException (base class) - NodeConnectionInterfaceException - NodeConnectionDatabaseException - NodeConnectionDataException - NodeConnectionOperationalException - NodeConnectionIntegrityException - NodeConnectionInternalException - NodeConnectionProgrammingException - NodeConnectionNotSupportedException - NodeConnectionQueryCanceledException - NodeConnectionTransactionRollbackException - NodeConnectionArrayContentNotSupportedException
1 parent 0ce3d7e commit 236ad65

29 files changed

Lines changed: 790 additions & 124 deletions

File tree

‎src/connection.py‎

Lines changed: 15 additions & 109 deletions
Original file line numberDiff line numberDiff line change
@@ -1,117 +1,23 @@
11
# coding: utf-8
22
# 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+
310
try:
4-
import psycopg2 as pglib
11+
from .impl.connection_psycopg2 import NodeConnection, pglib
512
except ImportError:
613
try:
7-
import pg8000 as pglib
14+
from .impl.connection_pg8000 import NodeConnection, pglib
815
except ImportError:
916
raise ImportError("You must have psycopg2 or pg8000 modules installed")
1017

11-
from .enums import IsolationLevel
12-
13-
from .defaults import \
14-
default_dbname, \
15-
default_username2
16-
17-
from .exceptions import QueryException
18-
19-
# export some exceptions
20-
DatabaseError = pglib.DatabaseError
21-
InternalError = pglib.InternalError
22-
ProgrammingError = pglib.ProgrammingError
23-
OperationalError = pglib.OperationalError
24-
25-
26-
class NodeConnection:
27-
"""
28-
Transaction wrapper returned by Node
29-
"""
30-
def __init__(self,
31-
node,
32-
dbname=None,
33-
username=None,
34-
password=None,
35-
autocommit=False):
36-
37-
# Set default arguments
38-
dbname = dbname or default_dbname()
39-
username = username or default_username2(node.os_ops)
40-
41-
self._node = node
42-
43-
self._connection = pglib.connect(
44-
database=dbname,
45-
user=username,
46-
password=password,
47-
host=node.host,
48-
port=node.port
49-
)
50-
51-
self._connection.autocommit = autocommit
52-
self._cursor = self.connection.cursor()
53-
54-
@property
55-
def node(self):
56-
return self._node
57-
58-
@property
59-
def connection(self):
60-
return self._connection
61-
62-
@property
63-
def pid(self):
64-
return self.execute("select pg_catalog.pg_backend_pid()")[0][0]
65-
66-
@property
67-
def cursor(self):
68-
return self._cursor
69-
70-
def __enter__(self):
71-
return self
72-
73-
def __exit__(self, type, value, traceback):
74-
self.close()
75-
76-
def begin(self, isolation_level=IsolationLevel.ReadCommitted):
77-
# Check if level isn't an IsolationLevel
78-
if not isinstance(isolation_level, IsolationLevel):
79-
# Get name of isolation level
80-
level_str = str(isolation_level).lower()
81-
82-
# Validate level string
83-
try:
84-
isolation_level = IsolationLevel(level_str)
85-
except ValueError:
86-
error = 'Invalid isolation level "{}"'
87-
raise QueryException(error.format(level_str))
88-
89-
# Set isolation level
90-
cmd = 'SET TRANSACTION ISOLATION LEVEL {}'
91-
self.cursor.execute(cmd.format(isolation_level.value))
92-
93-
return self
94-
95-
def commit(self):
96-
self.connection.commit()
97-
98-
return self
99-
100-
def rollback(self):
101-
self.connection.rollback()
102-
103-
return self
104-
105-
def execute(self, query, *args):
106-
self.cursor.execute(query, args)
107-
108-
if self.cursor.description is None:
109-
return None
110-
111-
# pg8000 might return tuples
112-
res = [tuple(t) for t in self.cursor.fetchall()]
113-
return res
114-
115-
def close(self):
116-
self.cursor.close()
117-
self.connection.close()
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

‎src/exceptions.py‎

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -302,6 +302,85 @@ def __repr__(self) -> str:
302302
return result
303303

304304

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+
305384
assert ExecUtilException.__name__ == "ExecUtilException"
306385
assert InvalidOperationException.__name__ == "InvalidOperationException"
307386
assert ExecTimeoutException.__name__ == "ExecTimeoutException"

0 commit comments

Comments
 (0)