Internal.hs

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

module Database.Sqlite.Easy.Internal where

import Database.SQLite3
import Data.String (IsString, fromString)
import Data.Text (Text)
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.
createSqlitePool :: ConnectionString -> IO (Pool Database)
createSqlitePool (ConnectionString connStr) =
  newPool PoolConfig
    { createResource = open connStr
    , freeResource = close
    , poolCacheTTL = 180
    , poolMaxResources = 50
    }

-- | Open a database, run some stuff, close the database.
withDb :: ConnectionString -> SQLite a -> IO a
withDb (ConnectionString connStr) =
  bracket (open connStr) 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 (prepare db stmt) 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 (prepare db stmt) finalize $ \preparedStmt -> do
      bind preparedStmt params
      fetchAll preparedStmt

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