UseSpec.hs

module Sharing.ByUnit.Connection.UseSpec (spec) where

import Data.Either
import Hasql.Connection qualified as Connection
import Hasql.Decoders qualified as Decoders
import Hasql.Encoders qualified as Encoders
import Hasql.Pipeline qualified as Pipeline
import Hasql.Session qualified as Session
import Hasql.Statement qualified as Statement
import Helpers.Dsls.Execution qualified as Execution
import Helpers.Scripts qualified as Scripts
import Helpers.Statements.SelectOne qualified as Statements.SelectOne
import Helpers.Statements.Sleep qualified as Statements
import Test.Hspec
import Prelude

spec :: SpecWith (Text, Word16)
spec = do
  describe "Transactions" do
    it "Do not cause \"in progress after error\"" \config -> do
      Scripts.onPreparableConnection config \connection -> do
        let sumStatement =
              Statement.preparable
                "select ($1 + $2)"
                ( mconcat
                    [ fst >$< Encoders.param (Encoders.nonNullable Encoders.int8),
                      snd >$< Encoders.param (Encoders.nonNullable Encoders.int8)
                    ]
                )
                (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int8)))

        result <-
          Connection.use connection do
            Session.script "."

        result `shouldSatisfy` isLeft

        result <-
          Connection.use connection do
            Session.script "begin;"
            s <- Session.statement (1 :: Int64, 2 :: Int64) sumStatement
            Session.script "end;"
            return s

        result `shouldBe` Right (3 :: Int64)

  describe "Pipeline Mode" do
    it "Leaves the connection usable after timeout in pipeline" \config -> do
      Scripts.onPreparableConnection config \connection -> do
        let selectStatement =
              Statement.preparable
                "select $1::int"
                (Encoders.param (Encoders.nonNullable Encoders.int4))
                (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int4)))

        -- Timeout during a pipeline operation
        result <-
          timeout 50_000 do
            Connection.use connection
              $ Session.pipeline
              $ (,)
              <$> Pipeline.statement 42 selectStatement
              <*> Execution.pipelineByParams (Statements.Sleep 0.1)

        result `shouldBe` Nothing

        -- Try to use pipeline again after timeout cleanup
        -- This should work but fails with "connection not idle" without the fix
        result2 <-
          Connection.use connection
            $ Session.pipeline
            $ Pipeline.statement 99 selectStatement

        result2 `shouldBe` Right 99

  describe "Timing out" do
    describe "On a statement" do
      it "Leaves the connection usable" \config -> Scripts.onPreparableConnection config \connection -> do
        result <-
          timeout 50_000 do
            Connection.use connection do
              Execution.sessionByParams (Statements.Sleep 0.1)

        result `shouldBe` Nothing

        result <-
          Connection.use connection do
            Execution.sessionByParams Statements.SelectOne.SelectOne

        result `shouldBe` Right 1

    describe "On a transaction" do
      it "Leaves the connection usable" \config -> Scripts.onPreparableConnection config \connection -> do
        -- Start a transaction and timeout during it
        result <-
          timeout 50_000 do
            Connection.use connection do
              Session.script "begin;"
              Execution.sessionByParams (Statements.Sleep 0.1)
              Session.script "commit;"

        result `shouldBe` Nothing

        -- Connection should still be usable after timeout in transaction
        result <-
          Connection.use connection do
            Execution.sessionByParams Statements.SelectOne.SelectOne

        result `shouldBe` Right 1

      it "Lets us start another transaction" do
        let checkTransactionStatus =
              Statement.preparable
                "select case when pg_advisory_lock(1) is null then 0 else 1 end"
                mempty
                (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int4)))
         in \config -> Scripts.onPreparableConnection config \connection -> do
              -- Timeout during a transaction
              result <-
                timeout 50_000 do
                  Connection.use connection do
                    Session.script "begin;"
                    Execution.sessionByParams (Statements.Sleep 0.1)

              result `shouldBe` Nothing

              -- Verify we can start a new transaction without "already in progress" error
              result <-
                Connection.use connection do
                  Session.script "begin;"
                  s <- Session.statement () checkTransactionStatus
                  Session.script "commit;"
                  return s

              result `shouldBe` Right 1

      it "Does not corrupt the prepared statement registry" do
        let returnIntStatement =
              Statement.preparable
                "select $1::int"
                (Encoders.param (Encoders.nonNullable Encoders.int4))
                (Decoders.singleRow (Decoders.column (Decoders.nonNullable Decoders.int4)))
         in \config -> Scripts.onPreparableConnection config \connection -> do
              -- Use a prepared statement first
              result <-
                Connection.use connection do
                  Session.statement 42 returnIntStatement

              result `shouldBe` Right 42

              -- Timeout during transaction (causes connection reset)
              result <-
                timeout 50_000 do
                  Connection.use connection do
                    Session.script "begin;"
                    Execution.sessionByParams (Statements.Sleep 0.1)
                    Session.script "commit;"

              result `shouldBe` Nothing

              -- The prepared statement should work again without "does not exist" error
              result <-
                Connection.use connection do
                  Session.statement 99 returnIntStatement

              result `shouldBe` Right 99