44 lines
1.0 KiB
Python
44 lines
1.0 KiB
Python
|
from __future__ import annotations
|
||
|
|
|
||
|
|
from contextlib import contextmanager
|
||
|
|
from time import sleep
|
||
|
|
from typing import Generator
|
||
|
|
|
||
|
|
import psycopg2
|
||
|
|
import psycopg2.extensions
|
||
|
|
import psycopg2.pool
|
||
|
|
|
||
|
|
from ajusta_bling.common import Args
|
||
|
|
|
||
|
|
|
||
|
|
class Database:
|
||
|
|
pool: psycopg2.pool.SimpleConnectionPool | None = None
|
||
|
|
|
||
|
|
def __init__(self, args: Args):
|
||
|
|
self.user = args.dbuser
|
||
|
|
self.pswd = args.dbpassword
|
||
|
|
self.host = args.dbhost
|
||
|
|
self.port = args.dbport
|
||
|
|
self.db = args.dbname
|
||
|
|
|
||
|
|
def connect(self):
|
||
|
|
self.pool = psycopg2.pool.SimpleConnectionPool(
|
||
|
|
1, 10,
|
||
|
|
user = self.user,
|
||
|
|
password = self.pswd,
|
||
|
|
host = self.host,
|
||
|
|
port = self.port,
|
||
|
|
database = self.db,
|
||
|
|
)
|
||
|
|
|
||
|
|
@contextmanager
|
||
|
|
def get_cur(self) -> Generator[psycopg2.extensions.cursor]:
|
||
|
|
if self.pool is None:
|
||
|
|
self.connect()
|
||
|
|
conn = self.pool.getconn()
|
||
|
|
try:
|
||
|
|
yield conn.cursor()
|
||
|
|
conn.commit()
|
||
|
|
finally:
|
||
|
|
self.pool.putconn(conn)
|