From 39791b4552b605fe60d03a5e98c5533bd741b729 Mon Sep 17 00:00:00 2001
From: Ben Sima <ben@bensima.com>
Date: Thu, 27 Nov 2025 15:21:18 -0500
Subject: Fix SQL query and add simple migration test

---
 Omni/Task.hs               |  4 +++-
 Omni/Task/Core.hs          |  3 ++-
 Omni/Task/MigrationTest.hs | 42 ++++++++++++++++++++++++++++++++++++++++++
 3 files changed, 47 insertions(+), 2 deletions(-)
 create mode 100644 Omni/Task/MigrationTest.hs

diff --git a/Omni/Task.hs b/Omni/Task.hs
index 8fb49be..b0dbe89 100644
--- a/Omni/Task.hs
+++ b/Omni/Task.hs
@@ -13,6 +13,7 @@ import qualified Data.Text as T
 import qualified Omni.Cli as Cli
 import qualified Omni.Namespace as Namespace
 import Omni.Task.Core
+import qualified Omni.Task.MigrationTest as MigrationTest
 import qualified Omni.Task.RaceTest as RaceTest
 import qualified Omni.Test as Test
 import qualified System.Console.Docopt as Docopt
@@ -415,7 +416,8 @@ test =
     "Omni.Task"
     [ unitTests,
       cliTests,
-      Tasty.after Tasty.AllSucceed "Unit tests" RaceTest.test
+      Tasty.after Tasty.AllSucceed "Unit tests" RaceTest.test,
+      Tasty.after Tasty.AllSucceed "Unit tests" MigrationTest.test
     ]
 
 unitTests :: Test.Tree
diff --git a/Omni/Task/Core.hs b/Omni/Task/Core.hs
index 4eac9b5..63bc533 100644
--- a/Omni/Task/Core.hs
+++ b/Omni/Task/Core.hs
@@ -437,7 +437,8 @@ migrateTable conn tableName expectedCols = do
 -- | Get list of column names for a table using PRAGMA table_info
 getTableColumns :: SQL.Connection -> Text -> IO [Text]
 getTableColumns conn tableName = do
-  rows <- SQL.query conn "PRAGMA table_info(?)" (SQL.Only tableName) :: IO [(Int, Text, Text, Int, Maybe Text, Int)]
+  let query = SQL.Query <| "PRAGMA table_info(" <> tableName <> ")"
+  rows <- SQL.query_ conn query :: IO [(Int, Text, Text, Int, Maybe Text, Int)]
   pure [colName | (_, colName, _, _, _, _) <- rows]
 
 -- | Add a column to a table
diff --git a/Omni/Task/MigrationTest.hs b/Omni/Task/MigrationTest.hs
new file mode 100644
index 0000000..f16f782
--- /dev/null
+++ b/Omni/Task/MigrationTest.hs
@@ -0,0 +1,42 @@
+{-# LANGUAGE OverloadedStrings #-}
+{-# LANGUAGE NoImplicitPrelude #-}
+
+module Omni.Task.MigrationTest where
+
+import Alpha
+import qualified Data.Set as Set
+import Omni.Task.Core
+import qualified Omni.Test as Test
+import System.Directory (doesFileExist, removeFile)
+import System.Environment (setEnv)
+
+test :: Test.Tree
+test = Test.group "Omni.Task.Migration" [migrationStartupTest]
+
+migrationStartupTest :: Test.Tree
+migrationStartupTest =
+  Test.unit "database initializes with schema migrations" <| do
+    setEnv "TASK_TEST_MODE" "1"
+
+    let testFile = "_/tmp/tasks-test.db"
+    exists <- doesFileExist testFile
+    when exists <| removeFile testFile
+
+    initTaskDb
+
+    withDb <| \conn -> do
+      tasksCols <- getTableColumns conn "tasks"
+      activityCols <- getTableColumns conn "task_activity"
+      retryCols <- getTableColumns conn "retry_context"
+
+      Set.fromList ["id", "title", "status"]
+        `Set.isSubsetOf` Set.fromList tasksCols
+        Test.@?= True
+      Set.fromList ["id", "task_id", "stage"]
+        `Set.isSubsetOf` Set.fromList activityCols
+        Test.@?= True
+      Set.fromList ["task_id", "attempt", "reason"]
+        `Set.isSubsetOf` Set.fromList retryCols
+        Test.@?= True
+
+    removeFile testFile
-- 
cgit v1.2.3