CORE-3249: Add migration type for modifying columns

src/Database/PostgreSQL/PQTypes/Checks.hs

@@ -792,6 +792,19 @@ checkDBConsistency options domains tablesWithVersions migrations = do
           runSQL_ "COMMIT"
           runQuery_ (sqlCreateIndexConcurrently tname idx) `finally` begin
           updateTableVersion
+
+        ModifyColumnMigration cursorSql updateSql batchSize -> do
+          logMigration
+          withCursorSQL "migration_cursor" NoScroll Hold cursorSql $ \cursor -> do
+            fix $ \loop -> do
+              cursorFetch_ cursor (CD_Forward batchSize)
+              primaryKeys <- fetchMany runIdentity
+              unless (null primaryKeys) $ do
+                updateSql primaryKeys
+                commit
+                loop
+          updateTableVersion
+
       where
         logMigration = do
           logInfo_ $ arrListTable mgrTableName

src/Database/PostgreSQL/PQTypes/Model/Migration.hs

@@ -31,6 +31,8 @@ module Database.PostgreSQL.PQTypes.Model.Migration (
 
 import Data.Int
 
+import Database.PostgreSQL.PQTypes.FromSQL (FromSQL)
+import Database.PostgreSQL.PQTypes.SQL (SQL)
 import Database.PostgreSQL.PQTypes.Model.Index
 import Database.PostgreSQL.PQTypes.Model.Table
 import Database.PostgreSQL.PQTypes.SQL.Raw
@@ -57,6 +59,17 @@ data MigrationAction m =
       TableIndex
 #endif
 
+  -- | Migration for modifying columns. Parameters are:
+  --
+  -- SQL that will be used for the cursor.
+  --
+  -- Function that takes a list of primary keys provided by the cursor SQL and
+  -- runs an arbitrary computation within MonadDB. The function might be called
+  -- repeatedly depending on the number of primary keys. See the last argument.
+  --
+  -- Number of primary keys fetched at once by the cursor SQL.
+  | forall t . FromSQL t => ModifyColumnMigration SQL ([t] -> m ()) Int
+
 -- | Migration object.
 data Migration m =
   Migration {
@@ -74,13 +87,11 @@ data Migration m =
 isStandardMigration :: Migration m -> Bool
 isStandardMigration Migration{..} =
   case mgrAction of
-    StandardMigration{}                -> True
-    DropTableMigration{}               -> False
-    CreateIndexConcurrentlyMigration{} -> False
+    StandardMigration{} -> True
+    _                   -> False
 
 isDropTableMigration :: Migration m -> Bool
 isDropTableMigration Migration{..} =
   case mgrAction of
-    StandardMigration{}                -> False
-    DropTableMigration{}               -> True
-    CreateIndexConcurrentlyMigration{} -> False
+    DropTableMigration{} -> True
+    _                    -> False

test/Main.hs

@@ -2,10 +2,12 @@ module Main
   where
 
 import Control.Exception.Lifted as E
+import Control.Monad (forM_)
 import Control.Monad.IO.Class
 import Control.Monad.Trans.Control
 
 import Data.Either
+import Data.List (zip4)
 import Data.Monoid
 import Prelude
 import qualified Data.Text as T
@@ -908,6 +910,137 @@ migrationTest4 connSource =
 
   freshTestDB         step
 
+migrationTest5 :: ConnectionSourceM (LogT IO) -> TestTree
+migrationTest5 connSource =
+  testCaseSteps' "Migration test 5" connSource $ \step -> do
+    freshTestDB step
+
+    step "Creating the database (schema version 1)..."
+    migrateDatabase defaultExtrasOptions ["pgcrypto"] [] [] [table1] [createTableMigration table1]
+    checkDatabase defaultExtrasOptions [] [] [table1]
+
+    step "Populating the 'bank' table..."
+    runQuery_ . sqlInsert "bank" $ do
+      sqlSetList "name" $ (\i -> "bank" <> show i) <$> numbers
+      sqlSetList "location" $ (\i -> "location" <> show i) <$> numbers
+
+    forM_ (zip4 tables migrations steps assertions) $
+      \(table, migration, step', assertion) -> do
+        step step'
+        migrateDatabase defaultExtrasOptions ["pgcrypto"] [] [] [table] [migration]
+        checkDatabase defaultExtrasOptions [] [] [table]
+        uncurry assertNoException assertion
+
+    freshTestDB step
+
+  where
+    -- Chosen by a fair dice roll.
+    numbers = [1..21] :: [Int]
+    table1 = tableBankSchema1
+    tables = [ table1 { tblVersion = 2
+                      , tblColumns = tblColumns table1 ++ [stringColumn]
+                      }
+             , table1 { tblVersion = 3
+                      , tblColumns = tblColumns table1 ++ [stringColumn]
+                      }
+             , table1 { tblVersion = 4
+                      , tblColumns = tblColumns table1 ++ [stringColumn, boolColumn]
+                      }
+             , table1 { tblVersion = 5
+                      , tblColumns = tblColumns table1 ++ [stringColumn, boolColumn]
+                      }
+             ]
+
+    migrations = [ addStringColumnMigration
+                 , copyStringColumnMigration
+                 , addBoolColumnMigration
+                 , modifyBoolColumnMigration
+                 ]
+
+    steps = [ "Adding string column (version 1 -> version 2)..."
+            , "Copying string column (version 2 -> version 3)..."
+            , "Adding bool column (version 3 -> version 4)..."
+            , "Modifying bool column (version 4 -> version 5)..."
+            ]
+
+    assertions =
+      [ ("Check that the string column has been added" :: String, checkAddStringColumn)
+      , ("Check that the string data has been copied", checkCopyStringColumn)
+      , ("Check that the bool column has been added", checkAddBoolColumn)
+      , ("Check that the bool column has been modified", checkModifyBoolColumn)
+      ]
+
+    stringColumn = tblColumn { colName = "name_new"
+                             , colType = TextT
+                             }
+
+    boolColumn = tblColumn { colName = "name_is_true"
+                            , colType = BoolT
+                            , colNullable = False
+                            , colDefault = Just "false"
+                            }
+
+    cursorSql = "SELECT id FROM bank" :: SQL
+
+    addStringColumnMigration = Migration
+      { mgrTableName = "bank"
+      , mgrFrom = 1
+      , mgrAction = StandardMigration $
+        runQuery_ $ sqlAlterTable "bank" [ sqlAddColumn stringColumn ]
+      }
+
+    copyStringColumnMigration = Migration
+      { mgrTableName = "bank"
+      , mgrFrom = 2
+      , mgrAction = ModifyColumnMigration cursorSql copyColumnSql 5
+      }
+    copyColumnSql :: MonadDB m => [UUID] -> m ()
+    copyColumnSql primaryKeys =
+      runQuery_ . sqlUpdate "bank" $ do
+        sqlSetCmd "name_new" "bank.name"
+        sqlWhereIn "bank.id" primaryKeys
+
+    addBoolColumnMigration = Migration
+      { mgrTableName = "bank"
+      , mgrFrom = 3
+      , mgrAction = StandardMigration $
+        runQuery_ $ sqlAlterTable "bank" [ sqlAddColumn boolColumn ]
+      }
+
+    modifyBoolColumnMigration = Migration
+      { mgrTableName = "bank"
+      , mgrFrom = 4
+      , mgrAction = ModifyColumnMigration cursorSql modifyColumnSql 100
+      }
+    modifyColumnSql :: MonadDB m => [UUID] -> m ()
+    modifyColumnSql primaryKeys =
+      runQuery_ . sqlUpdate "bank" $ do
+        sqlSet "name_is_true" True
+        sqlWhereIn "bank.id" primaryKeys
+
+    checkAddStringColumn = do
+      runQuery_ . sqlSelect "bank" $ sqlResult "name_new"
+      rows :: [Maybe T.Text] <- fetchMany runIdentity
+      liftIO . assertEqual "No name_new in empty column" True $ all (== Nothing) rows
+
+    checkCopyStringColumn = do
+      runQuery_ . sqlSelect "bank" $ sqlResult "name_new"
+      rows_new :: [Maybe T.Text] <- fetchMany runIdentity
+      runQuery_ . sqlSelect "bank" $ sqlResult "name_new"
+      rows_old :: [Maybe T.Text] <- fetchMany runIdentity
+      liftIO . assertEqual "All name_new are equal name" True $
+        all (uncurry (==)) $ zip rows_new rows_old
+
+    checkAddBoolColumn = do
+      runQuery_ . sqlSelect "bank" $ sqlResult "name_is_true"
+      rows :: [Maybe Bool] <- fetchMany runIdentity
+      liftIO . assertEqual "All name_is_true default to false" True $ all (== Just False) rows
+
+    checkModifyBoolColumn = do
+      runQuery_ . sqlSelect "bank" $ sqlResult "name_is_true"
+      rows :: [Maybe Bool] <- fetchMany runIdentity
+      liftIO . assertEqual "All name_is_true are true" True $ all (== Just True) rows
+
 eitherExc :: MonadBaseControl IO m =>
              (SomeException -> m ()) -> (a -> m ()) -> m a -> m ()
 eitherExc left right c = (E.try c) >>= either left right
@@ -946,6 +1079,7 @@ main = do
                          , migrationTest2 connSource
                          , migrationTest3 connSource
                          , migrationTest4 connSource
+                         , migrationTest5 connSource
                          ]
   where
     ings =