Merge pull request #2128 from satorg/in-not-in-fragments

Proposal: `IN` and `NOT IN` fragment builders for product types with arbitrary arities
typelevel · Nov 7, 2024 · cc635ff · cc635ff
2 parents 1f3c036 + 313e929
commit cc635ff
Show file tree

Hide file tree

Showing 2 changed files with 183 additions and 49 deletions.
diff --git a/modules/core/src/main/scala/doobie/util/fragments.scala b/modules/core/src/main/scala/doobie/util/fragments.scala
@@ -5,11 +5,12 @@
 package doobie
 package util
 
+import cats.Foldable
+import cats.Reducible
 import cats.data.NonEmptyList
 import cats.syntax.all.*
-import cats.{Foldable, Functor, Reducible}
-import doobie.implicits.*
 import doobie.Fragment.*
+import doobie.implicits.*
 
 /** Module of `Fragment` constructors. */
 object fragments {
@@ -27,33 +28,81 @@ object fragments {
   }
 
   /** Returns `(f IN (fs0, fs1, ...))`. */
-  def in[A: util.Put](f: Fragment, fs0: A, fs1: A, fs: A*): Fragment =
+  def in[A: util.Write](f: Fragment, fs0: A, fs1: A, fs: A*): Fragment =
     in(f, NonEmptyList(fs0, fs1 :: fs.toList))
 
-  /** Returns `(f IN (fs0, fs1, ...))`. */
-  def in[F[_]: Reducible: Functor, A: util.Put](f: Fragment, fs: F[A]): Fragment =
-    parentheses(f ++ fr" IN" ++ parentheses(comma(fs.map(a => fr"$a"))))
-
-  def inOpt[F[_]: Foldable, A: util.Put](f: Fragment, fs: F[A]): Option[Fragment] =
-    NonEmptyList.fromFoldable(fs).map(nel => in(f, nel))
-
-  /** Returns `(f IN ((fs0-A, fs0-B), (fs1-A, fs1-B), ...))`. */
-  def in[F[_]: Reducible: Functor, A: util.Put, B: util.Put](f: Fragment, fs: F[(A, B)]): Fragment =
-    parentheses(f ++ fr" IN" ++ parentheses(comma(fs.map { case (a, b) => fr0"($a,$b)" })))
-
   /** Returns `(f NOT IN (fs0, fs1, ...))`. */
-  def notIn[A: util.Put](f: Fragment, fs0: A, fs1: A, fs: A*): Fragment =
+  def notIn[A: util.Write](f: Fragment, fs0: A, fs1: A, fs: A*): Fragment =
     notIn(f, NonEmptyList(fs0, fs1 :: fs.toList))
 
-  /** Returns `(f NOT IN (fs0, fs1, ...))`. */
-  def notIn[F[_]: Reducible: Functor, A: util.Put](f: Fragment, fs: F[A]): Fragment = {
-    parentheses(f ++ fr" NOT IN" ++ parentheses(comma(fs.map(a => fr"$a"))))
+  @inline
+  private def mkRowFn[A](implicit A: util.Write[A]): A => Fragment =
+    if (A.length == 1) // no need for extra parentheses
+      a => values(a)
+    else
+      a => parentheses0(values(a))
+
+  @inline
+  private def constSubqueryExpr[F[_]: Reducible, A: util.Write](fs: F[A]): Fragment = {
+    val row = mkRowFn[A]
+    parentheses0(fs.reduceLeftTo(row) { _ ++ fr"," ++ row(_) })
   }
 
-  def notInOpt[F[_]: Foldable, A: util.Put](f: Fragment, fs: F[A]): Option[Fragment] = {
-    NonEmptyList.fromFoldable(fs).map(nel => notIn(f, nel))
+  @inline
+  private def constSubqueryExprOpt[F[_]: Foldable, A: util.Write](fs: F[A]): Option[Fragment] = {
+    val row = mkRowFn[A]
+    fs.reduceLeftToOption(row) { _ ++ fr"," ++ row(_) }
+      .map(parentheses0)
   }
 
+  /** Returns `f IN (fs0, fs1, ...)`.
+    *
+    * @param f
+    *   left-hand expression.
+    * @param fs
+    *   values of `Product` type to compare to the left-hand expression.
+    * @return
+    *   the result `IN` expression.
+    */
+  def in[F[_]: Reducible, A: util.Write](f: Fragment, fs: F[A]): Fragment =
+    parentheses(f ++ fr" IN" ++ constSubqueryExpr(fs))
+
+  /** Returns `f IN (fs0, fs1, ...)`.
+    *
+    * @param f
+    *   left-hand expression.
+    * @param fs
+    *   values of `Product` type to compare to the left-hand expression.
+    * @return
+    *   the result `IN` expression enclosed in `Some` or `None` if `fs` is empty.
+    */
+  def inOpt[F[_]: Foldable, A: util.Write](f: Fragment, fs: F[A]): Option[Fragment] =
+    constSubqueryExprOpt(fs).map(expr => parentheses(f ++ fr" IN" ++ expr))
+
+  /** Returns `f NOT IN (fs0, fs1, ...)`.
+    *
+    * @param f
+    *   left-hand expression.
+    * @param fs
+    *   values of `Product` type to compare to the left-hand expression.
+    * @return
+    *   the result `NOT IN` subquery expression.
+    */
+  def notIn[F[_]: Reducible, A: util.Write](f: Fragment, fs: F[A]): Fragment =
+    parentheses(f ++ fr" NOT IN" ++ constSubqueryExpr(fs))
+
+  /** Returns `f NOT IN (fs0, fs1, ...)`.
+    *
+    * @param f
+    *   left-hand expression.
+    * @param fs
+    *   values of `Product` type to compare to the left-hand expression.
+    * @return
+    *   the result `NOT IN` subquery expression enclosed in `Some` or `None` if `fs` is empty.
+    */
+  def notInOpt[F[_]: Foldable, A: util.Write](f: Fragment, fs: F[A]): Option[Fragment] =
+    constSubqueryExprOpt(fs).map(expr => parentheses(f ++ fr" NOT IN" ++ expr))
+
   /** Returns `(f1 AND f2 AND ... fn)`. */
   def and(f1: Fragment, f2: Fragment, fs: Fragment*): Fragment =
     and(NonEmptyList(f1, f2 :: fs.toList))

diff --git a/modules/core/src/test/scala/doobie/util/FragmentsSuite.scala b/modules/core/src/test/scala/doobie/util/FragmentsSuite.scala
@@ -5,8 +5,9 @@
 package doobie.util
 
 import cats.data.NonEmptyList
-import doobie.*, doobie.implicits.*
 import cats.effect.IO
+import doobie.*
+import doobie.implicits.*
 
 class FragmentsSuite extends munit.FunSuite {
   import Fragments.*
@@ -21,7 +22,11 @@ class FragmentsSuite extends munit.FunSuite {
   )
 
   val nelInt = NonEmptyList.of(1, 2, 3)
-  val listInt = nelInt.toList
+  val nelIntBool2 = NonEmptyList.of((1, true), (2, false))
+  val nelStrDblInt3 = NonEmptyList.of(("abc", 1.2, 3), ("def", 4.5, 6), ("ghi", 7.8, 9))
+  val listInt3 = nelInt.toList
+  val listIntBool2 = nelIntBool2.toList
+  val listStrDblInt3 = nelStrDblInt3.toList
   val nel1 = NonEmptyList.of(1).map(i => sql"$i")
   val nel = NonEmptyList.of(1, 2, 3).map(i => sql"$i")
   val fs = nel.toList
@@ -48,48 +53,128 @@ class FragmentsSuite extends munit.FunSuite {
     assertEquals(updateSetOpt(Fragment.const("Foo"), List.empty[Fragment]).map(_.query[Unit].sql), None)
   }
 
-  test("in (1-column varargs)") {
-    assertEquals(in(sql"foo", 1, 2, 3).query[Unit].sql, "(foo IN (? , ? , ? ) ) ")
+  test("in (1-column, varargs)") {
+    assertEquals(
+      in(sql"foo.bar", 1, 2).query[Unit].sql,
+      "(foo.bar IN (?, ?)) ")
+    assertEquals(
+      in(sql"foo.bar", 3, 4, 5).query[Unit].sql,
+      "(foo.bar IN (?, ?, ?)) ")
   }
-
-  test("in (1-column Reducible many)") {
-    assertEquals(in(sql"foo", nelInt).query[Unit].sql, "(foo IN (? , ? , ? ) ) ")
+  test("in (2-column, varargs)") {
+    assertEquals(
+      in(sql"foo.bar", 1 -> "one", 2 -> "two").query[Unit].sql,
+      "(foo.bar IN ((?,?), (?,?))) ")
+    assertEquals(
+      in(sql"foo.bar", 3 -> "three", 4 -> "four", 5 -> "five").query[Unit].sql,
+      "(foo.bar IN ((?,?), (?,?), (?,?))) ")
   }
-
-  test("inOpt (1-column Reducible empty)") {
-    assertEquals(inOpt(sql"foo", List.empty[Int]).map(_.query[Unit].sql), None)
+  test("in (3-column, varargs)") {
+    assertEquals(
+      in(sql"foo.bar", (1.2, "A", 3), (4.5, "B", 6)).query[Unit].sql,
+      "(foo.bar IN ((?,?,?), (?,?,?))) ")
+    assertEquals(
+      in(sql"foo.bar", (9.8, "Z", 7), (6.5, "Y", 4), (3.2, "X", 1)).query[Unit].sql,
+      "(foo.bar IN ((?,?,?), (?,?,?), (?,?,?))) ")
   }
 
-  test("inOpt (1-column Reducible many)") {
-    assertEquals(inOpt(sql"foo", listInt).map(_.query[Unit].sql), Some("(foo IN (? , ? , ? ) ) "))
+  test("notIn (1-column, varargs)") {
+    assertEquals(
+      notIn(sql"foo.bar", 1, 2).query[Unit].sql,
+      "(foo.bar NOT IN (?, ?)) ")
+    assertEquals(
+      notIn(sql"foo.bar", 3, 4, 5).query[Unit].sql,
+      "(foo.bar NOT IN (?, ?, ?)) ")
   }
-
-  test("in (2-column varargs)") {
-    assertEquals(in(sql"foo", NonEmptyList.of((1, true), (2, false))).query[Unit].sql, "(foo IN ((?,?), (?,?)) ) ")
+  test("notIn (2-column, varargs)") {
+    assertEquals(
+      notIn(sql"foo.bar", 1 -> "one", 2 -> "two").query[Unit].sql,
+      "(foo.bar NOT IN ((?,?), (?,?))) ")
+    assertEquals(
+      notIn(sql"foo.bar", 3 -> "three", 4 -> "four", 5 -> "five").query[Unit].sql,
+      "(foo.bar NOT IN ((?,?), (?,?), (?,?))) ")
   }
-
-  test("notIn (varargs many)") {
-    assertEquals(notIn(sql"foo", 1, 2, 3).query[Unit].sql, "(foo NOT IN (? , ? , ? ) ) ")
+  test("notIn (3-column, varargs)") {
+    assertEquals(
+      notIn(sql"foo.bar", (1.2, "A", 3), (4.5, "B", 6)).query[Unit].sql,
+      "(foo.bar NOT IN ((?,?,?), (?,?,?))) ")
+    assertEquals(
+      notIn(sql"foo.bar", (9.8, "Z", 7), (6.5, "Y", 4), (3.2, "X", 1)).query[Unit].sql,
+      "(foo.bar NOT IN ((?,?,?), (?,?,?), (?,?,?))) ")
   }
 
-  test("notIn (Reducible 1)") {
-    assertEquals(notIn(sql"foo", NonEmptyList.of(1)).query[Unit].sql, "(foo NOT IN (? ) ) ")
+  test("in (1-column)") {
+    assertEquals(
+      in(fr0"foo.bar", nelInt).query[Unit].sql,
+      "(foo.bar IN (?, ?, ?)) ")
   }
-
-  test("notIn (Reducible many)") {
-    assertEquals(notIn(sql"foo", nelInt).query[Unit].sql, "(foo NOT IN (? , ? , ? ) ) ")
+  test("in (2-columns)") {
+    assertEquals(
+      in(fr0"foo.bar", nelIntBool2).query[Unit].sql,
+      "(foo.bar IN ((?,?), (?,?))) ")
+  }
+  test("in (3-columns)") {
+    assertEquals(
+      in(fr0"foo.bar", nelStrDblInt3).query[Unit].sql,
+      "(foo.bar IN ((?,?,?), (?,?,?), (?,?,?))) ")
   }
 
-  test("notInOpt (Foldable empty)") {
-    assertEquals(notInOpt(sql"foo", List.empty[Int]).map(_.query[Unit].sql), None)
+  test("inOpt (1-column, many)") {
+    assertEquals(
+      inOpt(fr0"foo.bar", listInt3).map(_.query[Unit].sql),
+      Some("(foo.bar IN (?, ?, ?)) "))
+  }
+  test("inOpt (2-columns, many)") {
+    assertEquals(
+      inOpt(fr0"foo.bar", listIntBool2).map(_.query[Unit].sql),
+      Some("(foo.bar IN ((?,?), (?,?))) "))
+  }
+  test("inOpt (3-columns, many)") {
+    assertEquals(
+      inOpt(fr0"foo.bar", listStrDblInt3).map(_.query[Unit].sql),
+      Some("(foo.bar IN ((?,?,?), (?,?,?), (?,?,?))) "))
+  }
+  test("inOpt (empty)") {
+    assert(inOpt(fr0"foo.bar", List.empty[Int]).isEmpty, "1 column")
+    assert(inOpt(fr0"foo.bar", List.empty[(Int, Boolean)]).isEmpty, "2 columns")
+    assert(inOpt(fr0"foo.bar", List.empty[(String, Double, Int)]).isEmpty, "3 columns")
   }
 
-  test("notInOpt (Foldable 1)") {
-    assertEquals(notInOpt(sql"foo", List(1)).map(_.query[Unit].sql), Some("(foo NOT IN (? ) ) "))
+  test("notIn (1-column)") {
+    assertEquals(
+      notIn(fr0"foo.bar", nelInt).query[Unit].sql,
+      "(foo.bar NOT IN (?, ?, ?)) ")
+  }
+  test("notIn (2-columns)") {
+    assertEquals(
+      notIn(fr0"foo.bar", nelIntBool2).query[Unit].sql,
+      "(foo.bar NOT IN ((?,?), (?,?))) ")
+  }
+  test("notIn (3-columns)") {
+    assertEquals(
+      notIn(fr0"foo.bar", nelStrDblInt3).query[Unit].sql,
+      "(foo.bar NOT IN ((?,?,?), (?,?,?), (?,?,?))) ")
   }
 
-  test("notInOpt (Foldable many)") {
-    assertEquals(notInOpt(sql"foo", listInt).map(_.query[Unit].sql), Some("(foo NOT IN (? , ? , ? ) ) "))
+  test("notInOpt (1-column, many)") {
+    assertEquals(
+      notInOpt(fr0"foo.bar", listInt3).map(_.query[Unit].sql),
+      Some("(foo.bar NOT IN (?, ?, ?)) "))
+  }
+  test("notInOpt (2-columns, many)") {
+    assertEquals(
+      notInOpt(fr0"foo.bar", listIntBool2).map(_.query[Unit].sql),
+      Some("(foo.bar NOT IN ((?,?), (?,?))) "))
+  }
+  test("notInOpt (3-columns, many)") {
+    assertEquals(
+      notInOpt(fr0"foo.bar", listStrDblInt3).map(_.query[Unit].sql),
+      Some("(foo.bar NOT IN ((?,?,?), (?,?,?), (?,?,?))) "))
+  }
+  test("notInOpt (empty)") {
+    assert(notInOpt(fr0"foo.bar", List.empty[Int]).isEmpty, "1 column")
+    assert(notInOpt(fr0"foo.bar", List.empty[(Int, Boolean)]).isEmpty, "2 columns")
+    assert(notInOpt(fr0"foo.bar", List.empty[(String, Double, Int)]).isEmpty, "3 columns")
   }
 
   test("and (vararg 2)") {