Enable job insertion on both transactions and top level pool (#9)

A problem with the API currently is that is that inserts on a top-level
SQLAlchemy engine aren't supported. In the test suite, the driver is
always initialized with a transaction:

    with engine.begin() as tx:
        yield riversqlalchemy.Driver(tx)

But if trying to initialize the driver with the engine itself, we'd see
failures because it doesn't support the necessary SQL execution methods
required to insert a job. Additionally, when checking for a unique job,
the driver always opens a transaction with `self.conn.begin_nested()`,
which will error unless in a transaction already.

In the Ruby client there's only a single `River#insert` method because
convention in that language is to have a lot of global stuff going on,
and the current transaction is added to local thread storage by Sequel
or ActiveRecord so that we can determine whether or not we're already in
a transaction from within the client.

SQLAlchemy is a little different. After opening a transaction, you're
expected to track the session object for perform operations in it:

    with session.begin():
        session.add(some_object())
        session.add(some_other_object())

So here I'm proposing a client API that looks a little more like the Go
API, with a pair of insert functions for `#insert` and `#insert_tx`, the
former doing an insert on the connection originally passed to driver,
and the latter doing an insert on a transaction sent as argument.

So this style of top-level insertion is now possible:

    engine = sqlalchemy.create_engine(database_url)
    client = riverqueue.Client(riversqlalchemy.Driver(engine))

    insert_res = client.insert(
        MyJobArgs(),
        insert_opts=riverqueue.InsertOpts(
            unique_opts=riverqueue.UniqueOpts(by_period=900)
        ),
    )

The client opens a new transaction on the driver's engine, and inserts
the job there.

Insertion on a particular transaction continues to be possible, and now
without having to reinitialize the client/driver:

    with engine.begin() as session:
        insert_res = client.insert_tx(
            session,
            MyJobArgs(),
            insert_opts=riverqueue.InsertOpts(
                unique_opts=riverqueue.UniqueOpts(by_period=900)
            ),
        )
        print(insert_res)

Author: brandur

src/riverqueue/client.py

@@ -2,7 +2,7 @@
 from datetime import datetime, timezone, timedelta
 from typing import Any, Optional, Protocol, Tuple, List, Callable
 
-from .driver import GetParams, JobInsertParams, DriverProtocol
+from .driver import GetParams, JobInsertParams, DriverProtocol, ExecutorProtocol
 from .model import InsertResult
 from .fnv import fnv1_hash
 
@@ -53,29 +53,56 @@ def __init__(self, driver: DriverProtocol, advisory_lock_prefix=None):
     def insert(
         self, args: Args, insert_opts: Optional[InsertOpts] = None
     ) -> InsertResult:
+        with self.driver.executor() as exec:
+            if not insert_opts:
+                insert_opts = InsertOpts()
+            insert_params, unique_opts = self.__make_insert_params(args, insert_opts)
+
+            def insert():
+                return InsertResult(exec.job_insert(insert_params))
+
+            return self.__check_unique_job(exec, insert_params, unique_opts, insert)
+
+    def insert_tx(
+        self, tx, args: Args, insert_opts: Optional[InsertOpts] = None
+    ) -> InsertResult:
+        exec = self.driver.unwrap_executor(tx)
         if not insert_opts:
             insert_opts = InsertOpts()
         insert_params, unique_opts = self.__make_insert_params(args, insert_opts)
 
         def insert():
-            print(self.driver)
-            return InsertResult(self.driver.job_insert(insert_params))
+            return InsertResult(exec.job_insert(insert_params))
 
-        return self.__check_unique_job(insert_params, unique_opts, insert)
+        return self.__check_unique_job(exec, insert_params, unique_opts, insert)
 
     def insert_many(self, args: List[Args]) -> List[InsertResult]:
-        all_params = [
+        with self.driver.executor() as exec:
+            return [
+                InsertResult(x)
+                for x in exec.job_insert_many(self.__make_insert_params_many(args))
+            ]
+
+    def insert_many_tx(self, tx, args: List[Args]) -> List[InsertResult]:
+        exec = self.driver.unwrap_executor(tx)
+        return [
+            InsertResult(x)
+            for x in exec.job_insert_many(self.__make_insert_params_many(args))
+        ]
+
+    def __make_insert_params_many(self, args: List[Args]) -> List[JobInsertParams]:
+        return [
             self.__make_insert_params(
                 arg.args, arg.insert_opts or InsertOpts(), is_insert_many=True
             )[0]
             if isinstance(arg, InsertManyParams)
             else self.__make_insert_params(arg, InsertOpts(), is_insert_many=True)[0]
             for arg in args
         ]
-        return [InsertResult(x) for x in self.driver.job_insert_many(all_params)]
 
     def __check_unique_job(
         self,
+        exec: ExecutorProtocol,
         insert_params: JobInsertParams,
         unique_opts: Optional[UniqueOpts],
         insert_func: Callable[[], InsertResult],
@@ -125,25 +152,27 @@ def __check_unique_job(
         if not any_unique_opts:
             return insert_func()
 
-        with self.driver.transaction():
+        with exec.transaction():
             if self.advisory_lock_prefix is None:
                 lock_key = fnv1_hash(lock_str.encode("utf-8"), 64)
             else:
                 prefix = self.advisory_lock_prefix
                 lock_key = (prefix << 32) | fnv1_hash(lock_str.encode("utf-8"), 32)
 
             lock_key = self.__uint64_to_int64(lock_key)
-            self.driver.advisory_lock(lock_key)
+            exec.advisory_lock(lock_key)
 
-            existing_job = self.driver.job_get_by_kind_and_unique_properties(get_params)
+            existing_job = exec.job_get_by_kind_and_unique_properties(get_params)
             if existing_job:
                 return InsertResult(existing_job, unique_skipped_as_duplicated=True)
 
             return insert_func()
 
     @staticmethod
     def __make_insert_params(
-        args: Args, insert_opts: InsertOpts, is_insert_many: bool = False
+        args: Args,
+        insert_opts: InsertOpts,
+        is_insert_many: bool = False,
     ) -> Tuple[JobInsertParams, Optional[UniqueOpts]]:
         if not hasattr(args, "kind"):
             raise Exception("args should respond to `kind`")

src/riverqueue/driver/__init__.py

@@ -1,5 +1,6 @@
 # Reexport for more ergonomic use in calling code.
 from .driver_protocol import (
+    ExecutorProtocol as ExecutorProtocol,
     GetParams as GetParams,
     JobInsertParams as JobInsertParams,
     DriverProtocol as DriverProtocol,

src/riverqueue/driver/driver_protocol.py

@@ -1,6 +1,7 @@
+from contextlib import _GeneratorContextManager, contextmanager
 from dataclasses import dataclass, field
 from datetime import datetime
-from typing import Any, List, ContextManager, Optional, Protocol
+from typing import Any, Iterator, List, Optional, Protocol
 
 from ..model import Job
 
@@ -34,11 +35,11 @@ class JobInsertParams:
     finalized_at: Optional[datetime] = None
 
 
-class DriverProtocol(Protocol):
+class ExecutorProtocol(Protocol):
     def advisory_lock(self, lock: int) -> None:
         pass
 
-    def job_insert(self, insert_params: JobInsertParams) -> Optional[Job]:
+    def job_insert(self, insert_params: JobInsertParams) -> Job:
         pass
 
     def job_insert_many(self, all_params) -> List[Job]:
@@ -49,5 +50,14 @@ def job_get_by_kind_and_unique_properties(
     ) -> Optional[Job]:
         pass
 
-    def transaction(self) -> ContextManager:
+    def transaction(self) -> _GeneratorContextManager:
+        pass
+
+
+class DriverProtocol(Protocol):
+    @contextmanager
+    def executor(self) -> Iterator[ExecutorProtocol]:
+        pass
+
+    def unwrap_executor(self, tx) -> ExecutorProtocol:
         pass

src/riverqueue/driver/riversqlalchemy/sql_alchemy_driver.py

@@ -1,33 +1,64 @@
 from contextlib import contextmanager
-from typing import Optional, List, Generator
+from sqlalchemy import Engine
+from sqlalchemy.engine import Connection
+from typing import Iterator, Optional, List, Generator, cast
 
-from ...driver import DriverProtocol, GetParams, JobInsertParams
+from ...driver import DriverProtocol, ExecutorProtocol, GetParams, JobInsertParams
 from ...model import Job
 from . import river_job, pg_misc
 
 
-class Driver(DriverProtocol):
-    def __init__(self, session):
-        self.session = session
-        self.pg_misc_querier = pg_misc.Querier(session)
-        self.job_querier = river_job.Querier(session)
+class Executor(ExecutorProtocol):
+    def __init__(self, conn: Connection):
+        self.conn = conn
+        self.pg_misc_querier = pg_misc.Querier(conn)
+        self.job_querier = river_job.Querier(conn)
 
     def advisory_lock(self, key: int) -> None:
         self.pg_misc_querier.pg_advisory_xact_lock(key=key)
 
-    def job_insert(self, insert_params: JobInsertParams) -> Optional[Job]:
-        return self.job_querier.job_insert_fast(insert_params)
+    def job_insert(self, insert_params: JobInsertParams) -> Job:
+        return cast(
+            Job,
+            self.job_querier.job_insert_fast(
+                cast(river_job.JobInsertFastParams, insert_params)
+            ),
+        )
 
     def job_insert_many(self, all_params) -> List[Job]:
         raise NotImplementedError("sqlc doesn't implement copy in python yet")
 
     def job_get_by_kind_and_unique_properties(
         self, get_params: GetParams
     ) -> Optional[Job]:
-        return self.job_querier.job_get_by_kind_and_unique_properties(get_params)
+        return cast(
+            Optional[Job],
+            self.job_querier.job_get_by_kind_and_unique_properties(
+                cast(river_job.JobGetByKindAndUniquePropertiesParams, get_params)
+            ),
+        )
 
     @contextmanager
     def transaction(self) -> Generator:
-        session = self.session
-        with session.begin_nested():
-            yield
+        if self.conn.in_transaction():
+            with self.conn.begin_nested():
+                yield
+        else:
+            with self.conn.begin():
+                yield
+
+
+class Driver(DriverProtocol):
+    def __init__(self, conn: Connection | Engine):
+        self.conn = conn
+
+    @contextmanager
+    def executor(self) -> Iterator[ExecutorProtocol]:
+        if isinstance(self.conn, Engine):
+            with self.conn.begin() as tx:
+                yield Executor(tx)
+        else:
+            yield Executor(self.conn)
+
+    def unwrap_executor(self, tx) -> ExecutorProtocol:
+        return Executor(tx)