content: {-# LANGUAGE NamedFieldPuns #-}
  {-# LANGUAGE OverloadedStrings #-}
  {-# LANGUAGE QuasiQuotes #-}
  {-# LANGUAGE ScopedTypeVariables #-}
  
  module Simplex.Messaging.Agent.Store.Postgres.Util
    ( createDBAndUserIfNotExists,
      dropSchema,
      dropAllSchemasExceptSystem,
      dropDatabaseAndUser,
    )
  where
  
  import Control.Exception (bracket)
  import Control.Monad (forM_, unless, void, when)
  import Data.String (fromString)
  import Database.PostgreSQL.Simple (ConnectInfo (..), Only (..), defaultConnectInfo)
  import qualified Database.PostgreSQL.Simple as PSQL
  import Database.PostgreSQL.Simple.SqlQQ (sql)
  
  createDBAndUserIfNotExists :: ConnectInfo -> IO ()
  createDBAndUserIfNotExists ConnectInfo {connectUser = user, connectDatabase = dbName} = do
    -- connect to the default "postgres" maintenance database
    bracket (PSQL.connect defaultConnectInfo {connectUser = "postgres", connectDatabase = "postgres"}) PSQL.close $
      \postgresDB -> do
        void $ PSQL.execute_ postgresDB "SET client_min_messages TO WARNING"
        -- check if the user exists, create if not
        [Only userExists] <-
          PSQL.query
            postgresDB
            [sql|
              SELECT EXISTS (
                SELECT 1 FROM pg_catalog.pg_roles
                WHERE rolname = ?
              )
            |]
            (Only user)
        unless userExists $ void $ PSQL.execute_ postgresDB (fromString $ "CREATE USER " <> user)
        -- check if the database exists, create if not
        dbExists <- checkDBExists postgresDB dbName
        unless dbExists $ void $ PSQL.execute_ postgresDB (fromString $ "CREATE DATABASE " <> dbName <> " OWNER " <> user)
  
  checkDBExists :: PSQL.Connection -> String -> IO Bool
  checkDBExists postgresDB dbName = do
    [Only dbExists] <-
      PSQL.query
        postgresDB
        [sql|
          SELECT EXISTS (
            SELECT 1 FROM pg_catalog.pg_database
            WHERE datname = ?
          )
        |]
        (Only dbName)
    pure dbExists
  
  dropSchema :: ConnectInfo -> String -> IO ()
  dropSchema connectInfo schema =
    bracket (PSQL.connect connectInfo) PSQL.close $
      \db -> do
        void $ PSQL.execute_ db "SET client_min_messages TO WARNING"
        void $ PSQL.execute_ db (fromString $ "DROP SCHEMA IF EXISTS " <> schema <> " CASCADE")
  
  dropAllSchemasExceptSystem :: ConnectInfo -> IO ()
  dropAllSchemasExceptSystem connectInfo =
    bracket (PSQL.connect connectInfo) PSQL.close $
      \db -> do
        void $ PSQL.execute_ db "SET client_min_messages TO WARNING"
        schemaNames :: [Only String] <-
          PSQL.query_
            db
            [sql|
              SELECT schema_name
              FROM information_schema.schemata
              WHERE schema_name NOT IN ('public', 'pg_catalog', 'information_schema')
            |]
        forM_ schemaNames $ \(Only schema) ->
          PSQL.execute_ db (fromString $ "DROP SCHEMA " <> schema <> " CASCADE")
  
  dropDatabaseAndUser :: ConnectInfo -> IO ()
  dropDatabaseAndUser ConnectInfo {connectUser = user, connectDatabase = dbName} =
    bracket (PSQL.connect defaultConnectInfo {connectUser = "postgres", connectDatabase = "postgres"}) PSQL.close $
      \postgresDB -> do
        void $ PSQL.execute_ postgresDB "SET client_min_messages TO WARNING"
        dbExists <- checkDBExists postgresDB dbName
        when dbExists $ do
          void $ PSQL.execute_ postgresDB (fromString $ "ALTER DATABASE " <> dbName <> " WITH ALLOW_CONNECTIONS false")
          -- terminate all connections to the database
          _r :: [Only Bool] <-
            PSQL.query
              postgresDB
              [sql|
                SELECT pg_terminate_backend(pg_stat_activity.pid)
                FROM pg_stat_activity
                WHERE datname = ?
                  AND pid <> pg_backend_pid()
              |]
              (Only dbName)
          void $ PSQL.execute_ postgresDB (fromString $ "DROP DATABASE " <> dbName)
        void $ PSQL.execute_ postgresDB (fromString $ "DROP USER IF EXISTS " <> user)
  