SQLAlchemy async session scoping for background workers

<p>I have a background worker (not FastAPI) that processes traces in a loop. Each batch of traces needs its own database session. I need to avoid session reuse across batches, handle connection pool exhaustion, and ensure sessions are properly closed even if an exception occurs.</p>

<p>Use <code>async_sessionmaker</code> as a context manager per unit of work:</p> <div class="highlight"><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">sqlalchemy.ext.asyncio</span><span class="w"> </span><span class="kn">import</span> <span class="n">create_async_engine</span><span class="p">,</span> <span class="n">async_sessionmaker</span><span class="p">,</span> <span class="n">AsyncSession</span> <span class="n">engine</span> <span class="o">=</span> <span class="n">create_async_engine</span><span class="p">(</span> <span class="n">settings</span><span class="o">.</span><span class="n">database_url</span><span class="p">,</span> <span class="n">pool_size</span><span class="o">=</span><span class="mi">5</span><span class="p">,</span> <span class="n">max_overflow</span><span class="o">=</span><span class="mi">10</span><span class="p">,</span> <span class="n">pool_pre_ping</span><span class="o">=</span><span class="kc">True</span><span class="p">,</span> <span class="c1"># Test connection before use</span> <span class="p">)</span> <span class="n">async_session</span> <span class="o">=</span> <span class="n">async_sessionmaker</span><span class="p">(</span><span class="n">engine</span><span class="p">,</span> <span class="n">expire_on_commit</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">process_batch</span><span class="p">(</span><span class="n">trace_ids</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">str</span><span class="p">])</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span> <span class="c1"># New session per batch — automatically closed on exit</span> <span class="k">async</span> <span class="k">with</span> <span class="n">async_session</span><span class="p">()</span> <span class="k">as</span> <span class="n">session</span><span class="p">:</span> <span class="k">async</span> <span class="k">with</span> <span class="n">session</span><span class="o">.</span><span class="n">begin</span><span class="p">():</span> <span class="c1"># Auto-commit on success, rollback on exception</span> <span class="k">for</span> <span class="n">trace_id</span> <span class="ow">in</span> <span class="n">trace_ids</span><span class="p">:</span> <span class="n">trace</span> <span class="o">=</span> <span class="k">await</span> <span class="n">session</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="n">Trace</span><span class="p">,</span> <span class="n">trace_id</span><span class="p">)</span> <span class="k">if</span> <span class="n">trace</span><span class="p">:</span> <span class="n">trace</span><span class="o">.</span><span class="n">embedding</span> <span class="o">=</span> <span class="k">await</span> <span class="n">generate_embedding</span><span class="p">(</span><span class="n">trace</span><span class="p">)</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">run_worker</span><span class="p">():</span> <span class="k">while</span> <span class="kc">True</span><span class="p">:</span> <span class="n">pending</span> <span class="o">=</span> <span class="k">await</span> <span class="n">fetch_pending_traces</span><span class="p">()</span> <span class="k">if</span> <span class="n">pending</span><span class="p">:</span> <span class="k">await</span> <span class="n">process_batch</span><span class="p">(</span><span class="n">pending</span><span class="p">)</span> <span class="k">else</span><span class="p">:</span> <span class="k">await</span> <span class="n">asyncio</span><span class="o">.</span><span class="n">sleep</span><span class="p">(</span><span class="mi">5</span><span class="p">)</span> <span class="c1"># Graceful shutdown:</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">shutdown</span><span class="p">():</span> <span class="k">await</span> <span class="n">engine</span><span class="o">.</span><span class="n">dispose</span><span class="p">()</span> <span class="c1"># Close all connections in pool</span> </code></pre></div> <p>Key points: - Use <code>async with session.begin()</code> for automatic commit/rollback - <code>pool_pre_ping=True</code> re-establishes dropped connections automatically - One session per batch/transaction — never share across concurrent tasks - <code>engine.dispose()</code> cleanly closes all connections on shutdown - <code>max_overflow</code> controls burst capacity beyond <code>pool_size</code></p>