Skip to content

Commit a421dd8

Browse files
Add sync connection methods.
1 parent 2abb6bf commit a421dd8

File tree

1 file changed

+56
-8
lines changed

1 file changed

+56
-8
lines changed

db_wrapper/connection.py

Lines changed: 56 additions & 8 deletions
Original file line numberDiff line numberDiff line change
@@ -3,18 +3,18 @@
33
import asyncio
44
from dataclasses import dataclass
55
import logging
6-
from typing import Optional
6+
import time
7+
from typing import Optional, Any
78

9+
from psycopg2 import ( # type: ignore
10+
connect as psycopg2Connect,
11+
OperationalError as psycopg2OpError,
12+
)
813
from psycopg2.extras import RealDictCursor # type: ignore
9-
from psycopg2 import OperationalError # type: ignore
1014

1115
# no stubs available, starting out by just determining correct return
1216
# types & annotating in my wrappers here
13-
import aiopg # type: ignore
14-
15-
#
16-
# Postgres Connection helper
17-
#
17+
import aiopg
1818

1919
LOGGER = logging.getLogger(__name__)
2020

@@ -54,7 +54,7 @@ async def _try_connect(
5454
connection = await aiopg.connect(
5555
dsn,
5656
cursor_factory=RealDictCursor)
57-
except OperationalError as err:
57+
except psycopg2OpError as err:
5858
print(type(err))
5959
if retries > 12:
6060
raise ConnectionError(
@@ -71,10 +71,58 @@ async def _try_connect(
7171
return connection
7272

7373

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+
74115
# PENDS python 3.9 support in pylint
75116
# pylint: disable=unsubscriptable-object
76117
async def connect(
77118
connection_params: ConnectionParameters
78119
) -> aiopg.Connection:
79120
"""Establish database connection."""
80121
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

Comments
 (0)