Postgres COPY fails past minChunkSize when using non-Pure effects

[dleblanc]

We're seeing that the copy operation gets closed once we have a stream that's larger than the minChunkSize for the .copyIn() method. We only see this when using effects other than Pure.

This is true for both 1.0.0-M1 and 1.0.0-M5.

Failing test code follows:

package io.citrine.tbl.ingest

import cats.effect.IO
import cats.implicits.catsSyntaxApplicativeId
import doobie.ConnectionIO
import doobie.implicits._
import doobie.postgres.implicits._
import doobie.util.transactor.Transactor
import doobie.util.transactor.Transactor.Aux
import munit.CatsEffectSuite
import org.postgresql.ds.PGSimpleDataSource

import javax.sql.DataSource

class Fs2CopyTest extends CatsEffectSuite {

  val minChunkSize = 200

  val datasource: DataSource = {
    val ds = new PGSimpleDataSource
    ds.setURL(System.getenv("DB_URL"))
    ds.setUser(System.getenv("DB_USER"))
    ds.setPassword(System.getenv("DB_PASSWORD"))
    ds
  }
  val xa: Aux[IO, DataSource] = Transactor.fromDataSource[IO](datasource, scala.concurrent.ExecutionContext.global)

  test("A stream with a Pure effect inserts items properly") {

    sql"""
        DROP TABLE IF EXISTS demo;
        CREATE TABLE demo(id BIGSERIAL PRIMARY KEY NOT NULL, data BIGINT NOT NULL);
         """.update.run
      .transact(xa)
      .unsafeRunSync()

    // A pure stream is fine - can copy many items
    val count = 10000
    val stream = fs2.Stream.emits(1 to count)

    (sql"COPY demo(data) FROM STDIN").copyIn(stream, minChunkSize).transact(xa).unsafeRunSync()

    val queryCount =
      sql"SELECT count(*) from demo".query[Int].unique.transact(xa).unsafeRunSync()

    assertEquals(queryCount, count)
  }

  test("A stream with a ConnectionIO effect copies <= than minChunkSize items") {

    sql"""
        DROP TABLE IF EXISTS demo;
        CREATE TABLE demo(id BIGSERIAL PRIMARY KEY NOT NULL, data BIGINT NOT NULL);
         """.update.run
      .transact(xa)
      .unsafeRunSync()

    // Can copy up to minChunkSize just fine with ConnectionIO
    val inputs = (1 to minChunkSize)
    val stream = fs2.Stream.emits[ConnectionIO, Int](inputs)
      .evalMap(i => (i + 2).pure[ConnectionIO])

    (sql"COPY demo(data) FROM STDIN").copyIn(stream, minChunkSize).transact(xa).unsafeRunSync()

    val queryCount =
      sql"SELECT count(*) from demo".query[Int].unique.transact(xa).unsafeRunSync()

    assertEquals(queryCount, minChunkSize)
  }

  test("A stream with a ConnectionIO effect fails to copy items with count > minChunkSize") {

    sql"""
        DROP TABLE IF EXISTS demo;
        CREATE TABLE demo(id BIGSERIAL PRIMARY KEY NOT NULL, data BIGINT NOT NULL);
         """.update.run
      .transact(xa)
      .unsafeRunSync()

    // Can't copy over minChunkSize with ConnectionIO - copy operation gets closed
    val inputs = (1 to minChunkSize + 1)
    val stream = fs2.Stream.emits[ConnectionIO, Int](inputs)
      .evalMap(i => (i + 2).pure[ConnectionIO])

    (sql"COPY demo(data) FROM STDIN").copyIn(stream, minChunkSize).transact(xa).unsafeRunSync()

    val queryCount =
      sql"SELECT count(*) from demo".query[Int].unique.transact(xa).unsafeRunSync()

    assertEquals(queryCount, minChunkSize + 1)
  }
}