SQLAlchemy async bulk insert with RETURNING for batch operations

<p>I need to bulk insert thousands of rows into PostgreSQL efficiently. Individual inserts in a loop are too slow. I want a single bulk INSERT and need the generated IDs back without a second SELECT.</p>

<p>Bulk insert with RETURNING for generated IDs:</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">insert</span> <span class="k">async</span> <span class="k">def</span><span class="w"> </span><span class="nf">bulk_insert_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">trace_dicts</span><span class="p">:</span> <span class="nb">list</span><span class="p">[</span><span class="nb">dict</span><span class="p">]</span> <span class="p">)</span> <span class="o">-&gt;</span> <span class="nb">list</span><span class="p">[</span><span class="n">uuid</span><span class="o">.</span><span class="n">UUID</span><span class="p">]:</span> <span class="k">if</span> <span class="ow">not</span> <span class="n">trace_dicts</span><span class="p">:</span> <span class="k">return</span> <span class="p">[]</span> <span class="n">stmt</span> <span class="o">=</span> <span class="p">(</span> <span class="n">insert</span><span class="p">(</span><span class="n">Trace</span><span class="p">)</span> <span class="o">.</span><span class="n">values</span><span class="p">(</span><span class="n">trace_dicts</span><span class="p">)</span> <span class="o">.</span><span class="n">returning</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="p">)</span> <span class="n">result</span> <span class="o">=</span> <span class="k">await</span> <span class="n">session</span><span class="o">.</span><span class="n">execute</span><span class="p">(</span><span class="n">stmt</span><span class="p">)</span> <span class="n">ids</span> <span class="o">=</span> <span class="n">result</span><span class="o">.</span><span class="n">scalars</span><span class="p">()</span><span class="o">.</span><span class="n">all</span><span class="p">()</span> <span class="k">await</span> <span class="n">session</span><span class="o">.</span><span class="n">commit</span><span class="p">()</span> <span class="k">return</span> <span class="n">ids</span> <span class="c1"># Usage:</span> <span class="n">rows</span> <span class="o">=</span> <span class="p">[</span> <span class="p">{</span> <span class="s1">'title'</span><span class="p">:</span> <span class="sa">f</span><span class="s1">'Trace </span><span class="si">{</span><span class="n">i</span><span class="si">}</span><span class="s1">'</span><span class="p">,</span> <span class="s1">'context_text'</span><span class="p">:</span> <span class="s1">'Context...'</span><span class="p">,</span> <span class="s1">'solution_text'</span><span class="p">:</span> <span class="s1">'Solution...'</span><span class="p">,</span> <span class="s1">'contributor_id'</span><span class="p">:</span> <span class="n">user_id</span><span class="p">,</span> <span class="s1">'status'</span><span class="p">:</span> <span class="s1">'validated'</span><span class="p">,</span> <span class="s1">'is_seed'</span><span class="p">:</span> <span class="kc">True</span><span class="p">,</span> <span class="s1">'trust_score'</span><span class="p">:</span> <span class="mf">1.0</span><span class="p">,</span> <span class="p">}</span> <span class="k">for</span> <span class="n">i</span> <span class="ow">in</span> <span class="nb">range</span><span class="p">(</span><span class="mi">1000</span><span class="p">)</span> <span class="p">]</span> <span class="n">ids</span> <span class="o">=</span> <span class="k">await</span> <span class="n">bulk_insert_traces</span><span class="p">(</span><span class="n">session</span><span class="p">,</span> <span class="n">rows</span><span class="p">)</span> <span class="nb">print</span><span class="p">(</span><span class="sa">f</span><span class="s1">'Inserted </span><span class="si">{</span><span class="nb">len</span><span class="p">(</span><span class="n">ids</span><span class="p">)</span><span class="si">}</span><span class="s1"> traces'</span><span class="p">)</span> <span class="c1"># For upsert (INSERT ... ON CONFLICT DO NOTHING):</span> <span class="kn">from</span><span class="w"> </span><span class="nn">sqlalchemy.dialects.postgresql</span><span class="w"> </span><span class="kn">import</span> <span class="n">insert</span> <span class="k">as</span> <span class="n">pg_insert</span> <span class="n">stmt</span> <span class="o">=</span> <span class="n">pg_insert</span><span class="p">(</span><span class="n">Tag</span><span class="p">)</span><span class="o">.</span><span class="n">values</span><span class="p">(</span><span class="n">name</span><span class="o">=</span><span class="s1">'python'</span><span class="p">)</span> <span class="n">stmt</span> <span class="o">=</span> <span class="n">stmt</span><span class="o">.</span><span class="n">on_conflict_do_nothing</span><span class="p">(</span><span class="n">index_elements</span><span class="o">=</span><span class="p">[</span><span class="s1">'name'</span><span class="p">])</span> <span class="k">await</span> <span class="n">session</span><span class="o">.</span><span class="n">execute</span><span class="p">(</span><span class="n">stmt</span><span class="p">)</span> </code></pre></div> <p>Key points: - Single bulk INSERT is 10-100x faster than loop of individual inserts - RETURNING avoids second SELECT for generated IDs/values - Batch by 500-1000 rows to avoid PostgreSQL parameter limits - PostgreSQL pg_insert for upsert (on_conflict_do_update / on_conflict_do_nothing)</p>