Skip to content

Commit

Permalink
add get_connection function
Browse files Browse the repository at this point in the history
  • Loading branch information
sourBz committed Dec 12, 2024
1 parent 617d38a commit 9529979
Showing 1 changed file with 12 additions and 3 deletions.
15 changes: 12 additions & 3 deletions postgrestq/task_queue.py
Original file line number Diff line number Diff line change
Expand Up @@ -14,7 +14,7 @@
Sequence,
)

from psycopg import sql, connect
from psycopg import sql, connect, Connection

# supported only from 3.11 onwards:
# from datetime import UTC
Expand Down Expand Up @@ -69,19 +69,28 @@ def __init__(
# called when ttl <= 0 for a task
self.ttl_zero_callback = ttl_zero_callback

self.connect()
self.conn = self.connect()
if create_table:
self._create_queue_table()

if reset:
self._reset()

def get_connection(self):
connection = connect(self._dsn)
with connection.cursor() as cur:
cur.execute("SELECT 1+1")
cur.fetchone()

return connection

def connect(self) -> None:
"""
Establish a connection to Postgres.
If a connection already exists, it's overwritten.
"""
self.conn = connect(self._dsn)
if self.conn is None or self.conn.closed:
self.conn = self.get_connection()

def _create_queue_table(self) -> None:
"""
Expand Down

0 comments on commit 9529979

Please sign in to comment.