Internal.hs

{-# language RecordWildCards #-}
{-# language DerivingVia #-}
{-# language OverloadedStrings #-}
{-# language ScopedTypeVariables #-}

-- | The implementation of sqlite-easy.
--
-- This module is unstable and may change at any time.
module Database.Sqlite.Easy.Internal where

import Database.SQLite3 (Database, SQLData, Statement)
import qualified Database.SQLite3 as Direct
import Data.String (IsString, fromString)
import Data.Text (Text)
import qualified Data.Text as T
import Data.Typeable
import Data.Pool
import Control.Monad.Reader
import Control.Exception
import Control.Monad.IO.Unlift

-- * Connection

-- | A SQLite3 connection string
newtype ConnectionString
  = ConnectionString
    { unConnectionString :: Text
    }
  deriving IsString via Text
  deriving Show

-- | Create a pool of a sqlite3 db with a specific connection string.
--   This also sets a few default pragmas.
createSqlitePool :: ConnectionString -> IO (Pool Database)
createSqlitePool connStr =
  newPool $ defaultPoolConfig
    (openWith connStr defaultPragmas)
    Direct.close
    180
    50

-- | Create a pool of a sqlite3 db with a specific connection string.
--   This will also run the supplied SQL statements after establishing
--   each connection.
createSqlitePoolWith :: ConnectionString -> [SQL] -> IO (Pool Database)
createSqlitePoolWith connStr statements =
  newPool $ defaultPoolConfig
    (openWith connStr statements)
    Direct.close
    180
    50

-- | Open a connection to a database, run the supplied statements,
--   and return the connection.
openWith :: ConnectionString -> [SQL] -> IO Database
openWith (ConnectionString connStr) sqls = do
  conn <- Direct.open connStr
  Direct.exec
    conn
    (T.unlines $ map (\(SQL sql) -> sql <> "; ") sqls)
  pure conn

-- | Default pragmas to be set when opening a connection from a pool.
defaultPragmas :: [SQL]
defaultPragmas =
  [ -- https://www.sqlite.org/pragma.html#pragma_foreign_keys
    "PRAGMA foreign_keys = on"
  , -- https://www.sqlite.org/pragma.html#pragma_busy_timeout
    "PRAGMA busy_timeout = 30000"
  ]

-- | Open a database, run some stuff, close the database.
withDb :: ConnectionString -> SQLite a -> IO a
withDb (ConnectionString connStr) =
  bracket (Direct.open connStr) Direct.close . flip runSQLite

-- | Use an active database connection to run some stuff on a database.
withDatabase :: Database -> SQLite a -> IO a
withDatabase = runSQLite

-- | Use a resource pool to run some stuff on a database.
withPool :: Pool Database -> SQLite a -> IO a
withPool pool = withResource pool . flip runSQLite

-- * Execution

-- | A SQL statement
newtype SQL
  = SQL
    { unSQL :: Text
    }
  deriving (Semigroup, IsString) via Text
  deriving Show

-- | Run a SQL statement on a database and fetch the results.
run :: SQL -> SQLite [[SQLData]]
run (SQL stmt) = do
  db <- getDB
  liftIO $ bracket (Direct.prepare db stmt) Direct.finalize fetchAll

-- | Run a SQL statement with certain parameters on a database and fetch the results.
runWith :: SQL -> [SQLData] -> SQLite [[SQLData]]
runWith (SQL stmt) params = do
  db <- getDB
  liftIO $ do
    bracket (Direct.prepare db stmt) Direct.finalize $ \preparedStmt -> do
      Direct.bind preparedStmt params
      fetchAll preparedStmt

-- | Run a statement and fetch all of the data.
fetchAll :: Statement -> IO [[SQLData]]
fetchAll stmt = do
  res <- Direct.step stmt
  case res of
    Direct.Row -> do
      row <- Direct.columns stmt
      rows <- fetchAll stmt
      pure (row : rows)
    Direct.Done ->
      pure []

-- * Transaction

-- | The type of actions to run on a SQLite database.
--   In essence, it is almost the same as @Database -> IO a@.
--
--   'SQLite' actions can be created with the 'run' and 'runWith'
--   functions, and can be composed using the type class instances.
--
--   'SQLite' actions can be run with the 'withDb', 'withDatabase',
--   and 'withPool' functions.
newtype SQLite a
  = SQLite
    { unSQLite :: SQLiteStuff -> IO a
    }
  deriving (Functor, Applicative, Monad, MonadIO, MonadFail, MonadUnliftIO)
  via ReaderT SQLiteStuff IO

instance Semigroup a => Semigroup (SQLite a) where
  a <> b = (<>) <$> a <*> b

instance Monoid a => Monoid (SQLite a) where
  mempty = pure mempty

data SQLiteStuff
  = SQLiteStuff
    { dbConn :: Database
    , transactionNumber :: Maybe Int
    }

getDB :: SQLite Database
getDB = SQLite (\(SQLiteStuff dbConn _) -> pure dbConn)

runSQLite :: Database -> SQLite a -> IO a
runSQLite db t = (unSQLite t) (SQLiteStuff db Nothing)

-- | Run operations as a transaction.
--   If the action throws an error, the transaction is rolled back.
--   For more information, visit: <https://www.sqlite.org/lang_transaction.html>
transaction :: forall a. Typeable a => SQLite a -> SQLite a
transaction action = do
  SQLiteStuff {..} <- SQLite $ \stuff -> pure stuff
  case transactionNumber of
    Nothing -> do
      let
        runIO sql = (unSQLite (run sql)) (SQLiteStuff dbConn Nothing)
        commit = runIO "COMMIT"
        rollback' = runIO "ROLLBACK"
      [] <- run "BEGIN"
      liftIO $ catches
        ((unSQLite action) (SQLiteStuff dbConn (Just 1)) <* commit)
        [ Handler $ \(RollbackCurrent a) -> rollback' *> pure a
        , Handler $ \(RollbackAll a) -> rollback' *> pure a
        , Handler $ \(ex :: SomeException) -> rollback' *> throwIO ex
        ]
    Just n -> do
      let
        runIO sql = (unSQLite (run sql)) (SQLiteStuff dbConn Nothing)
        transactionName = "'sqlite_easy_transaction_" <> fromString (show n) <> "'"
        release = runIO $ "RELEASE SAVEPOINT " <> transactionName
        rollbackCurrent = runIO $ "ROLLBACK TRANSACTION TO SAVEPOINT " <> transactionName
      [] <- run $ "SAVEPOINT " <> transactionName
      liftIO $ catches
        ((unSQLite action) (SQLiteStuff dbConn (Just (n + 1))) <* release)
        [ Handler $ \(RollbackCurrent a) -> rollbackCurrent *> pure a
        , Handler $ \(ex :: RollbackAll a) -> rollbackCurrent *> throwIO ex
        , Handler $ \(ex :: SomeException) -> rollbackCurrent *> throwIO ex
        ]

asTransaction' :: Database -> IO a -> IO a
asTransaction' db action = do
  let
    runIO sql = (unSQLite (run sql)) (SQLiteStuff db Nothing)
  [] <- runIO "BEGIN"
  catches
    (action <* runIO "COMMIT")
    [ Handler $ \(ex :: SomeException) -> runIO "ROLLBACK" *> throwIO ex
    ]

-- | Rollback the current (inner-most) transaction by supplying the return value.
--   To be used inside transactions.
rollback :: Typeable a => a -> SQLite a
rollback = liftIO . throwIO . RollbackCurrent

-- | Rollback all transaction structure by supplying the return value.
--   To be used inside transactions.
rollbackAll :: Typeable a => a -> SQLite a
rollbackAll = liftIO . throwIO . RollbackAll

data RollbackCurrent a
  = RollbackCurrent a

instance Show (RollbackCurrent a) where
  show RollbackCurrent{} = "RollbackCurrent"

instance (Typeable a) => Exception (RollbackCurrent a)

data RollbackAll a
  = RollbackAll a

instance Show (RollbackAll a) where
  show RollbackAll{} = "RollbackAll"

instance (Typeable a) => Exception (RollbackAll a)