Skip to content

Commit 2eb1b2c

Browse files
Kushagra GuptaKushagra Gupta
Kushagra Gupta
authored and
Kushagra Gupta
committed
Add Postgres array functions: append, prepend, remove, replace, shuffle, sample, and to_string with tests
1 parent a5376c5 commit 2eb1b2c

File tree

2 files changed

+228
-0
lines changed

2 files changed

+228
-0
lines changed

‎beam-postgres/Database/Beam/Postgres/PgSpecific.hs

Lines changed: 149 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -81,6 +81,9 @@ module Database.Beam.Postgres.PgSpecific
8181
, arrayUpper_, arrayLower_
8282
, arrayUpperUnsafe_, arrayLowerUnsafe_
8383
, arrayLength_, arrayLengthUnsafe_
84+
, arrayAppend_, arrayPrepend_, arrayRemove_
85+
, arrayReplace_, arrayShuffle_, arraySample_
86+
, arrayToString_, arrayToStringWithNull_
8487

8588
, isSupersetOf_, isSubsetOf_
8689

@@ -410,6 +413,152 @@ isSubsetOf_ (QExpr needles) (QExpr haystack) =
410413
QExpr a ++. QExpr b =
411414
QExpr (pgBinOp "||" <$> a <*> b)
412415

416+
-- | Postgres array_append(value) function.
417+
--
418+
-- Appends an element to the end of an array. Equivalent to the
419+
-- @anycompatiblearray || anycompatible@ operator.
420+
--
421+
-- Notes:
422+
-- - The array must be empty or one-dimensional (per Postgres rules for
423+
-- concatenating an element with an array).
424+
-- - If the array is NULL, the result is NULL. If the element is NULL,
425+
-- a NULL element is appended.
426+
arrayAppend_
427+
:: QGenExpr ctxt Postgres s (V.Vector a)
428+
-> QGenExpr ctxt Postgres s a
429+
-> QGenExpr ctxt Postgres s (V.Vector a)
430+
arrayAppend_ (QExpr arr) (QExpr el) =
431+
QExpr (PgExpressionSyntax . mappend (emit "array_append") . pgParens . mconcat <$> sequenceA
432+
[ fromPgExpression <$> arr
433+
, pure (emit ", ")
434+
, fromPgExpression <$> el
435+
])
436+
437+
-- | Postgres array_prepend(value) function.
438+
--
439+
-- Prepends an element to the beginning of an array. Equivalent to the
440+
-- @anycompatible || anycompatiblearray@ operator.
441+
--
442+
-- Notes:
443+
-- - The array must be empty or one-dimensional.
444+
-- - If the array is NULL, the result is NULL. If the element is NULL,
445+
-- a NULL element is prepended.
446+
arrayPrepend_
447+
:: QGenExpr ctxt Postgres s a
448+
-> QGenExpr ctxt Postgres s (V.Vector a)
449+
-> QGenExpr ctxt Postgres s (V.Vector a)
450+
arrayPrepend_ (QExpr el) (QExpr arr) =
451+
QExpr (PgExpressionSyntax . mappend (emit "array_prepend") . pgParens . mconcat <$> sequenceA
452+
[ fromPgExpression <$> el
453+
, pure (emit ", ")
454+
, fromPgExpression <$> arr
455+
])
456+
457+
-- | Postgres array_remove(value) function.
458+
--
459+
-- Removes all elements equal to the given value from the array.
460+
-- Comparisons use @IS NOT DISTINCT FROM@ semantics, so this can remove NULLs.
461+
--
462+
-- Notes:
463+
-- - The array must be one-dimensional.
464+
-- - Returns NULL only if the array is NULL; if the value is not present,
465+
-- the original array is returned unchanged.
466+
arrayRemove_
467+
:: QGenExpr ctxt Postgres s (V.Vector a)
468+
-> QGenExpr ctxt Postgres s a
469+
-> QGenExpr ctxt Postgres s (V.Vector a)
470+
arrayRemove_ (QExpr arr) (QExpr el) =
471+
QExpr (PgExpressionSyntax . mappend (emit "array_remove") . pgParens . mconcat <$> sequenceA
472+
[ fromPgExpression <$> arr
473+
, pure (emit ", ")
474+
, fromPgExpression <$> el
475+
])
476+
477+
-- | Postgres array_replace(array, from, to) function.
478+
--
479+
-- Replaces each element equal to the second argument with the third.
480+
--
481+
-- Notes:
482+
-- - Comparisons use IS NOT DISTINCT FROM semantics; can replace NULLs.
483+
-- - The array must be one-dimensional.
484+
--
485+
-- Example:
486+
--
487+
-- @
488+
-- select_ $ pure $ arrayReplace_ (val_ $ V.fromList [1::Int32,2,5,4]) (val_ 5) (val_ 3)
489+
-- -- => {1,2,3,4}
490+
-- @
491+
arrayReplace_
492+
:: QGenExpr ctxt Postgres s (V.Vector a)
493+
-> QGenExpr ctxt Postgres s a
494+
-> QGenExpr ctxt Postgres s a
495+
-> QGenExpr ctxt Postgres s (V.Vector a)
496+
arrayReplace_ (QExpr arr) (QExpr fromVal) (QExpr toVal) =
497+
QExpr (PgExpressionSyntax . mappend (emit "array_replace") . pgParens . mconcat <$> sequenceA
498+
[ fromPgExpression <$> arr
499+
, pure (emit ", ")
500+
, fromPgExpression <$> fromVal
501+
, pure (emit ", ")
502+
, fromPgExpression <$> toVal
503+
])
504+
505+
-- | Postgres array_shuffle(array) function.
506+
-- Randomly shuffles the first dimension.
507+
arrayShuffle_
508+
:: QGenExpr ctxt Postgres s (V.Vector a)
509+
-> QGenExpr ctxt Postgres s (V.Vector a)
510+
arrayShuffle_ (QExpr arr) =
511+
QExpr (PgExpressionSyntax . mappend (emit "array_shuffle") . pgParens . fromPgExpression <$> arr)
512+
513+
-- | Postgres array_sample(array, n) function.
514+
-- Randomly selects @n@ items from the array. For multidimensional arrays,
515+
-- an "item" is a slice with a given first subscript.
516+
--
517+
-- Precondition: @n@ must not exceed the length of the first dimension.
518+
arraySample_
519+
:: Integral n
520+
=> QGenExpr ctxt Postgres s (V.Vector a)
521+
-> QGenExpr ctxt Postgres s n
522+
-> QGenExpr ctxt Postgres s (V.Vector a)
523+
arraySample_ (QExpr arr) (QExpr n) =
524+
QExpr (PgExpressionSyntax . mappend (emit "array_sample") . pgParens . mconcat <$> sequenceA
525+
[ fromPgExpression <$> arr
526+
, pure (emit ", ")
527+
, fromPgExpression <$> n
528+
])
529+
530+
-- | Postgres array_to_string(array, delimiter) function.
531+
-- Converts each element to text and joins with the delimiter. NULLs are omitted.
532+
arrayToString_
533+
:: BeamSqlBackendIsString Postgres text
534+
=> QGenExpr ctxt Postgres s (V.Vector a)
535+
-> QGenExpr ctxt Postgres s text
536+
-> QGenExpr ctxt Postgres s text
537+
arrayToString_ (QExpr arr) (QExpr delim) =
538+
QExpr (PgExpressionSyntax . mappend (emit "array_to_string") . pgParens . mconcat <$> sequenceA
539+
[ fromPgExpression <$> arr
540+
, pure (emit ", ")
541+
, fromPgExpression <$> delim
542+
])
543+
544+
-- | Postgres array_to_string(array, delimiter, null_string) function.
545+
-- Converts each element to text and joins with the delimiter. NULLs are
546+
-- represented by the provided @null_string@.
547+
arrayToStringWithNull_
548+
:: BeamSqlBackendIsString Postgres text
549+
=> QGenExpr ctxt Postgres s (V.Vector a)
550+
-> QGenExpr ctxt Postgres s text
551+
-> QGenExpr ctxt Postgres s text
552+
-> QGenExpr ctxt Postgres s text
553+
arrayToStringWithNull_ (QExpr arr) (QExpr delim) (QExpr nullStr) =
554+
QExpr (PgExpressionSyntax . mappend (emit "array_to_string") . pgParens . mconcat <$> sequenceA
555+
[ fromPgExpression <$> arr
556+
, pure (emit ", ")
557+
, fromPgExpression <$> delim
558+
, pure (emit ", ")
559+
, fromPgExpression <$> nullStr
560+
])
561+
413562
-- ** Array expressions
414563

