Alembic migration with data migration alongside schema change

<p>Need to split a single <code>name</code> column into <code>first_name</code> and <code>last_name</code> columns. This requires a schema change (add columns, remove old) AND a data migration (populate new columns from old). Pure DDL migrations don't handle the data step.</p>

<p>Write an Alembic migration with upgrade() containing both DDL and DML in the correct order:</p> <div class="highlight"><pre><span></span><code><span class="c1"># alembic/versions/0005_split_name_column.py</span> <span class="sd">"""split name into first_name and last_name</span> <span class="sd">Revision ID: 0005</span> <span class="sd">Downs revision: '0004'</span> <span class="sd">Create Date: 2024-01-15</span> <span class="sd">"""</span> <span class="kn">from</span><span class="w"> </span><span class="nn">alembic</span><span class="w"> </span><span class="kn">import</span> <span class="n">op</span> <span class="kn">import</span><span class="w"> </span><span class="nn">sqlalchemy</span><span class="w"> </span><span class="k">as</span><span class="w"> </span><span class="nn">sa</span> <span class="kn">from</span><span class="w"> </span><span class="nn">sqlalchemy.sql</span><span class="w"> </span><span class="kn">import</span> <span class="n">text</span> <span class="k">def</span><span class="w"> </span><span class="nf">upgrade</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span> <span class="c1"># 1. Add new columns as nullable first</span> <span class="n">op</span><span class="o">.</span><span class="n">add_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="n">sa</span><span class="o">.</span><span class="n">Column</span><span class="p">(</span><span class="s1">'first_name'</span><span class="p">,</span> <span class="n">sa</span><span class="o">.</span><span class="n">String</span><span class="p">(</span><span class="mi">100</span><span class="p">),</span> <span class="n">nullable</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span> <span class="n">op</span><span class="o">.</span><span class="n">add_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="n">sa</span><span class="o">.</span><span class="n">Column</span><span class="p">(</span><span class="s1">'last_name'</span><span class="p">,</span> <span class="n">sa</span><span class="o">.</span><span class="n">String</span><span class="p">(</span><span class="mi">100</span><span class="p">),</span> <span class="n">nullable</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span> <span class="c1"># 2. Migrate data (run before making columns NOT NULL)</span> <span class="n">conn</span> <span class="o">=</span> <span class="n">op</span><span class="o">.</span><span class="n">get_bind</span><span class="p">()</span> <span class="n">conn</span><span class="o">.</span><span class="n">execute</span><span class="p">(</span><span class="n">text</span><span class="p">(</span><span class="s2">"""</span> <span class="s2"> UPDATE users</span> <span class="s2"> SET</span> <span class="s2"> first_name = split_part(name, ' ', 1),</span> <span class="s2"> last_name = NULLIF(split_part(name, ' ', 2), '')</span> <span class="s2"> WHERE name IS NOT NULL</span> <span class="s2"> """</span><span class="p">))</span> <span class="c1"># 3. Make columns NOT NULL now that data is populated</span> <span class="n">op</span><span class="o">.</span><span class="n">alter_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="s1">'first_name'</span><span class="p">,</span> <span class="n">nullable</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span> <span class="c1"># 4. Drop old column (last — after data is safe)</span> <span class="n">op</span><span class="o">.</span><span class="n">drop_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="s1">'name'</span><span class="p">)</span> <span class="k">def</span><span class="w"> </span><span class="nf">downgrade</span><span class="p">()</span> <span class="o">-&gt;</span> <span class="kc">None</span><span class="p">:</span> <span class="c1"># Reverse the migration</span> <span class="n">op</span><span class="o">.</span><span class="n">add_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="n">sa</span><span class="o">.</span><span class="n">Column</span><span class="p">(</span><span class="s1">'name'</span><span class="p">,</span> <span class="n">sa</span><span class="o">.</span><span class="n">String</span><span class="p">(</span><span class="mi">200</span><span class="p">),</span> <span class="n">nullable</span><span class="o">=</span><span class="kc">True</span><span class="p">))</span> <span class="n">conn</span> <span class="o">=</span> <span class="n">op</span><span class="o">.</span><span class="n">get_bind</span><span class="p">()</span> <span class="n">conn</span><span class="o">.</span><span class="n">execute</span><span class="p">(</span><span class="n">text</span><span class="p">(</span><span class="s2">"""</span> <span class="s2"> UPDATE users</span> <span class="s2"> SET name = TRIM(COALESCE(first_name, '') || ' ' || COALESCE(last_name, ''))</span> <span class="s2"> """</span><span class="p">))</span> <span class="n">op</span><span class="o">.</span><span class="n">alter_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="s1">'name'</span><span class="p">,</span> <span class="n">nullable</span><span class="o">=</span><span class="kc">False</span><span class="p">)</span> <span class="n">op</span><span class="o">.</span><span class="n">drop_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="s1">'first_name'</span><span class="p">)</span> <span class="n">op</span><span class="o">.</span><span class="n">drop_column</span><span class="p">(</span><span class="s1">'users'</span><span class="p">,</span> <span class="s1">'last_name'</span><span class="p">)</span> </code></pre></div> <p>Key ordering: (1) add nullable columns → (2) migrate data → (3) add NOT NULL constraints → (4) drop old column. Never set NOT NULL before data migration. Always write a working <code>downgrade()</code>. Test on a production-size data clone before running — large tables may need batched updates.</p>