Postgres support

.github/workflows/build.yml

@@ -5,7 +5,7 @@ on:
 jobs:
   lint:
     if: github.event_name == 'push' && !startsWith(github.event.ref, 'refs/tags')
-    runs-on: ubuntu-latest
+    runs-on: ubuntu-22.04
     strategy:
       fail-fast: false
       matrix:
@@ -41,7 +41,7 @@ jobs:
           - '3.9'
           - '3.10'
         os:
-          - ubuntu-latest
+          - ubuntu-22.04
     services:
       redis-node-one:
         image: redis
@@ -69,6 +69,18 @@ jobs:
           MONGODB_REPLICA_SET_MODE: primary
           MONGODB_REPLICA_SET_NAME: test_replica_set
           ALLOW_EMPTY_PASSWORD: 'yes'
+      postgres-node-one:
+        image: 'postgres:latest'
+        ports:
+          - '5432:5432'
+        env:
+          POSTGRES_PASSWORD: 'mysecretpassword'
+      postgres-node-two:
+        image: 'postgres:latest'
+        ports:
+          - '5433:5432'
+        env:
+          POSTGRES_PASSWORD: 'mysecretpassword'
     steps:
       - name: Checkout
         uses: actions/checkout@v3

pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "poetry.core.masonry.api"
 
 [tool.poetry]
 name = "sergeant"
-version = "0.25.0"
+version = "0.26.0"
 readme = "README.md"
 homepage = "https://github.com/Intsights/sergeant"
 repository = "https://github.com/Intsights/sergeant"
@@ -39,6 +39,7 @@ psutil = "^5"
 pymongo = ">=3.0,<5.0"
 redis = "^4"
 typing_extensions = "^3.10"
+psycopg = {version = "^3", extras = ["binary"]}
 
 [tool.poetry.dev-dependencies]
 pytest = "^7"

sergeant/config.py

@@ -46,7 +46,7 @@ class Encoder:
 )
 class Connector:
     params: typing.Dict[str, typing.Any]
-    type: typing_extensions.Literal['redis', 'mongo', 'local'] = 'local'
+    type: typing_extensions.Literal['redis', 'mongo', 'local', 'postgres'] = 'local'
 
 
 @dataclasses.dataclass(

sergeant/connector/__init__.py

@@ -1,6 +1,7 @@
 from . import _connector
 from . import local
 from . import mongo
+from . import postgres
 from . import redis
 
 

sergeant/connector/local.py

@@ -35,7 +35,7 @@ def acquire(
                 expire_at = time.time() + ttl
                 self.connection.execute(
                     '''
-                        INSERT INTO locks (name, expireAt)
+                        INSERT INTO locks (name, expire_at)
                         VALUES(?, ?);
                     ''',
                     (
@@ -55,16 +55,16 @@ def acquire(
     def release(
         self,
     ) -> bool:
-        self.connection.execute(
-            '''
-                DELETE FROM locks WHERE expireAt < ?;
-            ''',
-            (
-                time.time(),
-            ),
-        )
-
         if self.acquired:
+            self.connection.execute(
+                '''
+                    DELETE FROM locks WHERE expire_at < ?;
+                ''',
+                (
+                    time.time(),
+                ),
+            )
+
             cursor = self.connection.execute(
                 '''
                     DELETE FROM locks WHERE name = ?;
@@ -85,7 +85,7 @@ def is_locked(
         cursor = self.connection.execute(
             '''
                 SELECT * FROM locks
-                WHERE name = ? AND expireAt > ?;
+                WHERE name = ? AND expire_at > ?;
             ''',
             (
                 self.name,
@@ -107,8 +107,8 @@ def set_ttl(
         cursor = self.connection.execute(
             '''
                 UPDATE locks
-                SET expireAt = ?
-                WHERE name = ? AND expireAt > ?;
+                SET expire_at = ?
+                WHERE name = ? AND expire_at > ?;
             ''',
             (
                 now + ttl,
@@ -125,7 +125,7 @@ def get_ttl(
         now = time.time()
         cursor = self.connection.execute(
             '''
-                SELECT expireAt FROM locks
+                SELECT expire_at FROM locks
                 WHERE name = ?;
             ''',
             (
@@ -173,9 +173,9 @@ def __init__(
                 CREATE TABLE IF NOT EXISTS keys (name TEXT, value BLOB);
                 CREATE UNIQUE INDEX IF NOT EXISTS key_by_name ON keys (name);
 
-                CREATE TABLE IF NOT EXISTS locks (name TEXT, expireAt REAL);
+                CREATE TABLE IF NOT EXISTS locks (name TEXT, expire_at REAL);
                 CREATE UNIQUE INDEX IF NOT EXISTS lock_by_name ON locks (name);
-                CREATE INDEX IF NOT EXISTS lock_by_expireAt ON locks (expireAt);
+                CREATE INDEX IF NOT EXISTS lock_by_expire_at ON locks (expire_at);
             '''
         )
 

sergeant/connector/mongo.py

@@ -412,15 +412,15 @@ def queue_length(
     ) -> int:
         queue_length = 0
 
-        for i in range(self.number_of_connections):
+        for connection in self.connections:
             if include_delayed:
-                queue_length += self.next_connection.sergeant.task_queue.count_documents(
+                queue_length += connection.sergeant.task_queue.count_documents(
                     filter={
                         'queue_name': queue_name,
                     },
                 )
             else:
-                queue_length += self.next_connection.sergeant.task_queue.count_documents(
+                queue_length += connection.sergeant.task_queue.count_documents(
                     filter={
                         'queue_name': queue_name,
                         'priority': {