Skip to content

Commit d089ee4

Browse files
jatcwangjatcwang
committed
Better postgres array of enum support
Add a helper `doobie.postgres.implicits.arrayOfEnum` with better typechecking support. For an array column, postgres prefixes `_` in front of the type name, therefore the expected vendor type name is e.g. `_myenum` Also added some tests / examples to test array of enum where the enum is defined in another schema. (e.g. Reported vendor type name is `"other_schema"."_other_enum"` for the column type `other_schema.other_enum[]`)
1 parent 0f190f4 commit d089ee4

File tree

7 files changed

+253
-15
lines changed

7 files changed

+253
-15
lines changed

‎init/postgres/test-db.sql‎

Lines changed: 8 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -6,6 +6,14 @@ create extension postgis;
66
create extension hstore;
77
create type myenum as enum ('foo', 'bar', 'invalid');
88

9+
create schema other_schema;
10+
11+
set search_path to other_schema;
12+
13+
create type other_enum as enum ('a', 'b');
14+
15+
set search_path to public;
16+
917
--
1018
-- The sample data used in the world database is Copyright Statistics
1119
-- Finland, http://www.stat.fi/worldinfigures.

‎modules/core/src/main/scala/doobie/util/meta/meta.scala‎

Lines changed: 5 additions & 5 deletions
Original file line numberDiff line numberDiff line change
@@ -130,13 +130,13 @@ trait MetaConstructors {
130130
)
131131

132132
def array[A >: Null <: AnyRef](
133-
elementType: String,
134-
schemaH: String,
135-
schemaT: String*
133+
elementTypeName: String, // Used in Put to set the array element type
134+
arrayTypeName: String,
135+
additionalArrayTypeNames: String*
136136
): Meta[Array[A]] =
137137
new Meta[Array[A]](
138-
Get.Advanced.array[A](NonEmptyList(schemaH, schemaT.toList)),
139-
Put.Advanced.array[A](NonEmptyList(schemaH, schemaT.toList), elementType)
138+
Get.Advanced.array[A](NonEmptyList(arrayTypeName, additionalArrayTypeNames.toList)),
139+
Put.Advanced.array[A](NonEmptyList(arrayTypeName, additionalArrayTypeNames.toList), elementTypeName)
140140
)
141141

