fix(pool): extended pipeline mixed messages for multi-shard

integration/go/go_pgx/pipeline_sharded_test.go

@@ -0,0 +1,126 @@
+package main
+
+import (
+	"context"
+	"testing"
+
+	"github.com/jackc/pgx/v5"
+	"github.com/jackc/pgx/v5/pgconn"
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+const (
+	textOID   = uint32(25)
+	bigintOID = uint32(20)
+)
+
+func fieldNames(description *pgconn.StatementDescription) []string {
+	names := []string{}
+	for _, field := range description.Fields {
+		names = append(names, field.Name)
+	}
+	return names
+}
+
+// pgx folds ParseComplete, ParameterDescription and RowDescription into one
+// StatementDescription, so dropping any of the three desynchronises the read.
+func nextStatement(t *testing.T, pipeline *pgconn.Pipeline) *pgconn.StatementDescription {
+	t.Helper()
+
+	result, err := pipeline.GetResults()
+	require.NoError(t, err)
+
+	description, ok := result.(*pgconn.StatementDescription)
+	require.Truef(t, ok, "expected a statement description, got %T", result)
+	return description
+}
+
+// A pipelined prepare phase describes several statements before one Sync, and
+// pgx reads the replies positionally. The aggregate and the ORDER BY send the
+// request to every shard, which is what puts it through the multi-shard path.
+func TestPipelinedPrepareDescribesEveryStatement(t *testing.T) {
+	ctx := context.Background()
+
+	conn, err := pgx.Connect(ctx, testConnStr)
+	require.NoError(t, err)
+	defer conn.Close(ctx)
+
+	pipeline := conn.PgConn().StartPipeline(ctx)
+
+	pipeline.SendPrepare("pipe1", "SELECT count(*) FROM sharded", nil)
+	pipeline.SendPrepare("pipe2", "SELECT id, $1::text AS second FROM sharded ORDER BY id", nil)
+	pipeline.SendPrepare(
+		"pipe3",
+		"SELECT id, $1::text AS second, $2::bigint AS third FROM sharded ORDER BY id",
+		nil,
+	)
+	// Binary results, as every pgx query asks for.
+	pipeline.SendQueryPrepared("pipe1", nil, nil, []int16{1})
+	require.NoError(t, pipeline.Sync())
+
+	first := nextStatement(t, pipeline)
+	assert.Empty(t, first.ParamOIDs)
+	assert.Equal(t, []string{"count"}, fieldNames(first))
+
+	second := nextStatement(t, pipeline)
+	assert.Equal(t, []uint32{textOID}, second.ParamOIDs)
+	assert.Equal(t, []string{"id", "second"}, fieldNames(second))
+
+	third := nextStatement(t, pipeline)
+	assert.Equal(t, []uint32{textOID, bigintOID}, third.ParamOIDs)
+	assert.Equal(t, []string{"id", "second", "third"}, fieldNames(third))
+
+	result, err := pipeline.GetResults()
+	require.NoError(t, err)
+	reader, ok := result.(*pgconn.ResultReader)
+	require.Truef(t, ok, "expected a result reader, got %T", result)
+	require.NoError(t, reader.Read().Err)
+
+	result, err = pipeline.GetResults()
+	require.NoError(t, err)
+	require.IsType(t, &pgconn.PipelineSync{}, result)
+
+	require.NoError(t, pipeline.Close())
+}
+
+// pgx asks for binary results, and pgdog decodes the rows itself to merge them.
+// Read as text these ids sort with 10 before 2.
+func TestShardedOrderByMergesBinaryRows(t *testing.T) {
+	ctx := context.Background()
+
+	conn, err := pgx.Connect(ctx, testConnStr)
+	require.NoError(t, err)
+	defer conn.Close(ctx)
+
+	_, err = conn.Exec(ctx, "TRUNCATE TABLE sharded")
+	require.NoError(t, err)
+
+	want := []int64{}
+	for id := int64(1); id <= 20; id++ {
+		_, err = conn.Exec(ctx, "INSERT INTO sharded (id, value) VALUES ($1, $2)", id, "row")
+		require.NoError(t, err)
+		want = append(want, id)
+	}
+
+	// Nothing is merged unless both shards hold rows.
+	assert.NotZero(t, countOnShardByComment(t, conn, 0, 1), "shard 0 holds no rows")
+	assert.NotZero(t, countOnShardByComment(t, conn, 1, 11), "shard 1 holds no rows")
+
+	rows, err := conn.Query(ctx, "SELECT id FROM sharded ORDER BY id")
+	require.NoError(t, err)
+
+	got := []int64{}
+	for rows.Next() {
+		var id int64
+		require.NoError(t, rows.Scan(&id))
+		got = append(got, id)
+	}
+	require.NoError(t, rows.Err())
+	rows.Close()
+
+	assert.Equal(t, want, got, "rows must arrive merge-sorted across shards")
+
+	_, err = conn.Exec(ctx, "TRUNCATE TABLE sharded")
+	require.NoError(t, err)
+}

pgdog/src/backend/pool/connection/aggregate.rs

@@ -515,7 +515,7 @@ mod test {
         let aggregate = parse("SELECT COUNT(*)::int FROM users");
 
         let rd = RowDescription::new(&[integer_field("count")]);
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         let mut rows = VecDeque::new();
         let mut shard0 = DataRow::new();
@@ -544,7 +544,7 @@ mod test {
         let aggregate = parse("SELECT AVG(price) FROM menu");
 
         let rd = RowDescription::new(&[Field::double("avg")]);
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         let mut rows = VecDeque::new();
         let mut shard0 = DataRow::new();
@@ -578,7 +578,7 @@ mod test {
         let aggregate = parse("SELECT price, SUM(quantity) FROM menu GROUP BY 1");
 
         let rd = RowDescription::new(&[Field::double("price"), Field::bigint("sum")]);
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         let mut rows = VecDeque::new();
         let mut shard0 = DataRow::new();
@@ -620,7 +620,7 @@ mod test {
         let aggregate = parse("SELECT matrix, COUNT(*) FROM samples GROUP BY 1");
 
         let rd = RowDescription::new(&[integer_array_field("matrix"), Field::bigint("count")]);
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         let mut rows = VecDeque::new();
 
@@ -670,7 +670,7 @@ mod test {
             interval_array_field("sample_interval_array"),
             Field::bigint("count"),
         ]);
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         let input = Bytes::from_static(br#"{"1 year 2 mons 1 day 04:05:06.7"}"#);
 

pgdog/src/backend/pool/connection/binding.rs

@@ -299,7 +299,9 @@ impl Binding {
         match self {
             Binding::Admin(admin) => !admin.done(),
             Binding::Direct(server, ..) => server.has_more_messages(),
-            Binding::MultiShard(servers, _state) => servers.iter().any(|s| s.has_more_messages()),
+            Binding::MultiShard(servers, state) => {
+                state.has_more_messages() || servers.iter().any(|s| s.has_more_messages())
+            }
             _ => false,
         }
     }

pgdog/src/backend/pool/connection/buffer.rs

@@ -216,6 +216,11 @@ impl Buffer {
         }
     }
 
+    /// The buffer holds rows the client can read now.
+    pub(super) fn can_take(&self) -> bool {
+        self.full && !self.is_empty()
+    }
+
     /// Execute LIMIT ... OFFSET ...
     pub(super) fn limit(&mut self, limit: &Limit) {
         let offset = limit.offset.unwrap_or(0);
@@ -230,7 +235,6 @@ impl Buffer {
         self.buffer.len()
     }
 
-    #[allow(dead_code)]
     pub(super) fn is_empty(&self) -> bool {
         self.len() == 0
     }
@@ -254,7 +258,7 @@ mod test {
             buf.add(dr.message().unwrap()).unwrap();
         }
 
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         buf.sort(&columns, &decoder);
         buf.full();
@@ -284,7 +288,7 @@ mod test {
             buf.add(dr.message().unwrap()).unwrap();
         }
 
-        buf.aggregate(&agg, &Decoder::from(&rd), &AggregateRewritePlan::default())
+        buf.aggregate(&agg, &Decoder::from(rd), &AggregateRewritePlan::default())
             .unwrap();
         buf.full();
 
@@ -311,7 +315,7 @@ mod test {
             }
         }
 
-        buf.aggregate(&agg, &Decoder::from(&rd), &AggregateRewritePlan::default())
+        buf.aggregate(&agg, &Decoder::from(rd), &AggregateRewritePlan::default())
             .unwrap();
         buf.full();
 
@@ -345,7 +349,7 @@ mod test {
             buf.add(dr.message().unwrap()).unwrap();
         }
 
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         buf.sort(&columns, &decoder);
         buf.full();
@@ -385,7 +389,7 @@ mod test {
             buf.add(dr.message().unwrap()).unwrap();
         }
 
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         buf.sort(&columns, &decoder);
         buf.full();
@@ -441,7 +445,7 @@ mod test {
             buf.add(dr.message().unwrap()).unwrap();
         }
 
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
         buf.sort(&columns, &decoder);
         buf.full();
 
@@ -485,7 +489,7 @@ mod test {
             buf.add(dr.message().unwrap()).unwrap();
         }
 
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
         buf.sort(&columns, &decoder);
         buf.full();
 
@@ -577,7 +581,7 @@ mod test {
     fn test_distinct() {
         let mut buf = Buffer::default();
         let rd = RowDescription::new(&[Field::bigint("id"), Field::text("email")]);
-        let decoder = Decoder::from(&rd);
+        let decoder = Decoder::from(rd);
 
         for email in ["test@test.com", "apples@test.com", "domain@test.com"] {
             for i in 0..5 {

pgdog/src/backend/pool/connection/mod.rs

@@ -414,7 +414,7 @@ impl Connection {
     pub(crate) fn bind(&mut self, bind: &Bind) -> Result<(), Error> {
         match self.binding {
             Binding::MultiShard(_, ref mut state) => {
-                state.set_context(bind);
+                state.set_bind_context(bind);
                 Ok(())
             }