415564
-- | An expression context that determines which types of expressions can be put

‎beam-postgres/test/Database/Beam/Postgres/Test/Select.hs

Lines changed: 79 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -5,6 +5,8 @@ module Database.Beam.Postgres.Test.Select (tests) where
55
import Data.Aeson
66
import Data.ByteString (ByteString)
77
import Data.Int
8+
import Data.List (sort)
9+
import qualified Data.Text as T
810
import qualified Data.Vector as V
911
import Test.Tasty
1012
import Test.Tasty.HUnit
@@ -24,6 +26,16 @@ tests getConn = testGroup "Selection Tests"
2426
[ testGroup "JSON"
2527
[ testPgArrayToJSON getConn
2628
]
29+
, testGroup "ARRAY functions"
30+
[ testArrayReplace getConn
31+
, testArrayShuffle getConn
32+
, testArraySample getConn
33+
, testArrayToStringBasic getConn
34+
, testArrayToStringWithNull getConn
35+
, testArrayAppend getConn
36+
, testArrayPrepend getConn
37+
, testArrayRemove getConn
38+
]
2739
, testGroup "UUID"
2840
[ testUuidFunction getConn "uuid_nil" $ \ext -> pgUuidNil ext
2941
, testUuidFunction getConn "uuid_ns_dns" $ \ext -> pgUuidNsDns ext
@@ -56,6 +68,73 @@ testPgArrayToJSON getConn = testFunction getConn "array_to_json" $ \conn -> do
5668
return $ pgArrayToJson $ val_ $ V.fromList values
5769
assertEqual "JSON list" [PgJSON $ toJSON values] actual
5870

71+
testArrayReplace :: IO ByteString -> TestTree
72+
testArrayReplace getConn = testFunction getConn "array_replace" $ \conn -> do
73+
let arr = V.fromList [1::Int32,2,5,4]
74+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
75+
pure $ arrayReplace_ (val_ arr) (val_ (5::Int32)) (val_ (3::Int32))
76+
assertEqual "array_replace" [V.fromList [1,2,3,4 :: Int32]] res
77+
78+
testArrayShuffle :: IO ByteString -> TestTree
79+
testArrayShuffle getConn = testFunction getConn "array_shuffle" $ \conn -> do
80+
let arr = V.fromList [1::Int32,2,3,4,5]
81+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
82+
pure $ arrayShuffle_ (val_ arr)
83+
-- shuffled result has same length and elements, order may change
84+
case res of
85+
[shuf] -> do
86+
assertEqual "length" (V.length arr) (V.length shuf)
87+
assertBool "is permutation"
88+
(sort (V.toList arr) == sort (V.toList shuf))
89+
_ -> assertFailure "unexpected result"
90+
91+
testArraySample :: IO ByteString -> TestTree
92+
testArraySample getConn = testFunction getConn "array_sample" $ \conn -> do
93+
let arr = V.fromList [1::Int32,2,3,4,5,6]
94+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
95+
pure $ arraySample_ (val_ arr) (val_ (3::Int32))
96+
case res of
97+
[samp] -> do
98+
assertEqual "length 3" 3 (V.length samp)
99+
assertBool "subset" (V.all (`V.elem` arr) samp)
100+
_ -> assertFailure "unexpected result"
101+
102+
testArrayToStringBasic :: IO ByteString -> TestTree
103+
testArrayToStringBasic getConn = testFunction getConn "array_to_string basic" $ \conn -> do
104+
let arr = V.fromList [1::Int32,2,3]
105+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
106+
pure $ arrayToString_ (val_ arr) (val_ ("," :: T.Text))
107+
assertEqual "join" ["1,2,3" :: T.Text] res
108+
109+
testArrayToStringWithNull :: IO ByteString -> TestTree
110+
testArrayToStringWithNull getConn = testFunction getConn "array_to_string with null" $ \conn -> do
111+
let arr :: V.Vector (Maybe T.Text)
112+
arr = V.fromList [Just "a", Nothing, Just "b"]
113+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
114+
pure $ arrayToStringWithNull_ (val_ arr) (val_ ("-" :: T.Text)) (val_ ("*" :: T.Text))
115+
assertEqual "join with null" ["a-*-b" :: T.Text] res
116+
117+
testArrayAppend :: IO ByteString -> TestTree
118+
testArrayAppend getConn = testFunction getConn "array_append" $ \conn -> do
119+
let arr = V.fromList [1::Int32,2]
120+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
121+
pure $ arrayAppend_ (val_ arr) (val_ (3::Int32))
122+
assertEqual "append" [V.fromList [1,2,3 :: Int32]] res
123+
124+
testArrayPrepend :: IO ByteString -> TestTree
125+
testArrayPrepend getConn = testFunction getConn "array_prepend" $ \conn -> do
126+
let arr = V.fromList [2::Int32,3]
127+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
128+
pure $ arrayPrepend_ (val_ (1::Int32)) (val_ arr)
129+
assertEqual "prepend" [V.fromList [1,2,3 :: Int32]] res
130+
131+
testArrayRemove :: IO ByteString -> TestTree
132+
testArrayRemove getConn = testFunction getConn "array_remove" $ \conn -> do
133+
let arr = V.fromList [1::Int32,2,3,2]
134+
res <- runBeamPostgres conn $ runSelectReturningList $ select $ do
135+
pure $ arrayRemove_ (val_ arr) (val_ (2::Int32))
136+
assertEqual "remove" [V.fromList [1,3 :: Int32]] res
137+
59138
data UuidSchema f = UuidSchema
60139
{ _uuidOssp :: f (PgExtensionEntity UuidOssp)
61140
} deriving (Generic, Database Postgres)

0 commit comments

Comments
 (0)