Merge pull request #741 from haskell-beam/fix-354

Support Postgres array migrations
haskell-beam · Jan 5, 2025 · 8bab511 · 8bab511
2 parents 00eeec4 + 2d2a30d
commit 8bab511
Show file tree

Hide file tree

Showing 3 changed files with 92 additions and 74 deletions.
diff --git a/beam-postgres/ChangeLog.md b/beam-postgres/ChangeLog.md
@@ -1,3 +1,9 @@
+# Unreleased
+
+## Bug fixes
+
+ * Added the ability to migrate Postgres' array types (#354).
+
 # 0.5.4.2
 
 ## Bug fixes

diff --git a/beam-postgres/Database/Beam/Postgres/Syntax.hs b/beam-postgres/Database/Beam/Postgres/Syntax.hs
@@ -580,8 +580,8 @@ instance IsSql99DataTypeSyntax PgDataTypeSyntax where
   binaryLargeObjectType = pgByteaType { pgDataTypeSerialized = binaryLargeObjectType }
   booleanType = PgDataTypeSyntax (PgDataTypeDescrOid (Pg.typoid Pg.bool) Nothing) (emit "BOOLEAN")
                                  booleanType
-  arrayType (PgDataTypeSyntax _ syntax serialized) sz =
-    PgDataTypeSyntax (error "TODO: array migrations")
+  arrayType (PgDataTypeSyntax descr syntax serialized) sz =
+    PgDataTypeSyntax (PgDataTypeDescrOid (fromMaybe (error "Unsupported array type") (arrayTypeDescr descr)) Nothing)
                      (syntax <> emit "[" <> emit (fromString (show sz)) <> emit "]")
                      (arrayType serialized sz)
   rowType = error "rowType"
