From: Mike Bayer <mike_mp@zzzcomputing.com>
Date: Sat, 7 Aug 2021 15:48:19 +0000 (-0400)
Subject: qualify sqlite batch add column for dynamic defaults
X-Git-Tag: rel_1_7_0~16^2
X-Git-Url: http://git.ipfire.org/cgi-bin/gitweb.cgi?a=commitdiff_plain;h=3110acba13c22c7d8a934c30dca650b1537a50e3;p=thirdparty%2Fsqlalchemy%2Falembic.git

qualify sqlite batch add column for dynamic defaults

Batch "auto" mode will now select for "recreate" if the ``add_column()``
operation is used on SQLite, and the column itself meets the criteria for
SQLite where ADD COLUMN is not allowed, in this case a functional or
parenthesized SQL expression or a ``Computed`` (i.e. generated) column.

Change-Id: Ie948b4b8ad8dc698b458403831e47bac4ad45b8a
Fixes: #883
---

diff --git a/alembic/ddl/sqlite.py b/alembic/ddl/sqlite.py
index 8b78966a..cb790ea7 100644
--- a/alembic/ddl/sqlite.py
+++ b/alembic/ddl/sqlite.py
@@ -2,6 +2,8 @@ import re
 
 from sqlalchemy import cast
 from sqlalchemy import JSON
+from sqlalchemy import schema
+from sqlalchemy import sql
 
 from .impl import DefaultImpl
 from .. import util
@@ -25,7 +27,19 @@ class SQLiteImpl(DefaultImpl):
 
         """
         for op in batch_op.batch:
-            if op[0] not in ("add_column", "create_index", "drop_index"):
+            if op[0] == "add_column":
+                col = op[1][1]
+                if isinstance(
+                    col.server_default, schema.DefaultClause
+                ) and isinstance(col.server_default.arg, sql.ClauseElement):
+                    return True
+                elif (
+                    isinstance(col.server_default, util.sqla_compat.Computed)
+                    and col.server_default.persisted
+                ):
+                    return True
+                return False
+            elif op[0] not in ("create_index", "drop_index"):
                 return True
         else:
             return False
diff --git a/docs/build/unreleased/883.rst b/docs/build/unreleased/883.rst
new file mode 100644
index 00000000..ecf38bc8
--- /dev/null
+++ b/docs/build/unreleased/883.rst
@@ -0,0 +1,8 @@
+.. change::
+    :tags: bug, sqlite, batch
+    :tickets: 883
+
+    Batch "auto" mode will now select for "recreate" if the ``add_column()``
+    operation is used on SQLite, and the column itself meets the criteria for
+    SQLite where ADD COLUMN is not allowed, in this case a functional or
+    parenthesized SQL expression or a ``Computed`` (i.e. generated) column.
diff --git a/tests/test_batch.py b/tests/test_batch.py
index 9e0491d7..91baef75 100644
--- a/tests/test_batch.py
+++ b/tests/test_batch.py
@@ -26,6 +26,7 @@ from sqlalchemy.schema import CreateTable
 from sqlalchemy.sql import column
 from sqlalchemy.sql import text
 
+from alembic import testing
 from alembic.ddl import sqlite
 from alembic.operations import Operations
 from alembic.operations.batch import ApplyBatchImpl
@@ -40,8 +41,16 @@ from alembic.testing import TestBase
 from alembic.testing.fixtures import op_fixture
 from alembic.util import exc as alembic_exc
 from alembic.util.sqla_compat import _select
+from alembic.util.sqla_compat import has_computed
+from alembic.util.sqla_compat import has_identity
 from alembic.util.sqla_compat import sqla_14
 
+if has_computed:
+    from alembic.util.sqla_compat import Computed
+
+if has_identity:
+    from alembic.util.sqla_compat import Identity
+
 
 class BatchApplyTest(TestBase):
     def setUp(self):
@@ -1931,6 +1940,77 @@ class BatchRoundTripTest(TestBase):
             ["id", "data", "x", "data2"],
         )
 
+    def test_add_column_auto_server_default_calculated(self):
+        """test #883"""
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.add_column(
+                Column(
+                    "data2",
+                    DateTime(),
+                    server_default=self._datetime_server_default_fixture(),
+                )
+            )
+
+        self._assert_data(
+            [
+                {"id": 1, "data": "d1", "x": 5, "data2": mock.ANY},
+                {"id": 2, "data": "22", "x": 6, "data2": mock.ANY},
+                {"id": 3, "data": "8.5", "x": 7, "data2": mock.ANY},
+                {"id": 4, "data": "9.46", "x": 8, "data2": mock.ANY},
+                {"id": 5, "data": "d5", "x": 9, "data2": mock.ANY},
+            ]
+        )
+        eq_(
+            [col["name"] for col in inspect(config.db).get_columns("foo")],
+            ["id", "data", "x", "data2"],
+        )
+
+    @testing.combinations((True,), (False,))
+    @testing.exclusions.only_on("sqlite")
+    def test_add_column_auto_generated(self, persisted):
+        """test #883"""
+        with self.op.batch_alter_table("foo") as batch_op:
+            batch_op.add_column(
+                Column(
+                    "data2", Integer, Computed("1 + 1", persisted=persisted)
+                )
+            )
+
+        self._assert_data(
+            [
+                {"id": 1, "data": "d1", "x": 5, "data2": 2},
+                {"id": 2, "data": "22", "x": 6, "data2": 2},
+                {"id": 3, "data": "8.5", "x": 7, "data2": 2},
+                {"id": 4, "data": "9.46", "x": 8, "data2": 2},
+                {"id": 5, "data": "d5", "x": 9, "data2": 2},
+            ]
+        )
+        eq_(
+            [col["name"] for col in inspect(config.db).get_columns("foo")],
+            ["id", "data", "x", "data2"],
+        )
+
+    @config.requirements.identity_columns
+    def test_add_column_auto_identity(self):
+        """test #883"""
+
+        self._no_pk_fixture()
+
+        with self.op.batch_alter_table("nopk") as batch_op:
+            batch_op.add_column(Column("id", Integer, Identity()))
+
+        self._assert_data(
+            [
+                {"a": 1, "b": 2, "c": 3, "id": 1},
+                {"a": 2, "b": 4, "c": 5, "id": 2},
+            ],
+            tablename="nopk",
+        )
+        eq_(
+            [col["name"] for col in inspect(config.db).get_columns("foo")],
+            ["id", "data", "x"],
+        )
+
     def test_add_column_insert_before_recreate(self):
         with self.op.batch_alter_table("foo", recreate="always") as batch_op:
             batch_op.add_column(