@@ -6,6 +6,7 @@
66{-# language FlexibleContexts #-}
77{-# language FlexibleInstances #-}
88{-# language MonoLocalBinds #-}
9+{-# language NamedFieldPuns #-}
910{-# language OverloadedStrings #-}
1011{-# language RecordWildCards #-}
1112{-# language ScopedTypeVariables #-}
@@ -37,6 +38,7 @@ import qualified Data.ByteString.Lazy
3738import Data.CaseInsensitive ( mk )
3839
3940-- containers
41+import Data.Containers.ListUtils ( nubOrdOn )
4042import qualified Data.Map.Strict as Map
4143
4244-- hasql
@@ -120,6 +122,7 @@ tests =
120122 , testLogicalFixities getTestDatabase
121123 , testUpdate getTestDatabase
122124 , testDelete getTestDatabase
125+ , testUpsert getTestDatabase
123126 , testSelectNestedPairs getTestDatabase
124127 , testSelectArray getTestDatabase
125128 , testNestedMaybeTable getTestDatabase
@@ -135,6 +138,7 @@ tests =
135138 flip run conn do
136139 sql "CREATE EXTENSION citext"
137140 sql "CREATE TABLE test_table ( column1 text not null, column2 bool not null )"
141+ sql "CREATE TABLE unique_table ( \"key\" text not null unique, \"value\" text not null )"
138142 sql "CREATE SEQUENCE test_seq"
139143
140144 return db
@@ -709,6 +713,72 @@ testDelete = databasePropertyTest "Can DELETE TestTable" \transaction -> do
709713 sort (deleted <> selected) === sort rows
710714
711715
716+data UniqueTable f = UniqueTable
717+ { uniqueTableKey :: Rel8.Column f Text
718+ , uniqueTableValue :: Rel8.Column f Text
719+ }
720+ deriving stock Generic
721+ deriving anyclass Rel8.Rel8able
722+
723+
724+deriving stock instance Eq (UniqueTable Result)
725+deriving stock instance Ord (UniqueTable Result)
726+deriving stock instance Show (UniqueTable Result)
727+
728+
729+uniqueTableSchema :: Rel8.TableSchema (UniqueTable Rel8.Name)
730+uniqueTableSchema =
731+ Rel8.TableSchema
732+ { name = "unique_table"
733+ , schema = Nothing
734+ , columns = UniqueTable
735+ { uniqueTableKey = "key"
736+ , uniqueTableValue = "value"
737+ }
738+ }
739+
740+
741+genUniqueTable :: Gen (UniqueTable Result)
742+genUniqueTable = do
743+ uniqueTableKey <- Gen.text (Range.linear 0 5) Gen.alphaNum
744+ uniqueTableValue <- Gen.text (Range.linear 0 5) Gen.alphaNum
745+ pure UniqueTable {..}
746+
747+
748+testUpsert :: IO TmpPostgres.DB -> TestTree
749+testUpsert = databasePropertyTest "Can UPSERT UniqueTable" \transaction -> do
750+ as <- unique $ forAll $ Gen.list (Range.linear 0 20) genUniqueTable
751+ bs <- unique $ forAll $ Gen.list (Range.linear 0 20) genUniqueTable
752+
753+ transaction do
754+ selected <- lift do
755+ statement () $ Rel8.insert Rel8.Insert
756+ { into = uniqueTableSchema
757+ , rows = Rel8.values $ Rel8.lit <$> as
758+ , onConflict = Rel8.DoNothing
759+ , returning = pure ()
760+ }
761+
762+ statement () $ Rel8.insert Rel8.Insert
763+ { into = uniqueTableSchema
764+ , rows = Rel8.values $ Rel8.lit <$> bs
765+ , onConflict = Rel8.DoUpdate Rel8.Upsert
766+ { index = uniqueTableKey
767+ , set = \UniqueTable {uniqueTableValue} old -> old {uniqueTableValue}
768+ , updateWhere = \_ _ -> Rel8.true
769+ }
770+ , returning = pure ()
771+ }
772+
773+ statement () $ Rel8.select do
774+ Rel8.each uniqueTableSchema
775+
776+ fromUniqueTables selected === fromUniqueTables bs <> fromUniqueTables as
777+ where
778+ unique = fmap (nubOrdOn uniqueTableKey)
779+ fromUniqueTables = Map.fromList . map \(UniqueTable key value) -> (key, value)
780+
781+
712782newtype HKNestedPair f = HKNestedPair { pairOne :: (TestTable f, TestTable f) }
713783 deriving stock Generic
714784 deriving anyclass Rel8.Rel8able