Skip to content

Instantly share code, notes, and snippets.

@KaneTW
Last active November 24, 2023 14:04
Show Gist options
  • Save KaneTW/0f321b6ccf0143611c43e2d087e939d3 to your computer and use it in GitHub Desktop.
Save KaneTW/0f321b6ccf0143611c43e2d087e939d3 to your computer and use it in GitHub Desktop.
--needs a bunch of extensions
module VerifySchema where
import Prelude hiding ( filter )
import qualified Data.List as L
import Rel8 hiding (run)
import qualified Rel8
import Hasql.Connection
import Hasql.Session
import Data.Functor.Contravariant ( (>$<) )
import Data.Int ( Int64 )
import Data.Text ( Text )
import qualified Data.Text as T
import GHC.Generics
import qualified Data.List.NonEmpty as NonEmpty
import qualified Data.Map as M
import Rel8.Schema.Null hiding (nullable)
import qualified Rel8.Schema.Null as Null
import Rel8.Schema.Name ( Name(Name) )
import Rel8.Schema.Spec
import Data.Functor.Const
import Rel8.Schema.HTable
import Control.Monad
data Relkind = RTable
deriving stock (Show)
deriving anyclass (DBEq)
instance DBType Relkind where
typeInformation = parseTypeInformation parser printer typeInformation
where
parser = \case
"r" -> pure RTable
(x :: Text) -> Left $ "Unknown relkind: " ++ show x
printer = \case
RTable -> "r"
newtype Oid = Oid Int64
deriving newtype (DBType, DBEq, Show)
data PGClass f = PGClass
{ oid :: Column f Oid
, relname :: Column f Text
, relkind :: Column f Relkind
, relnamespace :: Column f Oid
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (PGClass Result)
pgclass :: TableSchema (PGClass Name)
pgclass = TableSchema
{ name = QualifiedName "pg_class" (Just "pg_catalog")
, columns = namesFromLabelsWith NonEmpty.last
}
data PGAttribute f = PGAttribute
{ attrelid :: Column f Oid
, attname :: Column f Text
, atttypid :: Column f Oid
, attnum :: Column f Int64
, attnotnull :: Column f Bool
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (PGAttribute Result)
pgattribute :: TableSchema (PGAttribute Name)
pgattribute = TableSchema
{ name = QualifiedName "pg_attribute" (Just "pg_catalog")
, columns = namesFromLabelsWith NonEmpty.last
}
data PGType f = PGType
{ oid :: Column f Oid
, typname :: Column f Text
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (PGType Result)
pgtype :: TableSchema (PGType Name)
pgtype = TableSchema
{ name = QualifiedName "pg_type" (Just "pg_catalog")
, columns = namesFromLabelsWith NonEmpty.last
}
data PGNamespace f = PGNamespace
{ oid :: Column f Oid
, nspname :: Column f Text
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (PGNamespace Result)
pgnamespace :: TableSchema (PGNamespace Name)
pgnamespace = TableSchema
{ name = QualifiedName "pg_namespace" (Just "pg_catalog")
, columns = namesFromLabelsWith NonEmpty.last
}
data PGCast f = PGCast
{ oid :: Column f Oid
, castsource :: Column f Oid
, casttarget :: Column f Oid
, castfunc :: Column f Oid
, castcontext :: Column f Char
, castmethod :: Column f Char
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (PGCast Result)
pgcast :: TableSchema (PGCast Name)
pgcast = TableSchema
{ name = QualifiedName "pg_cast" (Just "pg_catalog")
, columns = namesFromLabelsWith NonEmpty.last
}
data PGTable f = PGTable
{ name :: Column f Text
, columns :: HList f (Attribute f)
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (PGTable Result)
data Attribute f = Attribute
{ attribute :: PGAttribute f
, typ :: PGType f
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (Attribute Result)
data Cast f = Cast
{ source :: PGType f
, target :: PGType f
, context :: Column f Char
}
deriving stock (Generic)
deriving anyclass (Rel8able)
deriving stock instance Show (Cast Result)
fetchTables :: Connection -> IO (Either QueryError [PGTable Result])
fetchTables c = do
flip run c $ statement () $ Rel8.run $ select do
PGClass{ oid = tableOid, relname } <- orderBy (relname >$< asc) do
each pgclass
>>= filter ((lit RTable ==.) . relkind)
columns <- many do
attribute@PGAttribute{ atttypid } <-
each pgattribute
>>= filter ((tableOid ==.) . attrelid)
>>= filter ((>. 0) . attnum)
typ <-
each pgtype
>>= filter (\PGType{ oid = typoid } -> atttypid ==. typoid)
return Attribute{ attribute, typ }
return PGTable
{ name = relname
, ..
}
fetchCasts :: Connection -> IO (Either QueryError [Cast Result])
fetchCasts c = do
flip run c $ statement () $ Rel8.run $ select do
PGCast {castsource, casttarget, castcontext} <- each pgcast
src <- each pgtype >>= filter (\PGType { oid = typoid } -> typoid ==. castsource)
tgt <- each pgtype >>= filter (\PGType { oid = typoid } -> typoid ==. casttarget)
return Cast { source = src, target = tgt, context = castcontext }
data CheckEnv = CheckEnv
{ ctx :: [String]
, schemaMap :: M.Map String [Attribute Result] -- map of schemas to attributes
, casts :: [(String, String)] -- list of implicit casts
} deriving (Show)
nulled :: forall t. Nullable t => Bool
nulled = nullableToBool $ Null.nullable @t
nullableToBool :: Nullity a -> Bool
nullableToBool Null = True
nullableToBool NotNull = False
attrsToMap :: [Attribute Result] -> M.Map String (Attribute Result)
attrsToMap = foldMap (\attr -> M.singleton (T.unpack $ attr.attribute.attname) attr)
data TypeInfo = TypeInfo
{ label :: String
, isNull :: Bool
, typeName :: QualifiedName
} deriving (Show, Eq)
schemaToTypeMap :: forall k. Rel8able k => k Name -> M.Map String TypeInfo
schemaToTypeMap cols = M.fromList . uncurry zip . getConst $
htabulateA @(Columns (k Name)) $ \field ->
case (hfield hspecs field, hfield (toColumns cols) field) of
(Spec {..}, Name name) -> Const ([name], [
TypeInfo { label = head labels
, isNull = nullableToBool nullity
, typeName = info.typeName.name}])
-- implicit casts are ok as long as they're bidirectional
checkTypeEquality :: CheckEnv -> Attribute Result -> TypeInfo -> Either String ()
checkTypeEquality env attr ty
| attrTyName == tyTyName = return ()
| (attrTyName, tyTyName) `elem` env.casts && (tyTyName, attrTyName) `elem` env.casts
= return ()
| otherwise = Left $ show env.ctx ++ ": Cannot convert between db type " ++ attrTyName ++ " and hs type " ++ tyTyName
where
attrTyName = T.unpack attr.typ.typname
tyTyName = ty.typeName.name
checkTypes :: CheckEnv -> M.Map String (Attribute Result) -> M.Map String TypeInfo -> Either String ()
checkTypes env attrMap typeMap = do
forM_ (M.assocs typeMap) $ \(key, ty) -> case M.lookup key attrMap of
Just attr -> checkTypeEquality env {ctx = env.ctx ++ [key]} attr ty
Nothing -> Left $ show env.ctx ++ ": Entry " ++ key ++ " not present in db"
forM_ (M.keys $ M.filter (\attr -> attr.attribute.attnotnull) attrMap) $
\key -> case M.lookup key typeMap of
Just _ -> return ()
Nothing -> Left $ show env.ctx ++ ": Entry " ++ key ++ " not null but not present in hs ty"
-- a schema is valid if:
-- 1. for every existing field, the types match
-- 2. all non-nullable columns are present in the hs type
-- 3. no nonexistent columns are present in the hs type
verifySchema :: Rel8able k => CheckEnv -> TableSchema (k Name) -> Either String ()
verifySchema env schema = go maybeTable
where
maybeTable = M.lookup schema.name.name env.schemaMap
typeMap = schemaToTypeMap schema.columns
go Nothing = Left $ "Table " ++ schema.name.name ++ " not found"
go (Just attrs) = do
checkTypes env {ctx = [schema.name.name]} attrMap typeMap
where
attrMap = attrsToMap attrs
fetchCheckEnv :: Connection -> IO CheckEnv
fetchCheckEnv c = do
tbls <- fetchTables c >>= either (fail . show) pure
casts <- fetchCasts c >>= either (fail . show) pure
let tblMap = foldMap (\PGTable {..} -> M.singleton (T.unpack name) columns) tbls
let castMap = map (\Cast {..} -> (T.unpack source.typname, T.unpack target.typname)) $ L.filter (\Cast {context} -> context == 'i') casts
return $ CheckEnv [] tblMap castMap
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment