<commit_before>{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveDataTypeable #-}

-- | A postgresql backend for persistent.
module Database.Persist.Postgresql
    ( withPostgresqlPool
    , withPostgresqlConn
    , createPostgresqlPool
    , module Database.Persist.Sql
    , ConnectionString
    , PostgresConf (..)
    , openSimpleConn
    ) where

import Database.Persist.Sql
import Data.Fixed (Pico)

import qualified Database.PostgreSQL.Simple as PG
import qualified Database.PostgreSQL.Simple.BuiltinTypes as PG
import qualified Database.PostgreSQL.Simple.Internal as PG
import qualified Database.PostgreSQL.Simple.ToField as PGTF
import qualified Database.PostgreSQL.Simple.FromField as PGFF
import qualified Database.PostgreSQL.Simple.Types as PG
import Database.PostgreSQL.Simple.Ok (Ok (..))

import qualified Database.PostgreSQL.LibPQ as LibPQ

import Control.Exception (throw)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Typeable
import Data.IORef
import qualified Data.Map as Map
import Data.Either (partitionEithers)
import Control.Arrow
import Data.List (find, intercalate, sort, groupBy, nub)
import Data.Function (on)
import Data.Conduit
import qualified Data.Conduit.List as CL
import Control.Monad.Logger (MonadLogger, runNoLoggingT)

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B8
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Blaze.ByteString.Builder.Char8 as BBB
import qualified Blaze.ByteString.Builder.ByteString as BBS

import Data.Time.LocalTime (localTimeToUTC, utc)
import Data.Text (Text, pack)
import Data.Aeson
import Control.Monad (forM, mzero)
import Control.Monad.Trans.Resource (Resource, mkResource, with)
import System.Environment (getEnvironment)
import Data.Int (Int64)
import Data.Maybe (mapMaybe, fromJust, isJust, fromMaybe)
import Data.Monoid ((<>))
-- | A @libpq@ connection string.  A simple example of connection
-- string would be @\"host=localhost port=5432 user=test
-- dbname=test password=<PASSWORD>\"@.  Please read libpq's
-- documentation at
-- <http://www.postgresql.org/docs/9.1/static/libpq-connect.html>
-- for more details on how to create such strings.
type ConnectionString = ByteString


-- | Create a PostgreSQL connection pool and run the given
-- action.  The pool is properly released after the action
-- finishes using it.  Note that you should not use the given
-- 'ConnectionPool' outside the action since it may be already
-- been released.
withPostgresqlPool :: (MonadBaseControl IO m, MonadLogger m, MonadIO m)
                   => ConnectionString
                   -- ^ Connection string to the database.
                   -> Int
                   -- ^ Number of connections to be kept open in
                   -- the pool.
                   -> (ConnectionPool -> m a)
                   -- ^ Action to be executed that uses the
                   -- connection pool.
                   -> m a
withPostgresqlPool ci = withSqlPool $ open' ci


-- | Create a PostgreSQL connection pool.  Note that it's your
-- responsability to properly close the connection pool when
-- unneeded.  Use 'withPostgresqlPool' for an automatic resource
-- control.
createPostgresqlPool :: (MonadIO m, MonadBaseControl IO m, MonadLogger m)
                     => ConnectionString
                     -- ^ Connection string to the database.
                     -> Int
                     -- ^ Number of connections to be kept open
                     -- in the pool.
                     -> m ConnectionPool
createPostgresqlPool ci = createSqlPool $ open' ci


-- | Same as 'withPostgresqlPool', but instead of opening a pool
-- of connections, only one connection is opened.
withPostgresqlConn :: (MonadIO m, MonadBaseControl IO m, MonadLogger m)
                   => ConnectionString -> (Connection -> m a) -> m a
withPostgresqlConn = withSqlConn . open'

open' :: ConnectionString -> LogFunc -> IO Connection
open' cstr logFunc = do
    PG.connectPostgreSQL cstr >>= openSimpleConn

-- | Generate a 'Connection' from a 'PG.Connection'
openSimpleConn :: PG.Connection -> IO Connection
openSimpleConn conn = do
    smap <- newIORef $ Map.empty
    return Connection
        { connPrepare    = prepare' conn
        , connStmtMap    = smap
        , connInsertSql  = insertSql'
        , connClose      = PG.close conn
        , connMigrateSql = migrate'
        , connBegin      = const $ PG.begin    conn
        , connCommit     = const $ PG.commit   conn
        , connRollback   = const $ PG.rollback conn
        , connEscapeName = escape
        , connNoLimit    = "LIMIT ALL"
        , connRDBMS      = "postgresql"
        , connLimitOffset = decorateSQLWithLimitOffset "LIMIT ALL"
        }

prepare' :: PG.Connection -> Text -> IO Statement
prepare' conn sql = do
    let query = PG.Query (T.encodeUtf8 sql)
    return Statement
        { stmtFinalize = return ()
        , stmtReset = return ()
        , stmtExecute = execute' conn query
        , stmtQuery = withStmt' conn query
        }
        
insertSql' :: EntityDef SqlType -> [PersistValue] -> InsertSqlResult
insertSql' ent vals =
  let sql = T.concat
                [ "INSERT INTO "
                , escape $ entityDB ent<commit_msg>Set connLogFunc for postgres<commit_after>{-# LANGUAGE OverloadedStrings #-}
{-# LANGUAGE TypeFamilies #-}
{-# LANGUAGE CPP #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE FlexibleInstances #-}
{-# LANGUAGE DeriveDataTypeable #-}

-- | A postgresql backend for persistent.
module Database.Persist.Postgresql
    ( withPostgresqlPool
    , withPostgresqlConn
    , createPostgresqlPool
    , module Database.Persist.Sql
    , ConnectionString
    , PostgresConf (..)
    , openSimpleConn
    ) where

import Database.Persist.Sql
import Data.Fixed (Pico)

import qualified Database.PostgreSQL.Simple as PG
import qualified Database.PostgreSQL.Simple.BuiltinTypes as PG
import qualified Database.PostgreSQL.Simple.Internal as PG
import qualified Database.PostgreSQL.Simple.ToField as PGTF
import qualified Database.PostgreSQL.Simple.FromField as PGFF
import qualified Database.PostgreSQL.Simple.Types as PG
import Database.PostgreSQL.Simple.Ok (Ok (..))

import qualified Database.PostgreSQL.LibPQ as LibPQ

import Control.Exception (throw)
import Control.Monad.IO.Class (MonadIO (..))
import Data.Typeable
import Data.IORef
import qualified Data.Map as Map
import Data.Either (partitionEithers)
import Control.Arrow
import Data.List (find, intercalate, sort, groupBy, nub)
import Data.Function (on)
import Data.Conduit
import qualified Data.Conduit.List as CL
import Control.Monad.Logger (MonadLogger, runNoLoggingT)

import Data.ByteString (ByteString)
import qualified Data.ByteString.Char8 as B8
import qualified Data.Text as T
import qualified Data.Text.Encoding as T
import qualified Blaze.ByteString.Builder.Char8 as BBB
import qualified Blaze.ByteString.Builder.ByteString as BBS

import Data.Time.LocalTime (localTimeToUTC, utc)
import Data.Text (Text, pack)
import Data.Aeson
import Control.Monad (forM, mzero)
import Control.Monad.Trans.Resource (Resource, mkResource, with)
import System.Environment (getEnvironment)
import Data.Int (Int64)
import Data.Maybe (mapMaybe, fromJust, isJust, fromMaybe)
import Data.Monoid ((<>))
-- | A @libpq@ connection string.  A simple example of connection
-- string would be @\"host=localhost port=5432 user=test
-- dbname=test password=<PASSWORD>\"@.  Please read libpq's
-- documentation at
-- <http://www.postgresql.org/docs/9.1/static/libpq-connect.html>
-- for more details on how to create such strings.
type ConnectionString = ByteString


-- | Create a PostgreSQL connection pool and run the given
-- action.  The pool is properly released after the action
-- finishes using it.  Note that you should not use the given
-- 'ConnectionPool' outside the action since it may be already
-- been released.
withPostgresqlPool :: (MonadBaseControl IO m, MonadLogger m, MonadIO m)
                   => ConnectionString
                   -- ^ Connection string to the database.
                   -> Int
                   -- ^ Number of connections to be kept open in
                   -- the pool.
                   -> (ConnectionPool -> m a)
                   -- ^ Action to be executed that uses the
                   -- connection pool.
                   -> m a
withPostgresqlPool ci = withSqlPool $ open' ci


-- | Create a PostgreSQL connection pool.  Note that it's your
-- responsability to properly close the connection pool when
-- unneeded.  Use 'withPostgresqlPool' for an automatic resource
-- control.
createPostgresqlPool :: (MonadIO m, MonadBaseControl IO m, MonadLogger m)
                     => ConnectionString
                     -- ^ Connection string to the database.
                     -> Int
                     -- ^ Number of connections to be kept open
                     -- in the pool.
                     -> m ConnectionPool
createPostgresqlPool ci = createSqlPool $ open' ci


-- | Same as 'withPostgresqlPool', but instead of opening a pool
-- of connections, only one connection is opened.
withPostgresqlConn :: (MonadIO m, MonadBaseControl IO m, MonadLogger m)
                   => ConnectionString -> (Connection -> m a) -> m a
withPostgresqlConn = withSqlConn . open'

open' :: ConnectionString -> LogFunc -> IO Connection
open' cstr logFunc = do
    PG.connectPostgreSQL cstr >>= openSimpleConn logFunc

-- | Generate a 'Connection' from a 'PG.Connection'
openSimpleConn :: LogFunc -> PG.Connection -> IO Connection
openSimpleConn logFunc conn = do
    smap <- newIORef $ Map.empty
    return Connection
        { connPrepare    = prepare' conn
        , connStmtMap    = smap
        , connInsertSql  = insertSql'
        , connClose      = PG.close conn
        , connMigrateSql = migrate'
        , connBegin      = const $ PG.begin    conn
        , connCommit     = const $ PG.commit   conn
        , connRollback   = const $ PG.rollback conn
        , connEscapeName = escape
        , connNoLimit    = "LIMIT ALL"
        , connRDBMS      = "postgresql"
        , connLimitOffset = decorateSQLWithLimitOffset "LIMIT ALL"
        , connLogFunc = logFunc
        }

prepare' :: PG.Connection -> Text -> IO Statement
prepare' conn sql = do
    let query = PG.Query (T.encodeUtf8 sql)
    return Statement
        { stmtFinalize = return ()
        , stmtReset = return ()
        , stmtExecute = execute' conn query
        , stmtQuery = withStmt' conn query
        }
        
insertSql' :: EntityDef SqlType -> [PersistValue] -> InsertSqlResult
insertSql' ent vals =
  let sql = T.concat
                [ "INSERT INTO "
                , escape $ entityDB ent
{'fraction_non_alphanumeric': 0.08734107241569929, 'fraction_numerical': 0.002211166390270868, 'mean_word_length': 2.716193084560082, 'pattern_counts': {'":': 0, '<': 11, '<?xml version=': 0, '>': 52, 'https://': 0, 'lorem ipsum': 0, 'www.': 2, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5102517', 'n_tokens_mistral': 3085, 'n_tokens_neox': 3010, 'n_words': 991}