142142
def other[A >: Null <: AnyRef: TypeName: ClassTag](

‎modules/docs/src/main/mdoc/docs/11-Arrays.md‎

Lines changed: 47 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -60,7 +60,7 @@ val create =
6060
(drop *> create).unsafeRunSync()
6161
```
6262

63-
**doobie** maps SQL array columns to `Array`, `List`, and `Vector` by default. No special handling is required, other than importing the vendor-specific array support above.
63+
**doobie** maps SQL array columns to `Array`, `List`, and `Vector` by default for standard types like `String` or `Int`. No special handling is required, other than importing the vendor-specific array support above.
6464

6565
```scala mdoc:silent
6666
case class Person(id: Long, name: String, pets: List[String])
@@ -93,3 +93,49 @@ sql"select array['foo','bar','baz']".query[Option[List[String]]].quick.unsafeRun
9393
sql"select array['foo',NULL,'baz']".query[List[Option[String]]].quick.unsafeRunSync()
9494
sql"select array['foo',NULL,'baz']".query[Option[List[Option[String]]]].quick.unsafeRunSync()
9595
```
96+
97+
### Array of enums
98+
99+
For reading from and writing to a column that is an array of enum, you can use `doobie.postgres.implicits.arrayOfEnum`
100+
to create a `Meta` instance for your enum type:
101+
102+
```scala mdoc
103+
import doobie.postgres.implicits.arrayOfEnum
104+
105+
sealed trait MyEnum
106+
107+
object MyEnum {
108+
case object Foo extends MyEnum
109+
110+
case object Bar extends MyEnum
111+
112+
private val typeName = "myenum"
113+
114+
def fromStrUnsafe(s: String): MyEnum = s match {
115+
case "foo" => Foo
116+
case "bar" => Bar
117+
case other => throw new RuntimeException(s"Unexpected value '$other' for MyEnum")
118+
}
119+
120+
def toStr(e: MyEnum): String = e match {
121+
case Foo => "foo"
122+
case Bar => "bar"
123+
}
124+
125+
implicit val MyEnumArrayMeta: Meta[Array[MyEnum]] =
126+
arrayOfEnum[MyEnum](
127+
enumTypeName = typeName,
128+
fromStr = fromStrUnsafe,
129+
toStr = toStr
130+
)
131+
132+
}
133+
```
134+
135+
and you can now map the array of enum column into an `Array[MyEnum]`, `List[MyEnum]`, `Vector[MyEnum]`:
136+
137+
```scala mdoc
138+
sql"select array['foo', 'bar'] :: myenum[]".query[List[MyEnum]].quick.unsafeRunSync()
139+
```
140+
141+
For an example of using an enum type from another schema, please see [OtherEnum.scala](https://github.com/typelevel/doobie/blob/main/modules/postgres/src/test/scala/doobie/postgres/enums/OtherEnum.scala)

‎modules/postgres/src/main/scala/doobie/postgres/Instances.scala‎

Lines changed: 22 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -170,6 +170,28 @@ trait Instances {
170170
.timap(_.map(_.map(a => if (a == null) null else BigDecimal.apply(a))))(_.map(_.map(a =>
171171
if (a == null) null else a.bigDecimal)))
172172

173+
/** Create a Meta instance to allow reading and writing into an array of enum, with stricter typechecking support to
174+
* verify that the column we're inserting into must match the enum array type.
175+
*
176+
* @param enumTypeName
177+
* Name of the enum type
178+
* @param fromStr
179+
* Function to convert each element to the Scala type when reading from the database
180+
* @param toStr
181+
* Function to convert each element to string when writing to the database
182+
* @return
183+
*/
184+
def arrayOfEnum[A: ClassTag](
185+
enumTypeName: String,
186+
fromStr: String => A,
187+
toStr: A => String
188+
): Meta[Array[A]] = {
189+
Meta.Advanced.array[String](
190+
enumTypeName,
191+
arrayTypeName = s"_$enumTypeName"
192+
).timap(arr => arr.map(fromStr))(arr => arr.map(toStr))
193+
}
194+
173195
// So, it turns out that arrays of structs don't work because something is missing from the
174196
// implementation. So this means we will only be able to support primitive types for arrays.
175197
//

‎modules/postgres/src/test/scala/doobie/postgres/PgArraySuite.scala‎

Lines changed: 114 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,114 @@
1+
// Copyright (c) 2013-2020 Rob Norris and Contributors
2+
// This software is licensed under the MIT License (MIT).
3+
// For more information see LICENSE or https://opensource.org/licenses/MIT
4+
5+
package doobie.postgres
6+
7+
import cats.effect.IO
8+
import doobie.Transactor
9+
import doobie.postgres.enums.{MyEnum, OtherEnum}
10+
import doobie.postgres.implicits.*
11+
import doobie.syntax.all.*
12+
import doobie.util.analysis.{ColumnTypeError, ParameterTypeError}
13+
14+
class PgArraySuite extends munit.CatsEffectSuite {
15+
16+
val transactor: Transactor[IO] = Transactor.fromDriverManager[IO](
17+
driver = "org.postgresql.Driver",
18+
url = "jdbc:postgresql:world",
19+
user = "postgres",
20+
password = "password",
21+
logHandler = None
22+
)
23+
24+
private val listOfMyEnums: List[MyEnum] = List(MyEnum.Foo, MyEnum.Bar)
25+
26+
private val listOfOtherEnums: List[OtherEnum] = List(OtherEnum.A, OtherEnum.B)
27+
28+
test("array of custom string type: read correctly and typechecks") {
29+
val q = sql"select array['foo', 'bar'] :: myenum[]".query[List[MyEnum]]
30+
(for {
31+
_ <- q.analysis
32+
.map(ana => assertEquals(ana.columnAlignmentErrors, List.empty))
33+
34+
_ <- q.unique.map(assertEquals(_, listOfMyEnums))
35+
36+
_ <- sql"select array['foo', 'bar']".query[List[MyEnum]].analysis.map(_.columnAlignmentErrors)
37+
.map {
38+
case List(e: ColumnTypeError) => assertEquals(e.schema.vendorTypeName, "_text")
39+
case other => fail(s"Unexpected typecheck result: $other")
40+
}
41+
} yield ())
42+
.transact(transactor)
43+
}
44+
45+
test("array of custom string type: writes correctly and typechecks") {
46+
val q = sql"insert into temp_myenum (arr) values ($listOfMyEnums)".update
47+
(for {
48+
_ <- sql"drop table if exists temp_myenum".update.run
49+
_ <- sql"create table temp_myenum(arr myenum[] not null)".update.run
50+
_ <- q.analysis.map(_.columnAlignmentErrors).map(ana => assertEquals(ana, List.empty))
51+
_ <- q.run
52+
_ <- sql"select arr from temp_myenum".query[List[MyEnum]].unique
53+
.map(assertEquals(_, listOfMyEnums))
54+
55+
_ <- sql"insert into temp_myenum (arr) values (${List("foo")})".update.analysis
56+
.map(_.parameterAlignmentErrors)
57+
.map {
58+
case List(e: ParameterTypeError) => assertEquals(e.vendorTypeName, "_myenum")
59+
case other => fail(s"Unexpected typecheck result: $other")
60+
}
61+
} yield ())
62+
.transact(transactor)
63+
}
64+
65+
test("array of custom type in another schema: read correctly and typechecks") {
66+
val q = sql"select array['a', 'b'] :: other_schema.other_enum[]".query[List[OtherEnum]]
67+
(for {
68+
_ <- q.analysis
69+
.map(ana => assertEquals(ana.columnAlignmentErrors, List.empty))
70+
71+
_ <- q.unique.map(assertEquals(_, listOfOtherEnums))
72+
73+
_ <- sql"select array['a', 'b']".query[List[OtherEnum]].analysis.map(_.columnAlignmentErrors)
74+
.map {
75+
case List(e: ColumnTypeError) => assertEquals(e.schema.vendorTypeName, "_text")
76+
case other => fail(s"Unexpected typecheck result: $other")
77+
}
78+
79+
_ <- sql"select array['a', 'b'] :: other_schema.other_enum[]".query[List[String]].analysis.map(
80+
_.columnAlignmentErrors)
81+
.map {
82+
case List(e: ColumnTypeError) => assertEquals(e.schema.vendorTypeName, """"other_schema"."_other_enum"""")
83+
case other => fail(s"Unexpected typecheck result: $other")
84+
}
85+
} yield ())
86+
.transact(transactor)
87+
}
88+
89+
test("array of custom type in another schema: writes correctly and typechecks") {
90+
val q = sql"insert into temp_otherenum (arr) values ($listOfOtherEnums)".update
91+
(for {
92+
_ <- sql"drop table if exists temp_otherenum".update.run
93+
_ <- sql"create table temp_otherenum(arr other_schema.other_enum[] not null)".update.run
94+
_ <- q.analysis.map(_.parameterAlignmentErrors).map(ana => assertEquals(ana, List.empty))
95+
_ <- q.run
96+
_ <- sql"select arr from temp_otherenum".query[List[OtherEnum]].to[List]
97+
.map(assertEquals(_, List(listOfOtherEnums)))
98+
99+
_ <- sql"insert into temp_otherenum (arr) values (${List("a")})".update.analysis
100+
.map(_.parameterAlignmentErrors)
101+
.map {
102+
case List(e: ParameterTypeError) => {
103+
// pgjdbc is a bit crazy. If you have inserted into the table already then it'll report the parameter type as
104+
// _other_enum, or otherwise "other_schema"."_other_enum"..
105+
assertEquals(e.vendorTypeName, "_other_enum")
106+
// assertEquals(e.vendorTypeName, s""""other_schema"."_other_enum"""")
107+
}
108+
case other => fail(s"Unexpected typecheck result: $other")
109+
}
110+
} yield ())
111+
.transact(transactor)
112+
}
113+
114+
}

‎modules/postgres/src/test/scala/doobie/postgres/enums/MyEnum.scala‎

Lines changed: 25 additions & 9 deletions
Original file line numberDiff line numberDiff line change
@@ -6,22 +6,38 @@ package doobie.postgres.enums
66

77
import doobie.Meta
88
import doobie.postgres.implicits.*
9+
import doobie.postgres.implicits.arrayOfEnum
910

1011
// create type myenum as enum ('foo', 'bar') <-- part of setup
1112
sealed trait MyEnum
1213
object MyEnum {
1314
case object Foo extends MyEnum
1415
case object Bar extends MyEnum
1516

17+
def fromStringUnsafe(s: String): MyEnum = s match {
18+
case "foo" => Foo
19+
case "bar" => Bar
20+
}
21+
22+
def asString(e: MyEnum): String = e match {
23+
case Foo => "foo"
24+
case Bar => "bar"
25+
}
26+
27+
private val typeName = "myenum"
28+
1629
implicit val MyEnumMeta: Meta[MyEnum] =
1730
pgEnumString(
18-
"myenum",
19-
{
20-
case "foo" => Foo
21-
case "bar" => Bar
22-
},
23-
{
24-
case Foo => "foo"
25-
case Bar => "bar"
26-
})
31+
typeName,
32+
fromStringUnsafe,
33+
asString
34+
)
35+
36+
implicit val MyEnumArrayMeta: Meta[Array[MyEnum]] =
37+
arrayOfEnum[MyEnum](
38+
typeName,
39+
fromStringUnsafe,
40+
asString
41+
)
42+
2743
}

‎modules/postgres/src/test/scala/doobie/postgres/enums/OtherEnum.scala‎

Lines changed: 32 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,32 @@
1+
// Copyright (c) 2013-2020 Rob Norris and Contributors
2+
// This software is licensed under the MIT License (MIT).
3+
// For more information see LICENSE or https://opensource.org/licenses/MIT
4+
5+
package doobie.postgres.enums
6+
7+
import doobie.Meta
8+
9+
// This is an enum type defined in another schema (See other_enum in test-db.sql)
10+
sealed abstract class OtherEnum(val strValue: String)
11+
12+
object OtherEnum {
13+
case object A extends OtherEnum("a")
14+
15+
case object B extends OtherEnum("b")
16+
17+
private def fromStrUnsafe(s: String): OtherEnum = s match {
18+
case "a" => A
19+
case "b" => B
20+
}
21+
22+
private val elementTypeNameUnqualified = "other_enum"
23+
private val elementTypeName = s""""other_schema"."$elementTypeNameUnqualified""""
24+
private val arrayTypeName = s""""other_schema"."_$elementTypeNameUnqualified""""
25+
26+
implicit val arrayMeta: Meta[Array[OtherEnum]] =
27+
Meta.Advanced.array[String](
28+
elementTypeName,
29+
arrayTypeName,
30+
s"_$elementTypeNameUnqualified"
31+
).timap(arr => arr.map(fromStrUnsafe))(arr => arr.map(_.strValue))
32+
}

0 commit comments

Comments
 (0)