Python async generator for streaming large datasets

<p>I need to stream large query results from PostgreSQL without loading everything into memory. I want to use a Python async generator that yields rows in batches, and I want to support both streaming HTTP responses and background processing.</p>

<p>Use SQLAlchemy's <code>stream_results</code> with async generators:</p> <div class="highlight"><pre><span></span><code><span class="kn">from</span><span class="w"> </span><span class="nn">sqlalchemy</span><span class="w"> </span><span class="kn">import</span> <span class="n">select</span> <span class="kn">from</span><span class="w"> </span><span class="nn">collections.abc</span><span class="w"> </span><span class="kn">import</span> <span class="n">AsyncIterator</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">stream_traces</span><span class="p">(</span> <span class="n">session</span><span class="p">:</span> <span class="n">AsyncSession</span><span class="p">,</span> <span class="n">batch_size</span><span class="p">:</span> <span class="nb">int</span> <span class="o">=</span> <span class="mi">100</span><span class="p">,</span> <span class="p">)</span> <span class="o">-&gt;</span> <span class="n">AsyncIterator</span><span class="p">[</span><span class="nb">list</span><span class="p">[</span><span class="n">Trace</span><span class="p">]]:</span> <span class="w"> </span><span class="sd">"""Yield traces in batches without loading all into memory."""</span> <span class="c1"># stream_results uses server-side cursor (postgresql)</span> <span class="k">async</span> <span class="k">with</span> <span class="n">session</span><span class="o">.</span><span class="n">stream</span><span class="p">(</span> <span class="n">select</span><span class="p">(</span><span class="n">Trace</span><span class="p">)</span><span class="o">.</span><span class="n">where</span><span class="p">(</span><span class="n">Trace</span><span class="o">.</span><span class="n">status</span> <span class="o">==</span> <span class="s1">'validated'</span><span class="p">)</span> <span class="p">)</span> <span class="k">as</span> <span class="n">result</span><span class="p">:</span> <span class="k">async</span> <span class="k">for</span> <span class="n">batch</span> <span class="ow">in</span> <span class="n">result</span><span class="o">.</span><span class="n">partitions</span><span class="p">(</span><span class="n">batch_size</span><span class="p">):</span> <span class="k">yield</span> <span class="p">[</span><span class="n">row</span><span class="p">[</span><span class="mi">0</span><span class="p">]</span> <span class="k">for</span> <span class="n">row</span> <span class="ow">in</span> <span class="n">batch</span><span class="p">]</span> <span class="c1"># Usage in background worker:</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">reindex_all</span><span class="p">():</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">for</span> <span class="n">batch</span> <span class="ow">in</span> <span class="n">stream_traces</span><span class="p">(</span><span class="n">session</span><span class="p">):</span> <span class="k">await</span> <span class="n">process_batch</span><span class="p">(</span><span class="n">batch</span><span class="p">)</span> <span class="c1"># Streaming HTTP response with FastAPI:</span> <span class="kn">from</span><span class="w"> </span><span class="nn">fastapi.responses</span><span class="w"> </span><span class="kn">import</span> <span class="n">StreamingResponse</span> <span class="kn">import</span><span class="w"> </span><span class="nn">json</span> <span class="nd">@router</span><span class="o">.</span><span class="n">get</span><span class="p">(</span><span class="s1">'/traces/export'</span><span class="p">)</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">export_traces</span><span class="p">(</span><span class="n">db</span><span class="p">:</span> <span class="n">DbSession</span><span class="p">):</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">generate</span><span class="p">():</span> <span class="k">async</span> <span class="k">for</span> <span class="n">batch</span> <span class="ow">in</span> <span class="n">stream_traces</span><span class="p">(</span><span class="n">db</span><span class="p">):</span> <span class="k">for</span> <span class="n">trace</span> <span class="ow">in</span> <span class="n">batch</span><span class="p">:</span> <span class="k">yield</span> <span class="n">json</span><span class="o">.</span><span class="n">dumps</span><span class="p">({</span><span class="s1">'id'</span><span class="p">:</span> <span class="nb">str</span><span class="p">(</span><span class="n">trace</span><span class="o">.</span><span class="n">id</span><span class="p">),</span> <span class="s1">'title'</span><span class="p">:</span> <span class="n">trace</span><span class="o">.</span><span class="n">title</span><span class="p">})</span> <span class="o">+</span> <span class="s1">'</span><span class="se">\n</span><span class="s1">'</span> <span class="k">return</span> <span class="n">StreamingResponse</span><span class="p">(</span><span class="n">generate</span><span class="p">(),</span> <span class="n">media_type</span><span class="o">=</span><span class="s1">'application/x-ndjson'</span><span class="p">)</span> </code></pre></div> <p>Key points: - <code>session.stream()</code> uses PostgreSQL server-side cursor — constant memory usage - <code>partitions(n)</code> yields in chunks — don't set this too small (overhead per round trip) - <code>StreamingResponse</code> with async generator streams HTTP response incrementally - Use NDJSON (newline-delimited JSON) for streaming — easier to parse than one big array</p>