Ledger.hs

module Database.PostgreSQL.Migrate.History.Codd.Ledger
  ( readCoddHistory,
    withLockedCoddHistory,
    SourceUnlockIssue,
    coddUnlockFailure,
    sourceUnlockCleanupIssue,
    readCoddHistoryOnConnection,
    classifyCoddSchema,
    validateCoddRows,
  )
where

import Control.Exception qualified as Exception
import Data.Bifunctor (first)
import Data.List qualified as List
import Data.Map.Strict qualified as Map
import Data.Set qualified as Set
import Data.Text qualified as Text
import Database.PostgreSQL.Migrate (CleanupIssue (..))
import Database.PostgreSQL.Migrate.History.Codd.Types
import Database.PostgreSQL.Migrate.Internal (useConnectionProvider)
import Hasql.Connection qualified as Connection
import Hasql.Decoders qualified as Decoders
import Hasql.Encoders qualified as Encoders
import Hasql.Errors qualified as Errors
import Hasql.Session qualified as Session
import Hasql.Statement (Statement)
import Hasql.Statement qualified as Statement
import PgMigrate.History.Codd.Prelude

-- | Read and verify selected legacy rows without mutating source objects.
readCoddHistory :: CoddSourceConfig -> IO (Either CoddImportError CoddHistory)
readCoddHistory config@CoddSourceConfig {sourceLockKey} = do
  (result, unlockIssues) <-
    withLockedCoddHistory config (readCoddHistoryOnConnection config)
  pure $ case (result, unlockIssues) of
    (Left primary, issue : _) -> Left (coddUnlockFailure sourceLockKey (Just primary) issue)
    (Left primary, []) -> Left primary
    (Right _, issue : _) -> Left (coddUnlockFailure sourceLockKey Nothing issue)
    (Right history, []) -> Right history

data SourceUnlockIssue
  = SourceUnlockReturnedFalse
  | SourceUnlockSessionFailed !Errors.SessionError

withLockedCoddHistory ::
  CoddSourceConfig ->
  (Connection.Connection -> IO (Either CoddImportError value)) ->
  IO (Either CoddImportError value, [SourceUnlockIssue])
withLockedCoddHistory CoddSourceConfig {sourceProvider, sourceLockKey} action = do
  acquired <-
    useConnectionProvider sourceProvider $ \connection ->
      Exception.mask $ \restore -> do
        locked <- runSession connection (Session.statement sourceLockKey tryLockStatement)
        case locked of
          Left err -> pure (Left err, [])
          Right False -> pure (Left (CoddLockUnavailable sourceLockKey), [])
          Right True -> do
            captured <- Exception.try @Exception.SomeException (restore (action connection))
            unlockIssues <- unlockSource connection sourceLockKey
            case captured of
              Left exception -> Exception.throwIO exception
              Right primary -> pure (primary, unlockIssues)
  pure $ case acquired of
    Left connectionError -> (Left (CoddConnectionFailed connectionError), [])
    Right result -> result

unlockSource :: Connection.Connection -> Int64 -> IO [SourceUnlockIssue]
unlockSource connection lockKey = do
  unlocked <- Connection.use connection (Session.statement lockKey unlockStatement)
  pure $ case unlocked of
    Left sessionError -> [SourceUnlockSessionFailed sessionError]
    Right False -> [SourceUnlockReturnedFalse]
    Right True -> []

coddUnlockFailure :: Int64 -> Maybe CoddImportError -> SourceUnlockIssue -> CoddImportError
coddUnlockFailure lockKey primary unlockIssue =
  case unlockIssue of
    SourceUnlockReturnedFalse -> CoddUnlockFailed lockKey primary Nothing
    SourceUnlockSessionFailed sessionError ->
      CoddUnlockFailed lockKey primary (Just (CoddSessionFailed sessionError))

sourceUnlockCleanupIssue :: SourceUnlockIssue -> CleanupIssue
sourceUnlockCleanupIssue unlockIssue =
  case unlockIssue of
    SourceUnlockReturnedFalse -> AdvisoryUnlockReturnedFalse
    SourceUnlockSessionFailed sessionError -> AdvisoryUnlockFailed sessionError

readCoddHistoryOnConnection ::
  CoddSourceConfig ->
  Connection.Connection ->
  IO (Either CoddImportError CoddHistory)
readCoddHistoryOnConnection config connection = do
  detected <- runSession connection (Session.statement () detectColumnsStatement)
  case detected >>= classifyCoddSchema of
    Left err -> pure (Left err)
    Right (version, schemaName) -> do
      loaded <- runSession connection (Session.statement () (loadRowsStatement version schemaName))
      pure (loaded >>= validateCoddRows config version)

classifyCoddSchema :: [(Text, Text)] -> Either CoddImportError (CoddSchemaVersion, Text)
classifyCoddSchema qualifiedColumns =
  case (Map.lookup "codd_schema" grouped, Map.lookup "codd" grouped) of
    (Nothing, Nothing) -> Left CoddLedgerMissing
    (Just _, Just _) -> Left CoddBothSchemasPresent
    (Just columns, Nothing) -> classifyLegacy columns
    (Nothing, Just columns)
      | columns == v4Columns -> Right (CoddV5, "codd")
      | otherwise -> Left (CoddUnsupportedShape "codd" columns)
  where
    grouped = Map.fromListWith (flip (<>)) [(schema, [column]) | (schema, column) <- qualifiedColumns]
    classifyLegacy columns
      | columns == v1Columns = Right (CoddV1, "codd_schema")
      | columns == v2Columns = Right (CoddV2, "codd_schema")
      | columns == v3Columns = Right (CoddV3, "codd_schema")
      | columns == v4Columns = Right (CoddV4, "codd_schema")
      | otherwise = Left (CoddUnsupportedShape "codd_schema" columns)

validateCoddRows ::
  CoddSourceConfig ->
  CoddSchemaVersion ->
  [CoddHistoryRow] ->
  Either CoddImportError CoddHistory
validateCoddRows CoddSourceConfig {selectedFilenames, strictSource} version rows = do
  case firstDuplicate (filename <$> rows) of
    Just duplicate -> Left (CoddDuplicateLedgerFilename duplicate)
    Nothing -> pure ()
  case List.find (\row -> isJust (noTransactionFailedAt row) || appliedAt row == Nothing) rows of
    Just partial -> Left (CoddPartialMigration (filename partial))
    Nothing -> pure ()
  selected <- traverse selectRow selectedFilenames
  let selectedSet = Set.fromList (toList selectedFilenames)
      unselectedRows = filter ((`Set.notMember` selectedSet) . filename) rows
  if strictSource && not (null unselectedRows)
    then Left (CoddStrictSourceHasUnselected (filename <$> unselectedRows))
    else Right CoddHistory {schemaVersion = version, selectedRows = selected, unselectedRows}
  where
    rowsByFilename = Map.fromList [(filename row, row) | row <- rows]
    selectRow selected =
      maybe (Left (CoddSelectedFilenameMissing selected)) Right (Map.lookup selected rowsByFilename)

firstDuplicate :: (Ord value) => [value] -> Maybe value
firstDuplicate = go Set.empty
  where
    go _ [] = Nothing
    go seen (value : rest)
      | Set.member value seen = Just value
      | otherwise = go (Set.insert value seen) rest

detectColumnsStatement :: Statement () [(Text, Text)]
detectColumnsStatement =
  Statement.preparable
    "SELECT n.nspname::text, a.attname::text FROM pg_catalog.pg_attribute a JOIN pg_catalog.pg_class c ON a.attrelid = c.oid JOIN pg_catalog.pg_namespace n ON c.relnamespace = n.oid WHERE c.relname = 'sql_migrations' AND n.nspname IN ('codd_schema', 'codd') AND a.attnum >= 1 AND NOT a.attisdropped ORDER BY n.nspname, a.attnum"
    Encoders.noParams
    (Decoders.rowList ((,) <$> required Decoders.text <*> required Decoders.text))
  where
    required = Decoders.column . Decoders.nonNullable

loadRowsStatement :: CoddSchemaVersion -> Text -> Statement () [CoddHistoryRow]
loadRowsStatement version schemaName =
  Statement.unpreparable
    ( Text.unwords
        [ "SELECT name, migration_timestamp, applied_at,",
          if version >= CoddV3 then "COALESCE(num_applied_statements, 0)::int4," else "0::int4,",
          if version >= CoddV3 then "no_txn_failed_at" else "NULL::timestamptz",
          "FROM " <> schemaName <> ".sql_migrations ORDER BY id"
        ]
    )
    Encoders.noParams
    (Decoders.rowList rowDecoder)
  where
    rowDecoder =
      CoddHistoryRow
        <$> (Text.unpack <$> required Decoders.text)
        <*> required Decoders.timestamptz
        <*> nullableColumn Decoders.timestamptz
        <*> required Decoders.int4
        <*> nullableColumn Decoders.timestamptz
    required = Decoders.column . Decoders.nonNullable
    nullableColumn = Decoders.column . Decoders.nullable

tryLockStatement :: Statement Int64 Bool
tryLockStatement =
  Statement.preparable
    "SELECT pg_try_advisory_lock($1)"
    (Encoders.param (Encoders.nonNullable Encoders.int8))
    (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.bool)))

unlockStatement :: Statement Int64 Bool
unlockStatement =
  Statement.preparable
    "SELECT pg_advisory_unlock($1)"
    (Encoders.param (Encoders.nonNullable Encoders.int8))
    (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.bool)))

runSession :: Connection.Connection -> Session.Session value -> IO (Either CoddImportError value)
runSession connection session =
  first CoddSessionFailed <$> Connection.use connection session

v1Columns :: [Text]
v1Columns = ["id", "migration_timestamp", "applied_at", "name"]

v2Columns :: [Text]
v2Columns = v1Columns <> ["application_duration"]

v3Columns :: [Text]
v3Columns = v2Columns <> ["num_applied_statements", "no_txn_failed_at"]

v4Columns :: [Text]
v4Columns = v3Columns <> ["txnid", "connid"]