diff --git a/query.go b/query.go index f99e3af6..5460e7ab 100644 --- a/query.go +++ b/query.go @@ -343,6 +343,19 @@ func (c *Client) sendInput(ctx context.Context, info proto.ColInfoInput, q Query if err := f(ctx); err != nil { if errors.Is(err, io.EOF) { // No more data. + if tailRows := q.Input[0].Data.Rows(); tailRows > 0 { + // Write data tail on next tick and break. + // + // This is required to resemble io.Reader behavior. + if ce := c.lg.Check(zap.DebugLevel, "Writing tail of input data (not empty and io.EOF)"); ce != nil { + ce.Write( + zap.Int("rows", tailRows), + ) + } + f = nil + continue + } + break } // ClickHouse server persists blocks after receive. diff --git a/query_test.go b/query_test.go index 8e9d5199..031f53ab 100644 --- a/query_test.go +++ b/query_test.go @@ -245,7 +245,7 @@ func TestClient_Query(t *testing.T) { OnInput: func(ctx context.Context) error { data = append(data[:0], uint8(total), 2, 3, 4) total++ - if total > blocks { + if total >= blocks { return io.EOF } return nil