Provisioning: Move apifmt, loki and safepath to provisioning app (#110226)

* Move apifmt * Move safepath * Move Loki package * Regenerate Loki mock * Missing file for Loki
2025-08-27 13:26:48 -05:00
parent e78f6b6b37
commit 93a35fc7be
42 changed files with 32 additions and 31 deletions
@@ -0,0 +1,185 @@
+package loki
+
+import (
+	"bytes"
+	"context"
+	"encoding/json"
+	"fmt"
+	"io"
+	"net/http"
+	"net/url"
+	"strconv"
+	"time"
+
+	"github.com/grafana/grafana-app-sdk/logging"
+)
+
+type Config struct {
+	ReadPathURL       *url.URL
+	WritePathURL      *url.URL
+	BasicAuthUser     string
+	BasicAuthPassword string
+	TenantID          string
+	ExternalLabels    map[string]string
+	MaxQuerySize      int
+}
+
+type Stream struct {
+	Stream map[string]string `json:"stream"`
+	Values []Sample          `json:"values"`
+}
+
+type Sample struct {
+	T time.Time
+	V string
+}
+
+func (r Sample) MarshalJSON() ([]byte, error) {
+	return json.Marshal([2]string{
+		fmt.Sprintf("%d", r.T.UnixNano()), r.V,
+	})
+}
+
+func (r *Sample) UnmarshalJSON(b []byte) error {
+	var tuple [2]string
+	if err := json.Unmarshal(b, &tuple); err != nil {
+		return fmt.Errorf("failed to deserialize sample in Loki response: %w", err)
+	}
+	nano, err := strconv.ParseInt(tuple[0], 10, 64)
+	if err != nil {
+		return fmt.Errorf("timestamp in Loki sample not convertible to nanosecond epoch: %v", tuple[0])
+	}
+	r.T = time.Unix(0, nano)
+	r.V = tuple[1]
+	return nil
+}
+
+type QueryRes struct {
+	Data QueryData `json:"data"`
+}
+
+type QueryData struct {
+	Result []Stream `json:"result"`
+}
+
+type PushRequest struct {
+	Streams []Stream `json:"streams"`
+}
+
+type Client struct {
+	cfg    Config
+	client *http.Client
+}
+
+func NewClient(cfg Config) *Client {
+	return &Client{
+		cfg:    cfg,
+		client: &http.Client{Timeout: 30 * time.Second},
+	}
+}
+
+func (c *Client) Push(ctx context.Context, streams []Stream) error {
+	log := logging.FromContext(ctx)
+
+	pushReq := PushRequest{Streams: streams}
+	body, err := json.Marshal(pushReq)
+	if err != nil {
+		return fmt.Errorf("failed to marshal push request: %w", err)
+	}
+
+	uri := c.cfg.WritePathURL.JoinPath("/loki/api/v1/push")
+	req, err := http.NewRequest(http.MethodPost, uri.String(), bytes.NewBuffer(body))
+	if err != nil {
+		return fmt.Errorf("failed to create Loki request: %w", err)
+	}
+
+	c.setAuthAndTenantHeaders(req)
+	req.Header.Set("Content-Type", "application/json")
+
+	req = req.WithContext(ctx)
+	res, err := c.client.Do(req)
+	if res != nil {
+		defer func() {
+			if err := res.Body.Close(); err != nil {
+				log.Warn("Failed to close response body", "err", err)
+			}
+		}()
+	}
+	if err != nil {
+		return fmt.Errorf("error sending request: %w", err)
+	}
+
+	if res.StatusCode < 200 || res.StatusCode >= 300 {
+		body, _ := io.ReadAll(res.Body)
+		log.Error("Error response from Loki", "response", string(body), "status", res.StatusCode)
+		return fmt.Errorf("received a non-200 response from loki, status: %d", res.StatusCode)
+	}
+
+	log.Debug("Successfully pushed streams to Loki", "status", res.StatusCode, "streams", len(streams))
+	return nil
+}
+
+func (c *Client) RangeQuery(ctx context.Context, logQL string, start, end, limit int64) (QueryRes, error) {
+	log := logging.FromContext(ctx)
+
+	uri := c.cfg.ReadPathURL.JoinPath("/loki/api/v1/query_range")
+	req, err := http.NewRequest(http.MethodGet, uri.String(), nil)
+	if err != nil {
+		return QueryRes{}, fmt.Errorf("error creating request: %w", err)
+	}
+
+	q := req.URL.Query()
+	q.Set("query", logQL)
+	q.Set("start", strconv.FormatInt(start, 10))
+	q.Set("end", strconv.FormatInt(end, 10))
+	if limit > 0 {
+		q.Set("limit", strconv.FormatInt(limit, 10))
+	}
+	req.URL.RawQuery = q.Encode()
+
+	c.setAuthAndTenantHeaders(req)
+	req = req.WithContext(ctx)
+
+	res, err := c.client.Do(req)
+	if res != nil {
+		defer func() {
+			if err := res.Body.Close(); err != nil {
+				log.Warn("Failed to close response body", "err", err)
+			}
+		}()
+	}
+	if err != nil {
+		return QueryRes{}, fmt.Errorf("error sending request: %w", err)
+	}
+
+	body, err := io.ReadAll(res.Body)
+	if err != nil {
+		return QueryRes{}, fmt.Errorf("error reading request response: %w", err)
+	}
+
+	if res.StatusCode < 200 || res.StatusCode >= 300 {
+		if len(body) > 0 {
+			log.Error("Error response from Loki", "response", string(body), "status", res.StatusCode)
+		} else {
+			log.Error("Error response from Loki with an empty body", "status", res.StatusCode)
+		}
+		return QueryRes{}, fmt.Errorf("received a non-200 response from loki, status: %d", res.StatusCode)
+	}
+
+	var queryRes QueryRes
+	if err := json.Unmarshal(body, &queryRes); err != nil {
+		return QueryRes{}, fmt.Errorf("error unmarshaling loki response: %w", err)
+	}
+
+	log.Debug("Successfully queried Loki", "status", res.StatusCode, "streams", len(queryRes.Data.Result))
+	return queryRes, nil
+}
+
+func (c *Client) setAuthAndTenantHeaders(req *http.Request) {
+	if c.cfg.BasicAuthUser != "" || c.cfg.BasicAuthPassword != "" {
+		req.SetBasicAuth(c.cfg.BasicAuthUser, c.cfg.BasicAuthPassword)
+	}
+	if c.cfg.TenantID != "" {
+		req.Header.Set("X-Scope-OrgID", c.cfg.TenantID)
+	}
+}
@@ -0,0 +1,314 @@
+package loki
+
+import (
+	"context"
+	"encoding/json"
+	"fmt"
+	"net/http"
+	"net/http/httptest"
+	"net/url"
+	"testing"
+	"time"
+
+	"github.com/stretchr/testify/assert"
+	"github.com/stretchr/testify/require"
+)
+
+func TestSample_MarshalJSON(t *testing.T) {
+	sample := Sample{
+		T: time.Unix(0, 1234567890000000000), // 1234567890 seconds in nanoseconds
+		V: "test log line",
+	}
+
+	data, err := json.Marshal(sample)
+	require.NoError(t, err)
+
+	expected := `["1234567890000000000","test log line"]`
+	assert.JSONEq(t, expected, string(data))
+}
+
+func TestSample_UnmarshalJSON(t *testing.T) {
+	t.Run("valid sample", func(t *testing.T) {
+		data := `["1234567890000000000","test log line"]`
+		var sample Sample
+
+		err := json.Unmarshal([]byte(data), &sample)
+		require.NoError(t, err)
+
+		assert.Equal(t, time.Unix(0, 1234567890000000000), sample.T)
+		assert.Equal(t, "test log line", sample.V)
+	})
+
+	t.Run("invalid format", func(t *testing.T) {
+		data := `"invalid"`
+		var sample Sample
+
+		err := json.Unmarshal([]byte(data), &sample)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "failed to deserialize sample")
+	})
+
+	t.Run("invalid timestamp", func(t *testing.T) {
+		data := `["not-a-number","test log line"]`
+		var sample Sample
+
+		err := json.Unmarshal([]byte(data), &sample)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "timestamp in Loki sample not convertible")
+	})
+}
+
+func TestClient_Push(t *testing.T) {
+	t.Run("successful push", func(t *testing.T) {
+		var receivedBody PushRequest
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			assert.Equal(t, "/loki/api/v1/push", r.URL.Path)
+			assert.Equal(t, http.MethodPost, r.Method)
+			assert.Equal(t, "application/json", r.Header.Get("Content-Type"))
+
+			err := json.NewDecoder(r.Body).Decode(&receivedBody)
+			require.NoError(t, err)
+
+			w.WriteHeader(http.StatusOK)
+		}))
+		defer server.Close()
+
+		client := createTestClient(t, server.URL, server.URL)
+
+		streams := []Stream{
+			{
+				Stream: map[string]string{"job": "test"},
+				Values: []Sample{
+					{T: time.Unix(0, 1234567890000000000), V: "log line 1"},
+					{T: time.Unix(0, 1234567891000000000), V: "log line 2"},
+				},
+			},
+		}
+
+		err := client.Push(context.Background(), streams)
+		assert.NoError(t, err)
+
+		// Verify the request body
+		assert.Len(t, receivedBody.Streams, 1)
+		assert.Equal(t, "test", receivedBody.Streams[0].Stream["job"])
+		assert.Len(t, receivedBody.Streams[0].Values, 2)
+	})
+
+	t.Run("push failure", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusBadRequest)
+			_, _ = w.Write([]byte("Bad request"))
+		}))
+		defer server.Close()
+
+		client := createTestClient(t, server.URL, server.URL)
+
+		streams := []Stream{{Stream: map[string]string{"job": "test"}}}
+		err := client.Push(context.Background(), streams)
+
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "non-200 response")
+	})
+}
+
+func TestClient_RangeQuery(t *testing.T) {
+	t.Run("successful query", func(t *testing.T) {
+		expectedResponse := QueryRes{
+			Data: QueryData{
+				Result: []Stream{
+					{
+						Stream: map[string]string{"job": "test"},
+						Values: []Sample{
+							{T: time.Unix(0, 1234567890000000000), V: "log line 1"},
+							{T: time.Unix(0, 1234567891000000000), V: "log line 2"},
+						},
+					},
+				},
+			},
+		}
+
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			assert.Equal(t, "/loki/api/v1/query_range", r.URL.Path)
+			assert.Equal(t, http.MethodGet, r.Method)
+
+			// Check query parameters
+			params := r.URL.Query()
+			assert.Equal(t, `{job="test"}`, params.Get("query"))
+			assert.Equal(t, "1000000000", params.Get("start"))
+			assert.Equal(t, "2000000000", params.Get("end"))
+			assert.Equal(t, "100", params.Get("limit"))
+
+			w.Header().Set("Content-Type", "application/json")
+			_ = json.NewEncoder(w).Encode(expectedResponse)
+		}))
+		defer server.Close()
+
+		client := createTestClient(t, server.URL, server.URL)
+
+		result, err := client.RangeQuery(
+			context.Background(),
+			`{job="test"}`,
+			1000000000, // start
+			2000000000, // end
+			100,        // limit
+		)
+
+		assert.NoError(t, err)
+		assert.Len(t, result.Data.Result, 1)
+		assert.Equal(t, "test", result.Data.Result[0].Stream["job"])
+		assert.Len(t, result.Data.Result[0].Values, 2)
+	})
+
+	t.Run("query without limit", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			params := r.URL.Query()
+			assert.Equal(t, "", params.Get("limit")) // Should not be set
+
+			w.Header().Set("Content-Type", "application/json")
+			_ = json.NewEncoder(w).Encode(QueryRes{})
+		}))
+		defer server.Close()
+
+		client := createTestClient(t, server.URL, server.URL)
+
+		_, err := client.RangeQuery(context.Background(), `{job="test"}`, 1000000000, 2000000000, 0)
+		assert.NoError(t, err)
+	})
+
+	t.Run("query failure", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.WriteHeader(http.StatusBadRequest)
+			_, _ = w.Write([]byte("Bad query"))
+		}))
+		defer server.Close()
+
+		client := createTestClient(t, server.URL, server.URL)
+
+		_, err := client.RangeQuery(context.Background(), `{job="test"}`, 1000000000, 2000000000, 100)
+
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "non-200 response")
+	})
+
+	t.Run("invalid JSON response", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			w.Header().Set("Content-Type", "application/json")
+			_, _ = w.Write([]byte("invalid json"))
+		}))
+		defer server.Close()
+
+		client := createTestClient(t, server.URL, server.URL)
+
+		_, err := client.RangeQuery(context.Background(), `{job="test"}`, 1000000000, 2000000000, 100)
+
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "error unmarshaling loki response")
+	})
+}
+
+func TestClient_setAuthAndTenantHeaders(t *testing.T) {
+	t.Run("with basic auth and tenant", func(t *testing.T) {
+		cfg := createTestConfig(t, "http://localhost", "http://localhost")
+		cfg.BasicAuthUser = "testuser"
+		cfg.BasicAuthPassword = "testpass"
+		cfg.TenantID = "test-tenant"
+
+		client := NewClient(cfg)
+
+		req, _ := http.NewRequest(http.MethodGet, "http://localhost", nil)
+		client.setAuthAndTenantHeaders(req)
+
+		username, password, ok := req.BasicAuth()
+		assert.True(t, ok)
+		assert.Equal(t, "testuser", username)
+		assert.Equal(t, "testpass", password)
+		assert.Equal(t, "test-tenant", req.Header.Get("X-Scope-OrgID"))
+	})
+
+	t.Run("without auth", func(t *testing.T) {
+		cfg := createTestConfig(t, "http://localhost", "http://localhost")
+		client := NewClient(cfg)
+
+		req, _ := http.NewRequest(http.MethodGet, "http://localhost", nil)
+		client.setAuthAndTenantHeaders(req)
+
+		_, _, ok := req.BasicAuth()
+		assert.False(t, ok)
+		assert.Equal(t, "", req.Header.Get("X-Scope-OrgID"))
+	})
+}
+
+func TestStream_JSONRoundtrip(t *testing.T) {
+	original := Stream{
+		Stream: map[string]string{
+			"job":       "test-job",
+			"instance":  "test-instance",
+			"namespace": "test-ns",
+		},
+		Values: []Sample{
+			{T: time.Unix(0, 1234567890000000000), V: "log line 1"},
+			{T: time.Unix(0, 1234567891000000000), V: "log line 2"},
+			{T: time.Unix(0, 1234567892000000000), V: "log line 3"},
+		},
+	}
+
+	// Marshal to JSON
+	data, err := json.Marshal(original)
+	require.NoError(t, err)
+
+	// Unmarshal back
+	var restored Stream
+	err = json.Unmarshal(data, &restored)
+	require.NoError(t, err)
+
+	// Verify all fields match
+	assert.Equal(t, original.Stream, restored.Stream)
+	assert.Len(t, restored.Values, len(original.Values))
+
+	for i, sample := range original.Values {
+		assert.True(t, sample.T.Equal(restored.Values[i].T),
+			fmt.Sprintf("Timestamp mismatch at index %d: expected %v, got %v", i, sample.T, restored.Values[i].T))
+		assert.Equal(t, sample.V, restored.Values[i].V)
+	}
+}
+
+// Helper functions
+
+func createTestClient(t *testing.T, readURL, writeURL string) *Client {
+	cfg := createTestConfig(t, readURL, writeURL)
+	return NewClient(cfg)
+}
+
+func createTestConfig(t *testing.T, readURL, writeURL string) Config {
+	readParsed, err := url.Parse(readURL)
+	require.NoError(t, err)
+
+	writeParsed, err := url.Parse(writeURL)
+	require.NoError(t, err)
+
+	return Config{
+		ReadPathURL:    readParsed,
+		WritePathURL:   writeParsed,
+		ExternalLabels: map[string]string{"source": "test"},
+		MaxQuerySize:   1000,
+	}
+}
+
+func TestClient_ContextCancellation(t *testing.T) {
+	t.Run("push with cancelled context", func(t *testing.T) {
+		server := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, r *http.Request) {
+			t.Error("Handler should not be called with cancelled context")
+		}))
+		defer server.Close()
+
+		client := createTestClient(t, server.URL, server.URL)
+
+		ctx, cancel := context.WithCancel(context.Background())
+		cancel()
+
+		streams := []Stream{{Stream: map[string]string{"job": "test"}}}
+		err := client.Push(ctx, streams)
+		assert.Error(t, err)
+		assert.Contains(t, err.Error(), "context canceled")
+	})
+}