use count_bytes_reader from plugin-sdk-go (#92553)
* use count_bytes_reader from plugin-sdk-go * run `make update-workspace` * update postgres tests * update mysql tests * time back to utc * make update-workspace done --------- Co-authored-by: Kyle Brandt <kyle@grafana.com>
This commit is contained in:
co-authored by
Kyle Brandt
parent
09f102b72e
commit
abca0380a8
@@ -1,39 +0,0 @@
|
||||
package httpclient
|
||||
|
||||
import (
|
||||
"io"
|
||||
)
|
||||
|
||||
type CloseCallbackFunc func(bytesRead int64)
|
||||
|
||||
// CountBytesReader counts the total amount of bytes read from the underlying reader.
|
||||
//
|
||||
// The provided callback func will be called before the underlying reader is closed.
|
||||
func CountBytesReader(reader io.ReadCloser, callback CloseCallbackFunc) io.ReadCloser {
|
||||
if reader == nil {
|
||||
panic("reader cannot be nil")
|
||||
}
|
||||
|
||||
if callback == nil {
|
||||
panic("callback cannot be nil")
|
||||
}
|
||||
|
||||
return &countBytesReader{reader: reader, callback: callback}
|
||||
}
|
||||
|
||||
type countBytesReader struct {
|
||||
reader io.ReadCloser
|
||||
callback CloseCallbackFunc
|
||||
counter int64
|
||||
}
|
||||
|
||||
func (r *countBytesReader) Read(p []byte) (int, error) {
|
||||
n, err := r.reader.Read(p)
|
||||
r.counter += int64(n)
|
||||
return n, err
|
||||
}
|
||||
|
||||
func (r *countBytesReader) Close() error {
|
||||
r.callback(r.counter)
|
||||
return r.reader.Close()
|
||||
}
|
||||
@@ -1,38 +0,0 @@
|
||||
package httpclient
|
||||
|
||||
import (
|
||||
"fmt"
|
||||
"io"
|
||||
"strings"
|
||||
"testing"
|
||||
|
||||
"github.com/stretchr/testify/require"
|
||||
)
|
||||
|
||||
func TestCountBytesReader(t *testing.T) {
|
||||
tcs := []struct {
|
||||
body string
|
||||
expectedBytesCount int64
|
||||
}{
|
||||
{body: "d", expectedBytesCount: 1},
|
||||
{body: "dummy", expectedBytesCount: 5},
|
||||
}
|
||||
|
||||
for index, tc := range tcs {
|
||||
t.Run(fmt.Sprintf("Test CountBytesReader %d", index), func(t *testing.T) {
|
||||
body := io.NopCloser(strings.NewReader(tc.body))
|
||||
var actualBytesRead int64
|
||||
|
||||
readCloser := CountBytesReader(body, func(bytesRead int64) {
|
||||
actualBytesRead = bytesRead
|
||||
})
|
||||
|
||||
bodyBytes, err := io.ReadAll(readCloser)
|
||||
require.NoError(t, err)
|
||||
err = readCloser.Close()
|
||||
require.NoError(t, err)
|
||||
require.Equal(t, tc.expectedBytesCount, actualBytesRead)
|
||||
require.Equal(t, string(bodyBytes), tc.body)
|
||||
})
|
||||
}
|
||||
}
|
||||
@@ -6,7 +6,6 @@ import (
|
||||
"time"
|
||||
|
||||
sdkhttpclient "github.com/grafana/grafana-plugin-sdk-go/backend/httpclient"
|
||||
"github.com/grafana/grafana/pkg/infra/httpclient"
|
||||
"github.com/grafana/grafana/pkg/infra/metrics/metricutil"
|
||||
"github.com/prometheus/client_golang/prometheus"
|
||||
"github.com/prometheus/client_golang/prometheus/promauto"
|
||||
@@ -113,7 +112,7 @@ func executeMiddleware(next http.RoundTripper, labels prometheus.Labels) http.Ro
|
||||
}
|
||||
|
||||
if res != nil && res.StatusCode != http.StatusSwitchingProtocols {
|
||||
res.Body = httpclient.CountBytesReader(res.Body, func(bytesRead int64) {
|
||||
res.Body = sdkhttpclient.CountBytesReader(res.Body, func(bytesRead int64) {
|
||||
responseSizeHistogram.Observe(float64(bytesRead))
|
||||
})
|
||||
}
|
||||
|
||||
Reference in New Issue
Block a user