@@ -647,12 +647,73 @@ pgLineType = PgDataTypeSyntax (PgDataTypeDescrOid (Pg.typoid Pg.line) Nothing) (
 pgLineSegmentType = PgDataTypeSyntax (PgDataTypeDescrOid (Pg.typoid Pg.lseg) Nothing) (emit "LSEG") (pgDataTypeJSON "lseg")
 pgBoxType = PgDataTypeSyntax (PgDataTypeDescrOid (Pg.typoid Pg.box) Nothing) (emit "BOX") (pgDataTypeJSON "box")
 
+-- TODO: better mechanism to tell, at compile time, that some type
+-- cannot be placed in an array
 pgUnboundedArrayType :: PgDataTypeSyntax -> PgDataTypeSyntax
-pgUnboundedArrayType (PgDataTypeSyntax _ syntax serialized) =
-    PgDataTypeSyntax (error "Can't do array migrations yet")
+pgUnboundedArrayType (PgDataTypeSyntax descr syntax serialized) =
+    PgDataTypeSyntax (PgDataTypeDescrOid (fromMaybe (error "Unsupported array type") (arrayTypeDescr descr)) Nothing)
                      (syntax <> emit "[]")
                      (pgDataTypeJSON (object [ "unbounded-array" .= fromBeamSerializedDataType serialized ]))
 
+-- TODO: define CPP macro to make sure the left hand side (e.g. `Pg.recordOid`) 
+--       always matches right hand side (e.g. `Pg.array_recordOid)
+
+-- | Get the Oid of Pg arrays which contains elements of a certain type
+arrayTypeDescr :: PgDataTypeDescr -> Maybe Pg.Oid
+arrayTypeDescr (PgDataTypeDescrDomain _) = Nothing
+arrayTypeDescr (PgDataTypeDescrOid elemOid _)
+  | elemOid == Pg.recordOid    = Just $ Pg.array_recordOid
+  | elemOid == Pg.xmlOid    = Just $ Pg.array_xmlOid
+  | elemOid == Pg.jsonOid   = Just $ Pg.array_jsonOid
+  | elemOid == Pg.lineOid   = Just $ Pg.array_lineOid
+  | elemOid == Pg.cidrOid   = Just $ Pg.array_cidOid
+  | elemOid == Pg.circleOid = Just $ Pg.array_circleOid
+  | elemOid == Pg.moneyOid  = Just $ Pg.array_moneyOid
+  | elemOid == Pg.boolOid   = Just $ Pg.array_boolOid
+  | elemOid == Pg.byteaOid  = Just $ Pg.array_byteaOid
+  | elemOid == Pg.charOid   = Just $ Pg.array_charOid
+  | elemOid == Pg.nameOid   = Just $ Pg.array_nameOid
+  | elemOid == Pg.int2Oid   = Just $ Pg.array_int2Oid
+  | elemOid == Pg.int2vectorOid = Just $ Pg.array_int2vectorOid
+  | elemOid == Pg.int4Oid   = Just $ Pg.array_int4Oid
+  | elemOid == Pg.regprocOid = Just $ Pg.array_regprocOid
+  | elemOid == Pg.textOid   = Just $ Pg.array_textOid
+  | elemOid == Pg.tidOid   = Just $ Pg.array_tidOid
+  | elemOid == Pg.xidOid   = Just $ Pg.array_xidOid
+  | elemOid == Pg.cidOid   = Just $ Pg.array_cidOid
+  | elemOid == Pg.bpcharOid = Just $ Pg.array_bpcharOid
+  | elemOid == Pg.varcharOid = Just $ Pg.array_varcharOid
+  | elemOid == Pg.int8Oid = Just $ Pg.array_int8Oid
+  | elemOid == Pg.pointOid = Just $ Pg.array_pointOid
+  | elemOid == Pg.lsegOid = Just $ Pg.array_lsegOid
+  | elemOid == Pg.pathOid = Just $ Pg.array_pathOid
+  | elemOid == Pg.boxOid = Just $ Pg.array_boxOid
+  | elemOid == Pg.float4Oid = Just $ Pg.array_float4Oid
+  | elemOid == Pg.float8Oid = Just $ Pg.array_float8Oid
+  | elemOid == Pg.polygonOid = Just $ Pg.array_polygonOid
+  | elemOid == Pg.oidOid = Just $ Pg.array_oidOid
+  | elemOid == Pg.macaddrOid = Just $ Pg.array_macaddrOid
+  | elemOid == Pg.inetOid = Just $ Pg.array_inetOid
+  | elemOid == Pg.timestampOid = Just $ Pg.array_timestampOid
+  | elemOid == Pg.dateOid = Just $ Pg.array_dateOid
+  | elemOid == Pg.timeOid = Just $ Pg.array_timeOid
+  | elemOid == Pg.timestamptzOid = Just $ Pg.array_timestamptzOid
+  | elemOid == Pg.intervalOid = Just $ Pg.array_intervalOid
+  | elemOid == Pg.numericOid = Just $ Pg.array_numericOid
+  | elemOid == Pg.timetzOid = Just $ Pg.array_timetzOid
+  | elemOid == Pg.bitOid = Just $ Pg.array_bitOid
+  | elemOid == Pg.varbitOid = Just $ Pg.array_varbitOid
+  | elemOid == Pg.refcursorOid = Just $ Pg.array_refcursorOid
+  | elemOid == Pg.regprocedureOid = Just $ Pg.array_regprocedureOid
+  | elemOid == Pg.regoperOid = Just $ Pg.array_regoperOid
+  | elemOid == Pg.regoperatorOid = Just $ Pg.array_regoperatorOid
+  | elemOid == Pg.regclassOid = Just $ Pg.array_regclassOid
+  | elemOid == Pg.regtypeOid = Just $ Pg.array_regtypeOid
+  | elemOid == Pg.uuidOid = Just $ Pg.array_uuidOid
+  | elemOid == Pg.jsonbOid = Just $ Pg.array_jsonbOid
+  | otherwise = Nothing
+
+
 pgTsQueryTypeInfo :: Pg.TypeInfo
 pgTsQueryTypeInfo = Pg.Basic (Pg.Oid 3615) 'U' ',' "tsquery"
 

diff --git a/beam-postgres/test/Database/Beam/Postgres/Test/Marshal.hs b/beam-postgres/test/Database/Beam/Postgres/Test/Marshal.hs
@@ -51,6 +51,10 @@ boxGen = do PgPoint x1 y1 <- pointGen
             pure (PgBox (PgPoint (min x1 x2) (min y1 y2))
                         (PgPoint (max x1 x2) (max y1 y2)))
 
+arrayGen :: Hedgehog.Gen a -> Hedgehog.Gen (Vector.Vector a)
+arrayGen = fmap Vector.fromList 
+         . Gen.list (Range.linear 0 5) -- small arrays == quick tests
+
 boxCmp :: PgBox -> PgBox -> Bool
 boxCmp (PgBox a1 b1) (PgBox a2 b2) =
     (a1 `ptCmp` a2 && b1 `ptCmp` b2) ||
@@ -90,13 +94,27 @@ tests postgresConn =
     , marshalTest (Gen.maybe (Gen.integral (Range.constantBounded @Word64))) postgresConn
     , marshalTest (Gen.maybe textGen) postgresConn
     , marshalTest (Gen.maybe uuidGen) postgresConn
-    , marshalTest692 postgresConn
 
     , marshalTest' (\a b -> Hedgehog.assert (liftEq ptCmp a b))  (Gen.maybe pointGen) postgresConn
     , marshalTest' (\a b -> Hedgehog.assert (liftEq boxCmp a b)) (Gen.maybe boxGen) postgresConn
 
---    , marshalTest (Gen.double  (Range.exponentialFloat 0 1e40))  postgresConn
---    , marshalTest (Gen.integral (Range.constantBounded @Word))   postgresConn
+    -- Arrays
+    --
+    -- Testing lots of element types for arrays is important, because 
+    -- the mapping between array Oid and element Oid is not type 
+    -- safe, and hence error-prone.
+    , marshalTest (arrayGen textGen) postgresConn
+    , marshalTest (arrayGen (Gen.double (Range.exponentialFloat 0 1e40))) postgresConn
+    , marshalTest (arrayGen ((Gen.integral (Range.constantBounded @Int16)))) postgresConn
+    , marshalTest (arrayGen ((Gen.integral (Range.constantBounded @Int32)))) postgresConn
+    , marshalTest (arrayGen ((Gen.integral (Range.constantBounded @Int64)))) postgresConn
+    , marshalTest (Gen.maybe (arrayGen textGen)) postgresConn
+    , marshalTest (Gen.maybe (arrayGen (Gen.double (Range.exponentialFloat 0 1e40)))) postgresConn
+    , marshalTest (Gen.maybe (arrayGen ((Gen.integral (Range.constantBounded @Int16))))) postgresConn
+    , marshalTest (Gen.maybe (arrayGen ((Gen.integral (Range.constantBounded @Int32))))) postgresConn
+    , marshalTest (Gen.maybe (arrayGen ((Gen.integral (Range.constantBounded @Int64))))) postgresConn
+
+    , marshalTest (Gen.double  (Range.exponentialFloat 0 1e40))  postgresConn
 --    , marshalTest (Gen.integral (Range.constantBounded @Int))    postgresConn
 
 --    , marshalTest @Int8    postgresConn
@@ -164,70 +182,3 @@ marshalTest' cmp gen postgresConn =
       v' `cmp` a
 
     assertBool "Hedgehog test failed" passes
-
-
--- Ensure that both `Vector Text` and `Maybe (Vector Text)` can be 
--- marshalled correctly (see issue 692).
---
--- At this time, the postgres migration backend can't create columns of arrays,
--- and hence this test does not use `marshalTest`.
-marshalTest692 :: IO ByteString -> TestTree
-marshalTest692 postgresConn = 
-  testCase "Can marshal Vector Text and Maybe (Vector Text) (#692)" $ 
-  withTestPostgres ("db_marshal_maybe_vector_text_issue_692") postgresConn $ \conn -> do
-    liftIO $ execute_ conn $ "CREATE TABLE mytable (\nmyid SERIAL PRIMARY KEY, mycolumn text[], mynullablecolumn text[]\n);"
-
-    passes <- Hedgehog.check . Hedgehog.property $ do
-      nullable <- Hedgehog.forAll (Gen.maybe (Vector.fromList <$> (Gen.list (Range.linear 0 10) textGen)))
-      nonnull <- Hedgehog.forAll (Vector.fromList <$> (Gen.list (Range.linear 0 10) textGen))
-
-      [MkTbl692 rowId v vnull] <-
-        liftIO . runBeamPostgres conn 
-          $ runInsertReturningList 
-            $ insert (_myTable myDB) 
-              $ insertExpressions [ MkTbl692 default_ (val_ nonnull) (val_ nullable) ]
-
-      v === nonnull
-      vnull === nullable
-
-      Just (MkTbl692 _ v' vnull') <-
-          liftIO . runBeamPostgres conn 
-            $ runSelectReturningOne (lookup_ (_myTable myDB) (Tbl692Key rowId))
-      v' === nonnull
-      vnull' === nullable
-
-    assertBool "Hedgehog test failed" passes
-  where
-    myDB :: DatabaseSettings Postgres MyDB692
-    myDB = defaultDbSettings `withDbModification`
-            MkMyDB692 {
-              _myTable =
-                setEntityName "mytable" <>
-                modifyTableFields
-                  tableModification {
-                    myid = fieldNamed "myid",
-                    mycolumn = fieldNamed "mycolumn",
-                    mynullablecolumn = fieldNamed "mynullablecolumn"
-                  }
-            }
-
-data Tbl692 f
-    = MkTbl692 
-    { myid             :: C f (SqlSerial Int32)
-    , mycolumn         :: C f (Vector.Vector T.Text)
-    , mynullablecolumn :: C f (Maybe (Vector.Vector T.Text)) 
-    }
-      deriving (Generic, Beamable)
-
-deriving instance Show (Tbl692 Identity)
-deriving instance Eq (Tbl692 Identity)
-
-instance Table Tbl692 where
-    data PrimaryKey Tbl692 f = Tbl692Key (C f (SqlSerial Int32))
-      deriving (Generic, Beamable)
-    primaryKey = Tbl692Key <$> myid
-data MyDB692 entity
-    = MkMyDB692
-    { _myTable :: entity (TableEntity Tbl692)
-    } deriving (Generic)
-instance Database Postgres MyDB692