From eaf3e438df57b0bc90a867ae5f0672f385843072 Mon Sep 17 00:00:00 2001
From: Fraser Murray <fraser.m.murray@gmail.com>
Date: Mon, 4 Dec 2023 11:40:16 +0000
Subject: [PATCH] Added support for extended result codes

The library now exposes sqlite3_extended_result_codes() and
SQLITE_OPEN_EXRESCODE.
---
 Database/SQLite3.hs          |  2 ++
 Database/SQLite3/Bindings.hs |  5 +++++
 Database/SQLite3/Direct.hs   |  6 ++++++
 test/Main.hs                 | 33 +++++++++++++++++++++++++++++++++
 4 files changed, 46 insertions(+)

diff --git a/Database/SQLite3.hs b/Database/SQLite3.hs
index c644229..1926555 100644
--- a/Database/SQLite3.hs
+++ b/Database/SQLite3.hs
@@ -232,6 +232,7 @@ data SQLOpenFlag
     | SQLOpenPrivateCache  -- Ok for sqlite3_open_v2()
     | SQLOpenWAL           -- VFS only
     | SQLOpenNoFollow      -- Ok for sqlite3_open_v2()
+    | SQLOpenExResCode     -- Extended result codes
     deriving (Eq, Show)
 
 -- | These VFS names are used when using the `open2` function.
@@ -371,6 +372,7 @@ open2 path flags zvfs =
         toNum SQLOpenPrivateCache   = 0x00040000
         toNum SQLOpenWAL            = 0x00080000
         toNum SQLOpenNoFollow       = 0x01000000
+        toNum SQLOpenExResCode      = 0x02000000
 
 -- | <https://www.sqlite.org/c3ref/close.html>
 close :: Database -> IO ()
diff --git a/Database/SQLite3/Bindings.hs b/Database/SQLite3/Bindings.hs
index 9f2ce33..9e0a466 100644
--- a/Database/SQLite3/Bindings.hs
+++ b/Database/SQLite3/Bindings.hs
@@ -9,6 +9,7 @@ module Database.SQLite3.Bindings (
     c_sqlite3_errcode,
     c_sqlite3_extended_errcode,
     c_sqlite3_errmsg,
+    c_sqlite3_extended_result_codes,
     c_sqlite3_interrupt,
     c_sqlite3_trace,
     CTraceCallback,
@@ -166,6 +167,10 @@ foreign import ccall unsafe "sqlite3_extended_errcode"
 foreign import ccall unsafe "sqlite3_errmsg"
     c_sqlite3_errmsg :: Ptr CDatabase -> IO CString
 
+-- | <https://www.sqlite.org/c3ref/extended_result_codes.html>
+foreign import ccall unsafe "sqlite3_extended_result_codes"
+    c_sqlite3_extended_result_codes :: Ptr CDatabase -> Bool -> IO CError
+
 -- | <https://www.sqlite.org/c3ref/interrupt.html>
 foreign import ccall "sqlite3_interrupt"
     c_sqlite3_interrupt :: Ptr CDatabase -> IO ()
diff --git a/Database/SQLite3/Direct.hs b/Database/SQLite3/Direct.hs
index 470d851..f42cc8c 100644
--- a/Database/SQLite3/Direct.hs
+++ b/Database/SQLite3/Direct.hs
@@ -16,6 +16,7 @@ module Database.SQLite3.Direct (
     errcode,
     extendedErrcode,
     errmsg,
+    setExtendedResultCodes,
     setTrace,
     getAutoCommit,
     setSharedCacheEnabled,
@@ -314,6 +315,11 @@ extendedErrcode :: Database -> IO Error
 extendedErrcode (Database db) =
     decodeError <$> c_sqlite3_extended_errcode db
 
+-- | <https://www.sqlite.org/c3ref/extended_result_codes.html>
+setExtendedResultCodes :: Database -> Bool -> IO (Either Error ())
+setExtendedResultCodes (Database db) enabled =
+    toResult () <$> c_sqlite3_extended_result_codes db enabled
+
 -- | <https://www.sqlite.org/c3ref/errcode.html>
 errmsg :: Database -> IO Utf8
 errmsg (Database db) =
diff --git a/test/Main.hs b/test/Main.hs
index 855a34e..11b0407 100644
--- a/test/Main.hs
+++ b/test/Main.hs
@@ -55,6 +55,7 @@ regressionTests1 =
     , TestLabel "TypedColumns"  . testTypedColumns
     , TestLabel "ColumnName"    . testColumnName
     , TestLabel "Errors"        . testErrors
+    , TestLabel "ExtendedErrors". testExtendedErrors
     , TestLabel "Integrity"     . testIntegrity
     , TestLabel "DecodeError"   . testDecodeError
     , TestLabel "ResultStats"   . testResultStats
@@ -632,6 +633,38 @@ testErrors TestEnv{..} = TestCase $ do
                 \INSERT INTO foo VALUES (3, 4); \
                 \INSERT INTO foo VALUES (5, 6)"
 
+testExtendedErrors :: forall f . TestEnv f -> Test
+testExtendedErrors TestEnv{..} = TestCase $ do
+  -- opening a connection with extended results mode
+  conn <- open2 ":memory:" [SQLOpenReadWrite, SQLOpenCreate, SQLOpenExResCode] SQLVFSDefault
+  exec conn "CREATE TABLE foo (a INT NOT NULL, b INT);"
+  res <- Direct.exec conn "INSERT INTO foo (a, b) VALUES (NULL, 0);"
+  case res of
+    Left err ->
+      assertEqual "testExtendedErrors: expected an extended error code, but got the wrong error code" err
+        (ErrorConstraintNotNull, "NOT NULL constraint failed: foo.a")
+    Right () ->
+      assertFailure "testExtendedErrors: exec should have return extended error code, but succeeded"
+  close conn
+
+  -- setting a connection to extended results mode after it's opened
+  withConn $ \conn -> do
+    exec conn "CREATE TABLE foo (a INT UNIQUE);"
+    exec conn "INSERT INTO foo (a) VALUES (1);"
+
+    res <- Direct.exec conn "INSERT INTO foo (a) VALUES (1);"
+    assertEqual "testExtendedErrors: expected a primary error code" res
+      (Left (ErrorConstraint, "UNIQUE constraint failed: foo.a"))
+
+    res2 <- Direct.setExtendedResultCodes conn True
+    case res2 of
+      Left err -> assertFailure $
+        "testExtendedErrors: failed to enable extended result codes, got error " <> show err
+      Right () -> do
+        res3 <- Direct.exec conn "INSERT INTO foo (a) VALUES (1);"
+        assertEqual "testExtendedErrors: expected an extended error code" res3
+          (Left (ErrorConstraintUnique, "UNIQUE constraint failed: foo.a"))
+
 -- Make sure data stored in a table comes back as-is.
 testIntegrity :: forall f. TestEnv f -> Test
 testIntegrity TestEnv{..} = TestCase $ do