Better structure by putting types in more appropriate files and reexporting (#8)

Follows up #5 to keep its API, but undo a lot of the file structure
changes. A big downside of #5 is that a lot o the core code ends up in
`__init__.py` files, which are really hard to fast find because they all
have the same name, and generally just awkward to work with.

Python docs make this impossible to find, but I found a trick in
SQLAlchemy where it's possible to reexport types from a module using an
`as` in an `import` like this:

    # Reexport for more ergonomic use in calling code.
    from .driver_protocol import (
        GetParams as GetParams,
        JobInsertParams as JobInsertParams,
        DriverProtocol as DriverProtocol,
    )

Here, move code back into files with names better suited for it (e.g.
`Client` goes back in `client.py`) and use this trick to reexport
user-facing types so that we keep the same import ergonomic benefits as
laid out in #5.

Author: brandur

src/riverqueue/__init__.py

@@ -1,189 +1,7 @@
-from dataclasses import dataclass
-from datetime import datetime, timezone, timedelta
-from typing import Any, Optional, Protocol, Tuple, List, Callable
-
-from .driver import GetParams, JobInsertParams, DriverProtocol
-from .model import InsertResult
-from .fnv import fnv1_hash
-
-MAX_ATTEMPTS_DEFAULT = 25
-PRIORITY_DEFAULT = 1
-QUEUE_DEFAULT = "default"
-
-DEFAULT_UNIQUE_STATES = ["available", "completed", "running", "retryable", "scheduled"]
-
-
-class Args(Protocol):
-    kind: str
-
-    def to_json(self) -> str:
-        pass
-
-
-@dataclass
-class InsertManyParams:
-    args: Args
-    insert_opts: Optional["InsertOpts"] = None
-
-
-@dataclass
-class InsertOpts:
-    scheduled_at: Optional[datetime] = None
-    unique_opts: Optional["UniqueOpts"] = None
-    max_attempts: Optional[int] = None
-    priority: Optional[int] = None
-    queue: Optional[str] = None
-    tags: Optional[List[Any]] = None
-
-
-@dataclass
-class UniqueOpts:
-    by_args: Optional[Any] = None
-    by_period: Optional[Any] = None
-    by_queue: Optional[Any] = None
-    by_state: Optional[Any] = None
-
-
-class Client:
-    def __init__(self, driver: DriverProtocol, advisory_lock_prefix=None):
-        self.driver = driver
-        self.advisory_lock_prefix = advisory_lock_prefix
-        self.time_now_utc = lambda: datetime.now(timezone.utc)  # for test time stubbing
-
-    def insert(
-        self, args: Args, insert_opts: Optional[InsertOpts] = None
-    ) -> InsertResult:
-        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 self.__check_unique_job(insert_params, unique_opts, insert)
-
-    def insert_many(self, args: List[Args]) -> List[InsertResult]:
-        all_params = [
-            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,
-        insert_params: JobInsertParams,
-        unique_opts: Optional[UniqueOpts],
-        insert_func: Callable[[], InsertResult],
-    ) -> InsertResult:
-        if unique_opts is None:
-            return insert_func()
-
-        any_unique_opts = False
-        get_params = GetParams(kind=insert_params.kind)
-
-        lock_str = f"unique_keykind={insert_params.kind}"
-
-        if unique_opts.by_args:
-            any_unique_opts = True
-            get_params.by_args = True
-            get_params.args = insert_params.args
-            lock_str += f"&args={insert_params.args}"
-
-        if unique_opts.by_period:
-            lower_period_bound = self.__truncate_time(
-                self.time_now_utc(), unique_opts.by_period
-            )
-
-            any_unique_opts = True
-            get_params.by_created_at = True
-            get_params.created_at = [
-                lower_period_bound,
-                lower_period_bound + timedelta(seconds=unique_opts.by_period),
-            ]
-            lock_str += f"&period={lower_period_bound.strftime('%FT%TZ')}"
-
-        if unique_opts.by_queue:
-            any_unique_opts = True
-            get_params.by_queue = True
-            get_params.queue = insert_params.queue
-            lock_str += f"&queue={insert_params.queue}"
-
-        if unique_opts.by_state:
-            any_unique_opts = True
-            get_params.by_state = True
-            get_params.state = unique_opts.by_state
-            lock_str += f"&state={','.join(unique_opts.by_state)}"
-        else:
-            get_params.state = DEFAULT_UNIQUE_STATES
-            lock_str += f"&state={','.join(DEFAULT_UNIQUE_STATES)}"
-
-        if not any_unique_opts:
-            return insert_func()
-
-        with self.driver.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)
-
-            existing_job = self.driver.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
-    ) -> Tuple[JobInsertParams, Optional[UniqueOpts]]:
-        if not hasattr(args, "kind"):
-            raise Exception("args should respond to `kind`")
-
-        args_json = args.to_json()
-        if args_json is None:
-            raise Exception("args should return non-nil from `to_json`")
-
-        args_insert_opts = getattr(args, "insert_opts", InsertOpts())
-
-        scheduled_at = insert_opts.scheduled_at or args_insert_opts.scheduled_at
-        unique_opts = insert_opts.unique_opts or args_insert_opts.unique_opts
-
-        if is_insert_many and unique_opts:
-            raise ValueError("unique opts can't be used with `insert_many`")
-
-        insert_params = JobInsertParams(
-            args=args_json,
-            kind=args.kind,
-            max_attempts=insert_opts.max_attempts
-            or args_insert_opts.max_attempts
-            or MAX_ATTEMPTS_DEFAULT,
-            priority=insert_opts.priority
-            or args_insert_opts.priority
-            or PRIORITY_DEFAULT,
-            queue=insert_opts.queue or args_insert_opts.queue or QUEUE_DEFAULT,
-            scheduled_at=scheduled_at and scheduled_at.astimezone(timezone.utc),
-            state="scheduled" if scheduled_at else "available",
-            tags=insert_opts.tags or args_insert_opts.tags,
-        )
-
-        return insert_params, unique_opts
-
-    @staticmethod
-    def __truncate_time(time, interval_seconds) -> datetime:
-        return datetime.fromtimestamp(
-            (time.timestamp() // interval_seconds) * interval_seconds, tz=timezone.utc
-        )
-
-    @staticmethod
-    def __uint64_to_int64(uint64):
-        # Packs a uint64 then unpacks to int64 to fit within Postgres bigint
-        return (uint64 + (1 << 63)) % (1 << 64) - (1 << 63)
+# Reexport for more ergonomic use in calling code.
+from .client import (
+    Args as Args,
+    Client as Client,
+    InsertOpts as InsertOpts,
+    UniqueOpts as UniqueOpts,
+)

src/riverqueue/client.py

@@ -0,0 +1,189 @@
+from dataclasses import dataclass
+from datetime import datetime, timezone, timedelta
+from typing import Any, Optional, Protocol, Tuple, List, Callable
+
+from .driver import GetParams, JobInsertParams, DriverProtocol
+from .model import InsertResult
+from .fnv import fnv1_hash
+
+MAX_ATTEMPTS_DEFAULT = 25
+PRIORITY_DEFAULT = 1
+QUEUE_DEFAULT = "default"
+
+DEFAULT_UNIQUE_STATES = ["available", "completed", "running", "retryable", "scheduled"]
+
+
+class Args(Protocol):
+    kind: str
+
+    def to_json(self) -> str:
+        pass
+
+
+@dataclass
+class InsertManyParams:
+    args: Args
+    insert_opts: Optional["InsertOpts"] = None
+
+
+@dataclass
+class InsertOpts:
+    scheduled_at: Optional[datetime] = None
+    unique_opts: Optional["UniqueOpts"] = None
+    max_attempts: Optional[int] = None
+    priority: Optional[int] = None
+    queue: Optional[str] = None
+    tags: Optional[List[Any]] = None
+
+
+@dataclass
+class UniqueOpts:
+    by_args: Optional[Any] = None
+    by_period: Optional[Any] = None
+    by_queue: Optional[Any] = None
+    by_state: Optional[Any] = None
+
+
+class Client:
+    def __init__(self, driver: DriverProtocol, advisory_lock_prefix=None):
+        self.driver = driver
+        self.advisory_lock_prefix = advisory_lock_prefix
+        self.time_now_utc = lambda: datetime.now(timezone.utc)  # for test time stubbing
+
+    def insert(
+        self, args: Args, insert_opts: Optional[InsertOpts] = None
+    ) -> InsertResult:
+        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 self.__check_unique_job(insert_params, unique_opts, insert)
+
+    def insert_many(self, args: List[Args]) -> List[InsertResult]:
+        all_params = [
+            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,
+        insert_params: JobInsertParams,
+        unique_opts: Optional[UniqueOpts],
+        insert_func: Callable[[], InsertResult],
+    ) -> InsertResult:
+        if unique_opts is None:
+            return insert_func()
+
+        any_unique_opts = False
+        get_params = GetParams(kind=insert_params.kind)
+
+        lock_str = f"unique_keykind={insert_params.kind}"
+
+        if unique_opts.by_args:
+            any_unique_opts = True
+            get_params.by_args = True
+            get_params.args = insert_params.args
+            lock_str += f"&args={insert_params.args}"
+
+        if unique_opts.by_period:
+            lower_period_bound = self.__truncate_time(
+                self.time_now_utc(), unique_opts.by_period
+            )
+
+            any_unique_opts = True
+            get_params.by_created_at = True
+            get_params.created_at = [
+                lower_period_bound,
+                lower_period_bound + timedelta(seconds=unique_opts.by_period),
+            ]
+            lock_str += f"&period={lower_period_bound.strftime('%FT%TZ')}"
+
+        if unique_opts.by_queue:
+            any_unique_opts = True
+            get_params.by_queue = True
+            get_params.queue = insert_params.queue
+            lock_str += f"&queue={insert_params.queue}"
+
+        if unique_opts.by_state:
+            any_unique_opts = True
+            get_params.by_state = True
+            get_params.state = unique_opts.by_state
+            lock_str += f"&state={','.join(unique_opts.by_state)}"
+        else:
+            get_params.state = DEFAULT_UNIQUE_STATES
+            lock_str += f"&state={','.join(DEFAULT_UNIQUE_STATES)}"
+
+        if not any_unique_opts:
+            return insert_func()
+
+        with self.driver.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)
+
+            existing_job = self.driver.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
+    ) -> Tuple[JobInsertParams, Optional[UniqueOpts]]:
+        if not hasattr(args, "kind"):
+            raise Exception("args should respond to `kind`")
+
+        args_json = args.to_json()
+        if args_json is None:
+            raise Exception("args should return non-nil from `to_json`")
+
+        args_insert_opts = getattr(args, "insert_opts", InsertOpts())
+
+        scheduled_at = insert_opts.scheduled_at or args_insert_opts.scheduled_at
+        unique_opts = insert_opts.unique_opts or args_insert_opts.unique_opts
+
+        if is_insert_many and unique_opts:
+            raise ValueError("unique opts can't be used with `insert_many`")
+
+        insert_params = JobInsertParams(
+            args=args_json,
+            kind=args.kind,
+            max_attempts=insert_opts.max_attempts
+            or args_insert_opts.max_attempts
+            or MAX_ATTEMPTS_DEFAULT,
+            priority=insert_opts.priority
+            or args_insert_opts.priority
+            or PRIORITY_DEFAULT,
+            queue=insert_opts.queue or args_insert_opts.queue or QUEUE_DEFAULT,
+            scheduled_at=scheduled_at and scheduled_at.astimezone(timezone.utc),
+            state="scheduled" if scheduled_at else "available",
+            tags=insert_opts.tags or args_insert_opts.tags,
+        )
+
+        return insert_params, unique_opts
+
+    @staticmethod
+    def __truncate_time(time, interval_seconds) -> datetime:
+        return datetime.fromtimestamp(
+            (time.timestamp() // interval_seconds) * interval_seconds, tz=timezone.utc
+        )
+
+    @staticmethod
+    def __uint64_to_int64(uint64):
+        # Packs a uint64 then unpacks to int64 to fit within Postgres bigint
+        return (uint64 + (1 << 63)) % (1 << 64) - (1 << 63)