fix (#1): Use NullPool with SQLAlchemy in Celery tasks

Need to improve this for the future, but at least we don't run into
asyncio issues due to shared connections improperly closed and handled
across celery tasks.

server/polar/postgres.py

@@ -1,16 +1,29 @@
+from polar.config import settings
+from polar.ext.sqlalchemy import sql
 from sqlalchemy.ext.asyncio import AsyncSession, create_async_engine
 from sqlalchemy.orm import sessionmaker
+from sqlalchemy.pool import NullPool
 
-from polar.config import settings
-from polar.ext.sqlalchemy import sql
 
-engine = create_async_engine(settings.postgres_dsn, echo=True)
-AsyncSessionLocal = sessionmaker(
-    engine,
-    autocommit=False,
-    autoflush=False,
-    expire_on_commit=False,
-    class_=AsyncSession,
-)
+def create_sessionmaker(is_celery: bool = False) -> sessionmaker:
+    engine_options = dict(echo=settings.DEBUG)
+    if is_celery:
+        # TODO: Change pooling strategy for celery workers.
+        # In the meantime, we're using NullPool to avoid
+        # issues with asyncio in Celery.
+        engine_options.update(dict(poolclass=NullPool))
+
+    engine = create_async_engine(settings.postgres_dsn, **engine_options)
+    return sessionmaker(
+        engine,
+        autocommit=False,
+        autoflush=False,
+        expire_on_commit=False,
+        class_=AsyncSession,
+    )
+
+
+AsyncSessionLocal = create_sessionmaker()
+
 
-__all__ = ["sql", "engine", "AsyncSessionLocal", "AsyncSession"]
+__all__ = ["sql", "AsyncSessionLocal", "AsyncSession"]

server/polar/tasks/github/webhook.py

@@ -1,9 +1,6 @@
 from typing import Any
 
 import structlog
-
-from polar import actions
-from polar.clients import github
 from polar.models import Issue, Organization, PullRequest, Repository
 from polar.platforms import Platforms
 from polar.postgres import AsyncSession
@@ -13,6 +10,9 @@
 from polar.schema.repository import CreateRepository
 from polar.worker import asyncify_task, task
 
+from polar import actions
+from polar.clients import github
+
 log = structlog.get_logger()
 
 

server/polar/worker.py

@@ -3,9 +3,8 @@
 
 from asgiref.sync import async_to_sync
 from celery import Celery
-
 from polar.config import settings
-from polar.postgres import AsyncSessionLocal
+from polar.postgres import create_sessionmaker
 
 app = Celery(
     "polar", backend=settings.CELERY_BACKEND_URL, broker=settings.CELERY_BROKER_URL
@@ -20,6 +19,9 @@
 ReturnValue = TypeVar("ReturnValue")
 
 
+AsyncSessionLocal = create_sessionmaker(is_celery=True)
+
+
 def asyncify_task(
     with_session: bool = False,
 ) -> Callable[[Callable[Params, ReturnValue]], Callable[Params, ReturnValue]]: