SchemaCheck.hs

{-# LANGUAGE TemplateHaskell #-}

module Keiro.Migrations.SchemaCheck
  ( SchemaDrift (..),
    compareSchemaSnapshot,
    expectedSchemaSnapshot,
    renderSchemaDrift,
    snapshotSchema,
    verifyExpectedSchema,
  )
where

import Control.Exception (finally)
import Data.Int (Int32)
import Data.Map.Strict (Map)
import Data.Map.Strict qualified as Map
import Data.Maybe (mapMaybe)
import Data.Set qualified as Set
import Data.Text (Text)
import Data.Text qualified as Text
import Database.PostgreSQL.Migrate (MigrationError (..))
import Hasql.Connection qualified as Connection
import Hasql.Connection.Settings qualified as Settings
import Hasql.Decoders qualified as Decoders
import Hasql.Encoders qualified as Encoders
import Hasql.Session (Session)
import Hasql.Session qualified as Session
import Hasql.Statement (Statement)
import Hasql.Statement qualified as Statement
import Keiro.Migrations.Internal.EmbedFile (embedTextFile)

-- | One named difference between the expected and live schema snapshots.
data SchemaDrift
  = MissingObject Text
  | UnexpectedObject Text
  | ChangedObject
      { driftKey :: Text,
        expectedDefinition :: Text,
        actualDefinition :: Text
      }
  deriving stock (Eq, Show)

-- | Compare canonical snapshots by their @kind<TAB>name@ object identity.
compareSchemaSnapshot :: Text -> Text -> [SchemaDrift]
compareSchemaSnapshot expected actual =
  mapMaybe driftFor allKeys
  where
    expectedObjects = snapshotObjects expected
    actualObjects = snapshotObjects actual
    allKeys =
      Set.toAscList
        (Map.keysSet expectedObjects `Set.union` Map.keysSet actualObjects)

    driftFor key =
      case (Map.lookup key expectedObjects, Map.lookup key actualObjects) of
        (Just (expectedLine, _), Nothing) ->
          Just (MissingObject expectedLine)
        (Nothing, Just (actualLine, _)) ->
          Just (UnexpectedObject actualLine)
        (Just (_, expectedValue), Just (_, actualValue))
          | expectedValue /= actualValue ->
              Just
                ChangedObject
                  { driftKey = key,
                    expectedDefinition = expectedValue,
                    actualDefinition = actualValue
                  }
        _ -> Nothing

-- | Render a drift as one operator-facing line with the affected object name.
renderSchemaDrift :: SchemaDrift -> Text
renderSchemaDrift drift =
  case drift of
    MissingObject line ->
      let (kind, name, definition) = splitSnapshotLine line
       in "schema drift: missing "
            <> kind
            <> " "
            <> name
            <> " (expected: "
            <> definition
            <> ")"
    UnexpectedObject line ->
      let (kind, name, definition) = splitSnapshotLine line
       in "schema drift: unexpected "
            <> kind
            <> " "
            <> name
            <> " (actual: "
            <> definition
            <> ")"
    ChangedObject {driftKey, expectedDefinition, actualDefinition} ->
      let (kind, name) = splitSnapshotKey driftKey
       in "schema drift: changed "
            <> kind
            <> " "
            <> name
            <> " (expected: "
            <> expectedDefinition
            <> "; actual: "
            <> actualDefinition
            <> ")"

-- | Read a sorted canonical snapshot of tables, columns, constraints, and indexes.
snapshotSchema :: Text -> Session Text
snapshotSchema schema =
  Text.unlines <$> Session.statement schema schemaSnapshotStatement

-- | PostgreSQL 18 snapshot generated from the complete embedded migration plan.
expectedSchemaSnapshot :: Text
expectedSchemaSnapshot =
  $(embedTextFile "expected-schema/native/keiro-v18.txt")

-- | Compare the live @keiro@ schema with the embedded PostgreSQL 18 snapshot.
verifyExpectedSchema ::
  Settings.Settings ->
  IO (Either MigrationError [SchemaDrift])
verifyExpectedSchema settings = do
  acquired <- Connection.acquire settings
  case acquired of
    Left connectionError ->
      pure (Left (ConnectionAcquisitionFailed connectionError))
    Right connection -> do
      result <-
        Connection.use connection liveSnapshotSession
          `finally` Connection.release connection
      pure $ case result of
        Left sessionError -> Left (DatabaseSessionFailed sessionError)
        Right (Left migrationError) -> Left migrationError
        Right (Right actual) ->
          Right (compareSchemaSnapshot expectedSchemaSnapshot actual)
  where
    liveSnapshotSession :: Session (Either MigrationError Text)
    liveSnapshotSession = do
      serverVersionNumber <- Session.statement () serverVersionStatement
      let majorVersion = fromIntegral serverVersionNumber `div` 10000
      if majorVersion == (18 :: Int)
        then Right <$> snapshotSchema "keiro"
        else pure (Left (UnsupportedPostgresVersion majorVersion))

snapshotObjects :: Text -> Map Text (Text, Text)
snapshotObjects =
  Map.fromList . mapMaybe parseSnapshotLine . Text.lines
  where
    parseSnapshotLine line =
      case Text.splitOn "\t" line of
        kind : name : definitionParts ->
          Just
            ( kind <> "\t" <> name,
              (line, Text.intercalate "\t" definitionParts)
            )
        _ -> Nothing

splitSnapshotLine :: Text -> (Text, Text, Text)
splitSnapshotLine line =
  case Text.splitOn "\t" line of
    kind : name : definitionParts ->
      (kind, name, Text.intercalate "\t" definitionParts)
    _ -> ("object", line, line)

splitSnapshotKey :: Text -> (Text, Text)
splitSnapshotKey key =
  case Text.splitOn "\t" key of
    [kind, name] -> (kind, name)
    _ -> ("object", key)

schemaSnapshotStatement :: Statement Text [Text]
schemaSnapshotStatement =
  Statement.preparable
    """
    WITH configured AS MATERIALIZED (
      SELECT set_config('search_path', 'pg_catalog', true) AS search_path
    )
    SELECT line
    FROM configured
    CROSS JOIN LATERAL (
      SELECT 'table' || E'\t' || c.relname || E'\t' || 'kind=r' AS line
        FROM pg_class c
        JOIN pg_namespace n ON n.oid = c.relnamespace
        WHERE n.nspname = $1
          AND c.relkind = 'r'
          AND configured.search_path = 'pg_catalog'
      UNION ALL
      SELECT 'column' || E'\t' || c.relname || '.' || a.attname || E'\t'
             || format_type(a.atttypid, a.atttypmod)
             || CASE WHEN a.attnotnull THEN ' not null' ELSE '' END
             || coalesce(' default ' || pg_get_expr(d.adbin, d.adrelid), '')
        FROM pg_attribute a
        JOIN pg_class c ON c.oid = a.attrelid
        JOIN pg_namespace n ON n.oid = c.relnamespace
        LEFT JOIN pg_attrdef d
          ON d.adrelid = a.attrelid AND d.adnum = a.attnum
        WHERE n.nspname = $1
          AND c.relkind = 'r'
          AND a.attnum > 0
          AND NOT a.attisdropped
          AND configured.search_path = 'pg_catalog'
      UNION ALL
      SELECT 'constraint' || E'\t' || rel.relname || '.' || con.conname || E'\t'
             || pg_get_constraintdef(con.oid)
        FROM pg_constraint con
        JOIN pg_class rel ON rel.oid = con.conrelid
        JOIN pg_namespace n ON n.oid = rel.relnamespace
        WHERE n.nspname = $1
          AND configured.search_path = 'pg_catalog'
      UNION ALL
      SELECT 'index' || E'\t' || ci.relname || E'\t'
             || pg_get_indexdef(i.indexrelid)
        FROM pg_index i
        JOIN pg_class ci ON ci.oid = i.indexrelid
        JOIN pg_class ct ON ct.oid = i.indrelid
        JOIN pg_namespace n ON n.oid = ct.relnamespace
        WHERE n.nspname = $1
          AND configured.search_path = 'pg_catalog'
    ) snapshot
    ORDER BY line COLLATE "C"
    """
    (Encoders.param (Encoders.nonNullable Encoders.text))
    (Decoders.rowList (Decoders.column (Decoders.nonNullable Decoders.text)))

serverVersionStatement :: Statement () Int32
serverVersionStatement =
  Statement.preparable
    "SELECT current_setting('server_version_num')::integer"
    Encoders.noParams
    (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int4)))