Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
23 changes: 23 additions & 0 deletions docker-compose-testing.yml
Original file line number Diff line number Diff line change
Expand Up @@ -433,6 +433,29 @@ services:
--host 0.0.0.0 \
--port \
1199
otlp_mockserver:
networks:
testing-network:
aliases:
- otlp.test
image: stackql/testlib
build:
context: .
dockerfile: testLib.Dockerfile
ports:
- "1200:1200/tcp"
environment:
IS_DOCKER: "${IS_DOCKER:-false}"
command:
- bash
- -c
- |
flask \
--app=/opt/testlib/test/python/stackql_test_tooling/flask/otlp/app \
run \
--host 0.0.0.0 \
--port \
1200
native_test_mockserver:
networks:
testing-network:
Expand Down
13 changes: 13 additions & 0 deletions internal/stackql/cmd/root.go
Original file line number Diff line number Diff line change
Expand Up @@ -29,6 +29,7 @@ import (
"github.com/stackql/stackql/internal/stackql/config"
"github.com/stackql/stackql/internal/stackql/envfile"
"github.com/stackql/stackql/internal/stackql/intrinsic"
"github.com/stackql/stackql/internal/stackql/output"

"github.com/magiconair/properties"
"github.com/spf13/cobra"
Expand Down Expand Up @@ -77,6 +78,12 @@ var (
//nolint:gochecknoglobals // cobra binds flags to package scope
var previewCfgRaw string

// otelConfigRaw is the raw --otel.config argument; initConfig hands it to
// the output package once.
//
//nolint:gochecknoglobals // cobra binds flags to package scope
var otelConfigRaw string

// rootCmd represents the base command when called without any subcommands.
//
//nolint:gochecknoglobals // global vars are a pattern for this lib
Expand Down Expand Up @@ -177,6 +184,8 @@ func init() {
rootCmd.PersistentFlags().StringVar(&runtimeCtx.ACIDCfgRaw, dto.ACIDCfgRawKey, "{}", "JSON / YAML string representing ACID config")
rootCmd.PersistentFlags().StringVar(&previewCfgRaw, intrinsic.CfgRawKey, "{}", "JSON string configuring the "+intrinsic.ProviderName+
" provider backend; keys: batchSize, flushInterval, endpoint, unstable")
rootCmd.PersistentFlags().StringVar(&otelConfigRaw, output.OTelConfigKey, "{}", "JSON / YAML string configuring the otel output format, "+
`eg: '{ "exporter": { "endpoint": "http://localhost:4318/v1/logs", "headers": { "authorization": "<token>" } } }'`)
rootCmd.PersistentFlags().StringVar(&runtimeCtx.SessionCtxRaw, dto.SessionCtxKey, "{}", "JSON / YAML string representing session config")
rootCmd.PersistentFlags().IntVar(&runtimeCtx.APIRequestTimeout, dto.APIRequestTimeoutKey, 45, "API request timeout in seconds, 0 for no timeout.") //nolint:mnd // TODO: investigate
rootCmd.PersistentFlags().StringVar(&dummyString, dto.ColorSchemeKey, "", "DEPRECATED: color schems no longer active")
Expand Down Expand Up @@ -273,6 +282,10 @@ func initConfig() {
mergeConfigFromFile(&runtimeCtx, *rootCmd.PersistentFlags())

intrinsic.Init(previewCfgRaw)
if err := output.InitOTelConfig(otelConfigRaw); err != nil {
fmt.Fprintf(os.Stderr, "invalid --%s: %v\n", output.OTelConfigKey, err)
os.Exit(1)
}

// An absent --env.file is created empty (issue #691) so packaged installs
// have a credential store to populate; creation failure is non-fatal
Expand Down
74 changes: 70 additions & 4 deletions internal/stackql/output/otel.go
Original file line number Diff line number Diff line change
Expand Up @@ -19,8 +19,42 @@ import (
"github.com/stackql/stackql-parser/go/vt/sqlparser"
"github.com/stackql/stackql/internal/stackql/buildinfo"
"github.com/stackql/stackql/pkg/sink"
"gopkg.in/yaml.v2"
)

// OTelConfigKey is the flag carrying the OTelConfig JSON / YAML.
const OTelConfigKey = "otel.config"

// OTelConfig is the --otel.config DTO: settings for the otel output format.
type OTelConfig struct {
// Exporter, when set, also pushes every statement's records to the
// OTLP/HTTP logs endpoint it names.
Exporter *sink.OTLPHTTPConfig `json:"exporter,omitempty" yaml:"exporter,omitempty"`
}

// otelConfig is the parsed --otel.config, set once at startup.
//
//nolint:gochecknoglobals // process-wide configuration
var otelConfig OTelConfig

// InitOTelConfig parses the --otel.config argument; an empty value keeps the
// defaults and an exporter without an endpoint is rejected.
func InitOTelConfig(raw string) error {
var cfg OTelConfig
if strings.TrimSpace(raw) != "" {
if err := yaml.Unmarshal([]byte(raw), &cfg); err != nil {
return err
}
}
if cfg.Exporter != nil {
if _, err := sink.NewOTLPHTTPSink(*cfg.Exporter); err != nil {
return err
}
}
otelConfig = cfg
return nil
}

// The emitted attribute set is a versioned interface (issue #738): bump
// otelAttributeSchemaVersion when it changes and update the schema test.
const (
Expand Down Expand Up @@ -51,7 +85,8 @@ var otelInvocationTraceID = sync.OnceValue(func() string {
// OTelWriter emits one OTLP/JSON LogsData per result row plus a completion
// record per statement, so a result set can be shipped as a timestamped
// inventory snapshot. Every record of a statement shares the snapshot
// instant, snapshot id and span id.
// instant, snapshot id and span id. When --otel.config names an exporter
// the statement is also pushed there as a batch.
type OTelWriter struct {
errWriter io.Writer
encoder sink.Sink
Expand Down Expand Up @@ -103,6 +138,20 @@ func (s writerSink) Record(_ context.Context, payload any) error {

func (writerSink) Close() error { return nil }

// newOTelOutputSink streams to the output writer and, when --otel.config
// names an exporter, also pushes to it.
func newOTelOutputSink(writer io.Writer) sink.Sink {
out := writerSink{writer: writer}
if otelConfig.Exporter == nil {
return out
}
exporter, err := sink.NewOTLPHTTPSink(*otelConfig.Exporter)
if err != nil {
return out
}
return sink.NewMultiSink(out, exporter)
}

// NewOTelWriter returns the writer for one statement: query is the submitted
// text and startTime the snapshot instant stamped on every record.
func NewOTelWriter(writer io.Writer, errWriter io.Writer, query string, startTime time.Time) IOutputWriter {
Expand All @@ -113,7 +162,7 @@ func NewOTelWriter(writer io.Writer, errWriter io.Writer, query string, startTim
queryHash := sha256.Sum256([]byte(strings.TrimSpace(query)))
return &OTelWriter{
errWriter: errWriter,
encoder: sink.NewOTelSink(writerSink{writer: writer}, otelResource(statement), otelScope()),
encoder: sink.NewOTelSink(newOTelOutputSink(writer), otelResource(statement), otelScope()),
startTime: startTime,
snapshotID: uuid.NewString(),
spanID: sink.OTelRandomID(otelSpanIDBytes),
Expand Down Expand Up @@ -280,6 +329,17 @@ func (ow *OTelWriter) writeCompletion(rows int) error {
return ow.record(sink.OTelSeverityInfo, body, attrs)
}

// close flushes the sinks so a configured exporter ships the statement's
// batch; the failure is echoed to the error stream because the response
// handler discards writer errors.
func (ow *OTelWriter) close() error {
if err := ow.encoder.Close(); err != nil {
_ = writeStderrError(ow.errWriter, err)
return err
}
return nil
}

func (ow *OTelWriter) Write(res sqldata.ISQLResultStream) error {
rows := 0
for {
Expand All @@ -296,7 +356,10 @@ func (ow *OTelWriter) Write(res sqldata.ISQLResultStream) error {
}
}
if err != nil {
return ow.writeCompletion(rows)
if completionErr := ow.writeCompletion(rows); completionErr != nil {
return completionErr
}
return ow.close()
}
}
}
Expand All @@ -309,5 +372,8 @@ func (ow *OTelWriter) WriteError(err error, errorPresentation string) error {
sink.OTelString("error.type", fmt.Sprintf("%T", err)),
sink.OTelString("error.message", err.Error()),
)
return ow.record(sink.OTelSeverityError, err.Error(), attrs)
if recordErr := ow.record(sink.OTelSeverityError, err.Error(), attrs); recordErr != nil {
return recordErr
}
return ow.close()
}
98 changes: 98 additions & 0 deletions internal/stackql/output/otel_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -4,8 +4,11 @@ import (
"bytes"
"encoding/json"
"errors"
"net/http"
"net/http/httptest"
"reflect"
"strings"
"sync"
"testing"
"time"

Expand Down Expand Up @@ -236,6 +239,101 @@ func TestOTelWriter_WriteError(t *testing.T) {
}
}

func TestInitOTelConfig(t *testing.T) {
t.Cleanup(func() { otelConfig = OTelConfig{} })
for _, raw := range []string{"", " ", "{}"} {
if err := InitOTelConfig(raw); err != nil || otelConfig.Exporter != nil {
t.Fatalf("InitOTelConfig(%q): err=%v exporter=%+v", raw, err, otelConfig.Exporter)
}
}
yamlForm := "exporter:\n endpoint: http://c:4318/v1/logs\n headers:\n api-key: secret\n batch_size: 3\n"
if err := InitOTelConfig(yamlForm); err != nil || otelConfig.Exporter == nil ||
otelConfig.Exporter.Endpoint != "http://c:4318/v1/logs" || otelConfig.Exporter.BatchSize != 3 ||
otelConfig.Exporter.Headers["api-key"] != "secret" {
t.Fatalf("yaml form: err=%v exporter=%+v", err, otelConfig.Exporter)
}
// A rejected value leaves the previous configuration in place.
for _, raw := range []string{`{"exporter": {"headers": {"a": "b"}}}`, `{"exporter": [1]}`, `not: [valid`} {
if err := InitOTelConfig(raw); err == nil || otelConfig.Exporter == nil || otelConfig.Exporter.BatchSize != 3 {
t.Fatalf("InitOTelConfig(%q): err=%v exporter=%+v", raw, err, otelConfig.Exporter)
}
}
}

// TestOTelWriter_PushesStatementBatchToOTLPEndpoint covers issue #755: with
// an exporter in --otel.config, the records streamed to stdout are also
// POSTed to the endpoint, chunked by batch_size.
func TestOTelWriter_PushesStatementBatchToOTLPEndpoint(t *testing.T) {
var (
mu sync.Mutex
requests []*http.Request
bodies []sink.OTelLogsData
)
srv := httptest.NewServer(http.HandlerFunc(func(_ http.ResponseWriter, r *http.Request) {
var data sink.OTelLogsData
if err := json.NewDecoder(r.Body).Decode(&data); err != nil {
t.Errorf("body is not LogsData: %v", err)
}
mu.Lock()
defer mu.Unlock()
requests = append(requests, r)
bodies = append(bodies, data)
}))
defer srv.Close()
t.Cleanup(func() { otelConfig = OTelConfig{} })
if err := InitOTelConfig(`{"exporter": {"endpoint": "` + srv.URL + `/v1/logs", ` +
`"headers": {"api-key": "secret one", "x-tenant": "acme"}, "batch_size": 2}}`); err != nil {
t.Fatalf("InitOTelConfig() error = %v", err)
}
lines := otelSchemaFixture(t)
if len(requests) != 2 {
t.Fatalf("expected 2 export requests (2 rows + completion at batch size 2), got %d", len(requests))
}
first := requests[0]
if first.URL.Path != "/v1/logs" || first.Header.Get("Content-Type") != "application/json" ||
first.Header.Get("api-key") != "secret one" || first.Header.Get("x-tenant") != "acme" {
t.Fatalf("request: path %q headers %v", first.URL.Path, first.Header)
}
var pushed []sink.OTelWireLogRecord
for i, body := range bodies {
if len(body.ResourceLogs) != 1 || len(body.ResourceLogs[0].ScopeLogs) != 1 {
t.Fatalf("batch %d: expected one shared resource and scope envelope, got %+v", i, body)
}
pushed = append(pushed, body.ResourceLogs[0].ScopeLogs[0].LogRecords...)
}
if len(bodies[0].ResourceLogs[0].ScopeLogs[0].LogRecords) != 2 || len(pushed) != len(lines) {
t.Fatalf("expected batches of 2 + 1 records, got %d pushed", len(pushed))
}
for i, rl := range lines {
if !reflect.DeepEqual(otelRecord(rl), pushed[i]) {
t.Fatalf("pushed record %d differs from the streamed one:\n got %+v\nwant %+v", i, pushed[i], otelRecord(rl))
}
}
if !reflect.DeepEqual(bodies[0].ResourceLogs[0].Resource, lines[0].Resource) {
t.Fatalf("pushed resource differs from the streamed one: %+v", bodies[0].ResourceLogs[0].Resource)
}
}

// TestOTelWriter_ReportsExportFailureOnStderr: a push that keeps failing is
// surfaced on the error stream after the rows were streamed.
func TestOTelWriter_ReportsExportFailureOnStderr(t *testing.T) {
srv := httptest.NewServer(http.HandlerFunc(func(w http.ResponseWriter, _ *http.Request) {
w.WriteHeader(http.StatusBadRequest)
}))
defer srv.Close()
t.Cleanup(func() { otelConfig = OTelConfig{} })
otelConfig = OTelConfig{Exporter: &sink.OTLPHTTPConfig{Endpoint: srv.URL + "/v1/logs"}}
var out, errOut bytes.Buffer
w := NewOTelWriter(&out, &errOut, otelTestQuery, time.Time{})
err := w.Write(otelTestStream([]interface{}{"fw-a", "us-east1", false, nil}))
if err == nil || !strings.Contains(errOut.String(), "otlp http sink") {
t.Fatalf("expected the export failure on stderr, err=%v stderr=%q", err, errOut.String())
}
if len(otelLines(t, out.String())) != 2 {
t.Fatalf("rows must still stream to stdout: %q", out.String())
}
}

func TestGetOutputWriter_OTel(t *testing.T) {
ctx := internaldto.OutputContext{RuntimeContext: dto.RuntimeCtx{OutputFormat: "otel"}, Query: "select 1"}
w, err := GetOutputWriter(&bytes.Buffer{}, &bytes.Buffer{}, ctx)
Expand Down
2 changes: 1 addition & 1 deletion pkg/mcp_server/policy/policy.go
Original file line number Diff line number Diff line change
Expand Up @@ -73,7 +73,7 @@ func ClassifyQuery(sql string) QueryClass {
verb = trimmed
}
switch strings.ToUpper(verb) {
case "SELECT", "SHOW", "DESCRIBE", "EXPLAIN":
case "SELECT", "SHOW", "DESCRIBE", "DESC", "EXPLAIN":
return QueryClassSelect
case "WITH":
return classifyWith(trimmed)
Expand Down
2 changes: 2 additions & 0 deletions pkg/mcp_server/policy/policy_test.go
Original file line number Diff line number Diff line change
Expand Up @@ -17,6 +17,8 @@ func TestClassifyQuery(t *testing.T) {
{" SELECT 1", policy.QueryClassSelect},
{"Show methods in x.y.z;", policy.QueryClassSelect},
{"describe x.y", policy.QueryClassSelect},
{"desc x.y", policy.QueryClassSelect}, // issue #773: DESC aliases DESCRIBE
{"DESC EXTENDED x.y.z", policy.QueryClassSelect},
{"EXPLAIN select 1", policy.QueryClassSelect},
{"insert into t values (1)", policy.QueryClassMutationCreate},
{"UPDATE t SET a=1", policy.QueryClassMutationCreate},
Expand Down
Loading
Loading