3
3
import asyncio
4
4
from dataclasses import dataclass
5
5
import logging
6
- from typing import Optional
6
+ import time
7
+ from typing import Optional , Any
7
8
9
+ from psycopg2 import ( # type: ignore
10
+ connect as psycopg2Connect ,
11
+ OperationalError as psycopg2OpError ,
12
+ )
8
13
from psycopg2 .extras import RealDictCursor # type: ignore
9
- from psycopg2 import OperationalError # type: ignore
10
14
11
15
# no stubs available, starting out by just determining correct return
12
16
# types & annotating in my wrappers here
13
- import aiopg # type: ignore
14
-
15
- #
16
- # Postgres Connection helper
17
- #
17
+ import aiopg
18
18
19
19
LOGGER = logging .getLogger (__name__ )
20
20
@@ -54,7 +54,7 @@ async def _try_connect(
54
54
connection = await aiopg .connect (
55
55
dsn ,
56
56
cursor_factory = RealDictCursor )
57
- except OperationalError as err :
57
+ except psycopg2OpError as err :
58
58
print (type (err ))
59
59
if retries > 12 :
60
60
raise ConnectionError (
@@ -71,10 +71,58 @@ async def _try_connect(
71
71
return connection
72
72
73
73
74
+ def _sync_try_connect (
75
+ connection_params : ConnectionParameters ,
76
+ retries : int = 1
77
+ ) -> Any :
78
+ database = connection_params .database
79
+ user = connection_params .user
80
+ password = connection_params .password
81
+ host = connection_params .host
82
+ port = connection_params .port
83
+
84
+ dsn = f'dbname={ database } user={ user } password={ password } host={ host } port={ port } '
85
+
86
+ # PENDS python 3.9 support in pylint
87
+ # pylint: disable=unsubscriptable-object
88
+ connection : Optional [aiopg .Connection ] = None
89
+
90
+ LOGGER .info (
91
+ f'Attempting to connect to database { database } as { user } @{ host } :{ port } ...' )
92
+
93
+ while connection is None :
94
+ try :
95
+ connection = psycopg2Connect (
96
+ dsn ,
97
+ cursor_factory = RealDictCursor )
98
+ except psycopg2OpError as err :
99
+ print (type (err ))
100
+ if retries > 12 :
101
+ raise ConnectionError (
102
+ 'Max number of connection attempts has been reached (12)'
103
+ ) from err
104
+
105
+ LOGGER .info (
106
+ f'Connection failed ({ retries } time(s))'
107
+ 'retrying again in 5 seconds...' )
108
+
109
+ time .sleep (5 )
110
+ return _sync_try_connect (connection_params , retries + 1 )
111
+
112
+ return connection
113
+
114
+
74
115
# PENDS python 3.9 support in pylint
75
116
# pylint: disable=unsubscriptable-object
76
117
async def connect (
77
118
connection_params : ConnectionParameters
78
119
) -> aiopg .Connection :
79
120
"""Establish database connection."""
80
121
return await _try_connect (connection_params )
122
+
123
+
124
+ async def sync_connect (
125
+ connection_params : ConnectionParameters
126
+ ) -> Any :
127
+ """Establish database connection."""
128
+ return _sync_try_connect (connection_params )
0 commit comments