{-# LANGUAGE DeriveGeneric #-}
{-# LANGUAGE FlexibleContexts #-}
{-# LANGUAGE ScopedTypeVariables #-}
{-# LANGUAGE TupleSections #-}
{-# LANGUAGE ViewPatterns #-}
{-# OPTIONS_GHC -fno-warn-orphans #-}
module Database.Beam.AutoMigrate.Schema.Gen
( genSchema,
genSimilarSchemas,
SimilarSchemas (..),
shrinkSchema,
)
where
import Control.Monad
import Control.Monad.State.Strict
import Data.Foldable (foldlM)
import Data.Functor ((<&>))
import Data.Functor.Identity
import Data.Int (Int16, Int32, Int64)
import qualified Data.Map.Strict as M
import Data.Proxy
import Data.Scientific (Scientific, scientific)
import Data.Set (Set)
import qualified Data.Set as S
import Data.Text (Text)
import qualified Data.Text as T
import Data.Time (Day, LocalTime, TimeOfDay)
import Data.Word
import Database.Beam.AutoMigrate (HasColumnType, defaultColumnType, sqlSingleQuoted)
import Database.Beam.AutoMigrate.Annotated (pgDefaultConstraint)
import Database.Beam.AutoMigrate.Types
import Database.Beam.Backend.SQL (HasSqlValueSyntax, timestampType)
import qualified Database.Beam.Backend.SQL.AST as AST
import Database.Beam.Backend.SQL.Types (SqlSerial (..))
import qualified Database.Beam.Postgres as Pg
import qualified Database.Beam.Postgres.Syntax as Pg
import Database.Beam.Query (currentTimestamp_, val_)
import GHC.Generics
import Test.QuickCheck
import Test.QuickCheck.Instances.Time ()
import Text.Printf (printf)
instance Arbitrary Schema where
arbitrary = genSchema
shrink = shrinkSchema
newtype SimilarSchemas = SimilarSchemas {unSchemas :: (Schema, Schema)}
deriving (Generic, Show)
instance Arbitrary SimilarSchemas where
arbitrary = SimilarSchemas <$> genSimilarSchemas
shrink = genericShrink
genAlphaName :: Gen Text
genAlphaName = T.pack <$> vectorOf 10 (elements $ ['a' .. 'z'] ++ ['A' .. 'Z'])
genName :: (Text -> a) -> Gen a
genName f = f <$> genAlphaName
genTableName :: Gen TableName
genTableName = genName TableName
genColumnName :: Gen ColumnName
genColumnName = genName ColumnName
genUniqueConstraint :: Columns -> Gen (Set TableConstraint)
genUniqueConstraint allCols = do
someCols <- map fst . filter isStdType . take 32 <$> listOf1 (elements $ M.toList allCols)
case someCols of
[] -> pure mempty
_ -> do
constraintName <- runIdentity <$> genName Identity
pure $ S.singleton $ Unique (constraintName <> "_unique") (S.fromList someCols)
isStdType :: (ColumnName, Column) -> Bool
isStdType (_, columnType -> SqlStdType _) = True
isStdType _ = False
genPkConstraint :: Columns -> Gen (Set TableConstraint)
genPkConstraint allCols = do
someCols <- take 32 . filter (\x -> isStdType x && notNull x) <$> listOf1 (elements $ M.toList allCols)
case someCols of
[] -> pure mempty
_ -> do
constraintName <- runIdentity <$> genName Identity
pure $ S.singleton $ PrimaryKey (constraintName <> "_pk") (S.fromList $ map fst someCols)
where
notNull :: (ColumnName, Column) -> Bool
notNull (_, col) = NotNull `S.member` columnConstraints col
genTableConstraints :: Tables -> Columns -> Gen (Set TableConstraint)
genTableConstraints _allOtherTables ourColums =
frequency
[ (60, pure mempty),
(30, genUniqueConstraint ourColums),
(30, genPkConstraint ourColums),
(15, mappend <$> genPkConstraint ourColums <*> genUniqueConstraint ourColums)
]
genColumnType :: Gen (ColumnType, ColumnConstraint)
genColumnType =
oneof
[ genSqlStdType,
genPgSpecificType
]
genSqlStdType :: Gen (ColumnType, ColumnConstraint)
genSqlStdType =
oneof
[ genType arbitrary (Proxy @Int32),
genType arbitrary (Proxy @Int16),
genType arbitrary (Proxy @Int64),
genType arbitrary (Proxy @Word16),
genType arbitrary (Proxy @Word32),
genType arbitrary (Proxy @Word64),
genType genAlphaName (Proxy @Text),
genBitStringType,
genType (elements [-0.1, 3.5]) (Proxy @Double),
genType (pure (scientific (1 :: Integer) (1 :: Int))) (Proxy @Scientific),
genType arbitrary (Proxy @Day),
genType (pure (read "01:00:07.979173" :: TimeOfDay)) (Proxy @TimeOfDay),
genType arbitrary (Proxy @Bool),
genType (pure (read "1864-05-10 13:50:45.919197" :: LocalTime)) (Proxy @LocalTime),
pure (SqlStdType $ timestampType Nothing False, pgDefaultConstraint @LocalTime currentTimestamp_),
genType (fmap SqlSerial arbitrary) (Proxy @(SqlSerial Int64))
]
genType ::
forall a.
( HasColumnType a,
HasSqlValueSyntax Pg.PgValueSyntax a
) =>
Gen a ->
Proxy a ->
Gen (ColumnType, ColumnConstraint)
genType gen Proxy =
(,) <$> pure (defaultColumnType (Proxy @a))
<*> (gen <&> (\(x :: a) -> pgDefaultConstraint $ val_ x))
genBitStringType :: Gen (ColumnType, ColumnConstraint)
genBitStringType = do
varying <- arbitrary
charPrec <- elements [1 :: Word, 2, 4, 6, 8, 16, 32, 64]
string <- vectorOf (fromIntegral charPrec) (elements ['0', '1'])
let txt = sqlSingleQuoted (T.pack string) <> "::bit(" <> (T.pack . show $ charPrec) <> ")"
case varying of
False ->
pure
( SqlStdType $ AST.DataTypeBit False (Just charPrec),
Default txt
)
True ->
pure
( SqlStdType $ AST.DataTypeBit True (Just 1),
Default txt
)
genPgSpecificType :: Gen (ColumnType, ColumnConstraint)
genPgSpecificType =
oneof
[ genType (fmap Pg.PgJSON (arbitrary @Int)) (Proxy @(Pg.PgJSON Int)),
genType (fmap Pg.PgJSONB (arbitrary @Int)) (Proxy @(Pg.PgJSONB Int))
]
_genRangeType ::
forall a n.
( Ord a,
Num a,
Arbitrary a,
Pg.PgIsRange n,
HasColumnType (Pg.PgRange n a),
HasSqlValueSyntax Pg.PgValueSyntax a
) =>
Proxy n ->
Proxy a ->
Gen (ColumnType, ColumnConstraint)
_genRangeType Proxy Proxy = do
let colType = defaultColumnType (Proxy @(Pg.PgRange n a))
lowerBoundRange <- elements [Pg.Inclusive, Pg.Exclusive]
upperBoundRange <- elements [Pg.Inclusive, Pg.Exclusive]
mbLower <- arbitrary @(Maybe a)
mbUpper <-
arbitrary @(Maybe (Positive a)) <&> \u -> case liftM2 (,) u mbLower of
Nothing -> u
Just (ub, lb) -> Just $ Positive $ (getPositive ub) + lb + 1
let dVal =
pgDefaultConstraint $
Pg.range_ @n @a lowerBoundRange upperBoundRange (val_ mbLower) (val_ (fmap getPositive mbUpper))
pure $ (colType, dVal)
_genDbEnumeration :: Gen (ColumnType, Text)
_genDbEnumeration = do
vals <- map sqlSingleQuoted <$> listOf1 genAlphaName
dVal <- elements vals
name <- genName EnumerationName
pure (DbEnumeration name (Enumeration vals), dVal)
_defVal :: forall a. (Arbitrary a, Show a) => Proxy a -> Gen Text
_defVal Proxy = T.pack . show <$> (arbitrary :: Gen a)
_genFloatType :: Gen (AST.DataType, Text)
_genFloatType = do
floatPrec <- elements [Nothing, Just 4, Just 8]
def <- _defVal @Float Proxy
pure (AST.DataTypeFloat floatPrec, def)
_genRealType :: Gen (AST.DataType, Text)
_genRealType = do
v <- T.pack . printf "%.1f" <$> arbitrary @Float
pure (AST.DataTypeReal, sqlSingleQuoted v <> "::real")
_genIntType :: Gen (AST.DataType, Text)
_genIntType = do
v <- arbitrary @Int32
pure $
if v < 0
then (AST.DataTypeBigInt, sqlSingleQuoted (T.pack . show $ v) <> "::integer")
else (AST.DataTypeBigInt, T.pack . show $ v)
_genBigIntType :: Gen (AST.DataType, Text)
_genBigIntType = do
v <- arbitrary @Integer
pure $
if v < 0
then (AST.DataTypeBigInt, sqlSingleQuoted (T.pack . show $ v) <> "::integer")
else (AST.DataTypeBigInt, T.pack . show $ v)
_genDoublePrecisionType :: Gen (AST.DataType, Text)
_genDoublePrecisionType = do
v <- T.pack . printf "%.1f" <$> arbitrary @Double
pure (AST.DataTypeDoublePrecision, sqlSingleQuoted v <> "::double precision")
_genNumericType :: (Maybe (Word, Maybe Word) -> AST.DataType) -> Text -> Gen (AST.DataType, Text)
_genNumericType f _cast = do
numPrec <- choose (1 :: Word, 15)
numScale <- choose (1 :: Word, numPrec)
p <- elements [Nothing, Just (numPrec, Nothing), Just (numPrec, Just numScale)]
let renderNum (a, b) = (T.pack . show $ a) <> "." <> (T.pack . show $ b)
defaultValue <- case p of
Nothing ->
oneof
[ _defVal @Int32 Proxy,
fmap renderNum ((,) <$> choose (0 :: Word, 131072) <*> choose (0 :: Word, 16383))
]
Just (_, Nothing) ->
oneof
[ _defVal @Int32 Proxy,
fmap
renderNum
( (,) <$> (choose (0 :: Word, 131072))
<*> choose (0 :: Word, 16383)
)
]
Just (_, Just _) ->
oneof
[ _defVal @Int32 Proxy,
fmap
renderNum
( (,) <$> choose (0 :: Word, 131072)
<*> choose (0 :: Word, 16383)
)
]
pure (f p, defaultValue)
_pgSimplify :: (AST.DataType, Text) -> (AST.DataType, Text)
_pgSimplify = \case
(AST.DataTypeChar varying Nothing c, def) -> (AST.DataTypeChar varying (Just 1) c, def)
(AST.DataTypeNationalChar varying Nothing, def) -> (AST.DataTypeChar varying (Just 1) Nothing, def)
(AST.DataTypeNationalChar varying precision, def) -> (AST.DataTypeChar varying precision Nothing, def)
(AST.DataTypeDecimal v, def) -> (AST.DataTypeNumeric v, def)
(AST.DataTypeFloat (Just 4), def) -> (AST.DataTypeReal, def)
(AST.DataTypeFloat (Just 8), def) -> (AST.DataTypeDoublePrecision, def)
x -> x
_genCharType :: (Bool -> Maybe Word -> AST.DataType) -> Gen (AST.DataType, Text)
_genCharType f = do
varying <- arbitrary
text <- genAlphaName
charPrec <- choose (1, 2048)
case varying of
False -> pure (f False (Just charPrec), sqlSingleQuoted (T.take (fromIntegral charPrec) text) <> "::bpchar")
True -> pure (f True (Just 1), sqlSingleQuoted text <> "::character varying")
genColumn :: Columns -> Gen Column
genColumn _allColums = do
constNum <- choose (0, 2)
(cType, dVal) <- genColumnType
constrs <- vectorOf constNum (elements [NotNull, dVal])
pure $ Column cType (S.fromList constrs)
genColumns :: Gen Columns
genColumns = do
colNum <- choose (1, 50)
columnNames <- vectorOf colNum genColumnName
foldlM (\acc cName -> flip (M.insert cName) acc <$> genColumn acc) mempty columnNames
genTable :: Tables -> Gen Table
genTable currentTables = do
cols <- genColumns
Table <$> genTableConstraints currentTables cols <*> pure cols
genSchema :: Gen Schema
genSchema = sized $ \tableNum -> do
tableNames <- vectorOf tableNum genTableName
tbls <- foldlM (\acc tName -> flip (M.insert tName) acc <$> genTable acc) mempty tableNames
pure $ Schema tbls mempty mempty
data TablesEditAction
= AddTable
| DropTable
| ModifyTable
| LeaveTableAlone
data TableEditAction
= AddColumn
| DropColumn
| ModifyColumn
| LeaveColumnAlone
data ColumnEditAction
= ChangeType
| ChangeConstraints
| NoChange
genSimilarSchemas :: Gen (Schema, Schema)
genSimilarSchemas = do
initialSchema <- genSchema
(initialSchema,) <$> fmap (\tbs -> Schema tbs mempty mempty) (similarTables (schemaTables initialSchema))
similarTables :: Tables -> Gen Tables
similarTables tbls = flip execStateT tbls $
forM_ (M.toList tbls) $ \(tName, tbl) -> do
tableEditAction <-
lift $
frequency
[ (1, pure AddTable),
(1, pure DropTable),
(1, pure ModifyTable),
(15, pure LeaveTableAlone)
]
case tableEditAction of
AddTable -> do
s <- get
newTableName <- lift genTableName
newTable <- lift $ genTable s
modify' (M.insert newTableName newTable)
DropTable -> modify' (M.delete tName)
ModifyTable -> do
table' <- lift $ similarTable tbl
modify' (M.insert tName table')
LeaveTableAlone -> pure ()
similarTable :: Table -> Gen Table
similarTable tbl = flip execStateT tbl $
forM_ (M.toList . tableColumns $ tbl) $ \(cName, col) -> do
tableEditAction <-
lift $
frequency
[ (1, pure AddColumn),
(1, pure DropColumn),
(1, pure ModifyColumn),
(15, pure LeaveColumnAlone)
]
case tableEditAction of
AddColumn -> do
s <- get
newColumnName <- lift genColumnName
newColumn <- lift $ genColumn (tableColumns s)
modify' (\st -> st {tableColumns = M.insert newColumnName newColumn (tableColumns st)})
DropColumn ->
modify'
( \st ->
st
{ tableColumns = M.delete cName (tableColumns st),
tableConstraints = deleteConstraintReferencing cName (tableConstraints st)
}
)
ModifyColumn -> do
col' <- lift $ similarColumn col
modify'
( \st ->
st
{ tableColumns = M.insert cName col' (tableColumns st),
tableConstraints = deleteConstraintReferencing cName (tableConstraints st)
}
)
LeaveColumnAlone -> pure ()
deleteConstraintReferencing :: ColumnName -> Set TableConstraint -> Set TableConstraint
deleteConstraintReferencing cName conss = S.filter (not . doesReference) conss
where
doesReference :: TableConstraint -> Bool
doesReference = \case
PrimaryKey _ refs -> S.member cName refs
ForeignKey _ _ refs _ _ -> let ours = S.map snd refs in S.member cName ours
Unique _ refs -> S.member cName refs
similarColumn :: Column -> Gen Column
similarColumn col = do
editAction' <-
frequency
[ (15, pure ChangeType),
(10, pure ChangeConstraints),
(30, pure NoChange)
]
case editAction' of
ChangeType -> do
(newType, newDef) <- genColumnType
let oldConstraints = S.filter (\c -> case c of Default _ -> False; _ -> True) (columnConstraints col)
pure $
col
{ columnType = newType,
columnConstraints = S.insert newDef oldConstraints
}
ChangeConstraints -> do
let oldConstraints = columnConstraints col
let newConstraints = case S.toList oldConstraints of
[] -> S.singleton NotNull
[NotNull] -> mempty
_ -> oldConstraints
pure $ col {columnConstraints = newConstraints}
NoChange -> pure col
shrinkSchema :: Schema -> [Schema]
shrinkSchema s =
noSchema : concatMap shrinkTable (M.toList (schemaTables s))
where
shrinkTable :: (TableName, Table) -> [Schema]
shrinkTable (tName, tbl) =
s {schemaTables = M.delete tName (schemaTables s)} :
concatMap (shrinkColumns tName tbl) (M.toList (tableColumns tbl))
shrinkColumns :: TableName -> Table -> (ColumnName, Column) -> [Schema]
shrinkColumns tName tbl (cName, _col) =
let tbl' =
tbl
{ tableColumns = M.delete cName (tableColumns tbl),
tableConstraints = deleteConstraintReferencing cName (tableConstraints tbl)
}
in [s {schemaTables = M.insert tName tbl' (schemaTables s)}]