executor_test.go

package query_test

import (
	"errors"
	"fmt"
	"strings"
	"sync"
	"testing"
	"time"

	"github.com/influxdata/influxdb/query"
	"github.com/influxdata/influxdb/services/meta"
	"github.com/influxdata/influxql"
	"github.com/stretchr/testify/require"
)

var errUnexpected = errors.New("unexpected error")

type StatementExecutor struct {
	ExecuteStatementFn func(stmt influxql.Statement, ctx *query.ExecutionContext) error
}

func (e *StatementExecutor) ExecuteStatement(ctx *query.ExecutionContext, stmt influxql.Statement) error {
	return e.ExecuteStatementFn(stmt, ctx)
}

type StatementNormalizerExecutor struct {
	StatementExecutor
	NormalizeStatementFn func(stmt influxql.Statement, database, retentionPolicy string) error
}

func (e *StatementNormalizerExecutor) NormalizeStatement(stmt influxql.Statement, database, retentionPolicy string) error {
	return e.NormalizeStatementFn(stmt, database, retentionPolicy)
}

func NewQueryExecutor() *query.Executor {
	return query.NewExecutor()
}

func TestQueryExecutor_AttachQuery(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			if ctx.QueryID != 1 {
				t.Errorf("incorrect query id: exp=1 got=%d", ctx.QueryID)
			}
			return nil
		},
	}

	discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
}

func TestQueryExecutor_KillQuery(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	qid := make(chan uint64)

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			switch stmt.(type) {
			case *influxql.KillQueryStatement:
				return e.TaskManager.ExecuteStatement(ctx, stmt)
			}

			qid <- ctx.QueryID
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(100 * time.Millisecond):
				t.Error("killing the query did not close the channel after 100 milliseconds")
				return errUnexpected
			}
		},
	}

	results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
	q, err = influxql.ParseQuery(fmt.Sprintf("KILL QUERY %d", <-qid))
	if err != nil {
		t.Fatal(err)
	}
	discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))

	result := <-results
	if result.Err != query.ErrQueryInterrupted {
		t.Errorf("unexpected error: %s", result.Err)
	}
}

func TestQueryExecutor_KillQuery_Zombie(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	qid := make(chan uint64)
	done := make(chan struct{})

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			switch stmt.(type) {
			case *influxql.KillQueryStatement, *influxql.ShowQueriesStatement:
				return e.TaskManager.ExecuteStatement(ctx, stmt)
			}

			qid <- ctx.QueryID
			select {
			case <-ctx.Done():
				select {
				case <-done:
					// Keep the query running until we run SHOW QUERIES.
				case <-time.After(100 * time.Millisecond):
					// Ensure that we don't have a lingering goroutine.
				}
				return query.ErrQueryInterrupted
			case <-time.After(100 * time.Millisecond):
				t.Error("killing the query did not close the channel after 100 milliseconds")
				return errUnexpected
			}
		},
	}

	results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
	q, err = influxql.ParseQuery(fmt.Sprintf("KILL QUERY %d", <-qid))
	if err != nil {
		t.Fatal(err)
	}
	discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))

	// Display the queries and ensure that the original is still in there.
	q, err = influxql.ParseQuery("SHOW QUERIES")
	if err != nil {
		t.Fatal(err)
	}
	tasks := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)

	// The killed query should still be there.
	task := <-tasks
	if len(task.Series) != 1 {
		t.Errorf("expected %d series, got %d", 1, len(task.Series))
	} else if len(task.Series[0].Values) != 2 {
		t.Errorf("expected %d rows, got %d", 2, len(task.Series[0].Values))
	}
	close(done)

	// The original query should return.
	result := <-results
	if result.Err != query.ErrQueryInterrupted {
		t.Errorf("unexpected error: %s", result.Err)
	}
}

func TestQueryExecutor_KillQuery_CloseTaskManager(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	qid := make(chan uint64)

	// Open a channel to stall the statement executor forever. This keeps the statement executor
	// running even after we kill the query which can happen with some queries. We only close it once
	// the test has finished running.
	done := make(chan struct{})
	defer close(done)

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			switch stmt.(type) {
			case *influxql.KillQueryStatement, *influxql.ShowQueriesStatement:
				return e.TaskManager.ExecuteStatement(ctx, stmt)
			}

			qid <- ctx.QueryID
			<-done
			return nil
		},
	}

	// Kill the query. This should switch it into a zombie state.
	go discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
	q, err = influxql.ParseQuery(fmt.Sprintf("KILL QUERY %d", <-qid))
	if err != nil {
		t.Fatal(err)
	}
	discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))

	// Display the queries and ensure that the original is still in there.
	q, err = influxql.ParseQuery("SHOW QUERIES")
	if err != nil {
		t.Fatal(err)
	}
	tasks := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)

	// The killed query should still be there.
	task := <-tasks
	if len(task.Series) != 1 {
		t.Errorf("expected %d series, got %d", 1, len(task.Series))
	} else if len(task.Series[0].Values) != 2 {
		t.Errorf("expected %d rows, got %d", 2, len(task.Series[0].Values))
	}

	// Close the task manager to ensure it doesn't cause a panic.
	if err := e.TaskManager.Close(); err != nil {
		t.Errorf("unexpected error: %s", err)
	}
}

func TestQueryExecutor_KillQuery_AlreadyKilled(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	qid := make(chan uint64)

	// Open a channel to stall the statement executor forever. This keeps the statement executor
	// running even after we kill the query which can happen with some queries. We only close it once
	// the test has finished running.
	done := make(chan struct{})
	defer close(done)

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			switch stmt.(type) {
			case *influxql.KillQueryStatement, *influxql.ShowQueriesStatement:
				return e.TaskManager.ExecuteStatement(ctx, stmt)
			}

			qid <- ctx.QueryID
			<-done
			return nil
		},
	}

	// Kill the query. This should switch it into a zombie state.
	go discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
	q, err = influxql.ParseQuery(fmt.Sprintf("KILL QUERY %d", <-qid))
	if err != nil {
		t.Fatal(err)
	}
	discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))

	// Now attempt to kill it again. We should get an error.
	results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
	result := <-results
	if got, want := result.Err, query.ErrAlreadyKilled; got != want {
		t.Errorf("unexpected error: got=%v want=%v", got, want)
	}
}

func TestQueryExecutor_Interrupt(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(100 * time.Millisecond):
				t.Error("killing the query did not close the channel after 100 milliseconds")
				return errUnexpected
			}
		},
	}

	closing := make(chan struct{})
	results := e.ExecuteQuery(q, query.ExecutionOptions{}, closing)
	close(closing)
	result := <-results
	if result.Err != query.ErrQueryInterrupted {
		t.Errorf("unexpected error: %s", result.Err)
	}
}

func TestQueryExecutor_Abort(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	ch1 := make(chan struct{})
	ch2 := make(chan struct{})

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			<-ch1
			if err := ctx.Send(&query.Result{Err: errUnexpected}); err != query.ErrQueryAborted {
				t.Errorf("unexpected error: %v", err)
			}
			close(ch2)
			return nil
		},
	}

	done := make(chan struct{})
	close(done)

	results := e.ExecuteQuery(q, query.ExecutionOptions{AbortCh: done}, nil)
	close(ch1)

	<-ch2
	discardOutput(results)
}

func TestQueryExecutor_ShowQueries(t *testing.T) {
	const testUser = "Fred"
	// Column layout is queryFieldNames: host, qid, query, database, duration, status, user.
	const userColumn = 6
	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			switch stmt.(type) {
			case *influxql.ShowQueriesStatement:
				return e.TaskManager.ExecuteStatement(ctx, stmt)
			}

			t.Errorf("unexpected statement: %s", stmt)
			return errUnexpected
		},
	}

	q, err := influxql.ParseQuery(`SHOW QUERIES`)
	if err != nil {
		t.Fatal(err)
	}

	results := e.ExecuteQuery(q, query.ExecutionOptions{UserID: testUser}, nil)
	result := <-results
	if len(result.Series) != 1 {
		t.Errorf("expected %d series, got %d", 1, len(result.Series))
	} else if len(result.Series[0].Values) != 1 {
		t.Errorf("expected %d row, got %d", 1, len(result.Series[0].Values))
	} else if result.Series[0].Values[0][userColumn] != testUser {
		t.Errorf("unexpected user: %s", result.Series[0].Values[0][userColumn])
	} else if result.Series[0].Columns[userColumn] != "user" {
		t.Errorf("unexpected column: %s", result.Series[0].Columns[userColumn])
	}
	if result.Err != nil {
		t.Errorf("unexpected error: %s", result.Err)
	}
}

// TestQueryExecutor_ShowQueries_NonAdminFiltering verifies that SHOW QUERIES
// filters results based on the requesting user's database-level read
// permissions. A non-admin user should only see queries running against
// databases they have read access to, while an admin should see all queries.
func TestQueryExecutor_ShowQueries_NonAdminFiltering(t *testing.T) {
	const (
		// Column layout is queryFieldNames: host, qid, query, database, duration, status, user.
		dbColumn     = 3
		allowedDB    = "mydb"
		forbiddenDB  = "secretdb"
		nonAdminUser = "bar"
		adminUser    = "alice"
	)

	e := NewQueryExecutor()

	// blockCh keeps the "long-running" queries alive so they appear in SHOW QUERIES.
	blockCh := make(chan struct{})
	defer close(blockCh)

	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			switch stmt.(type) {
			case *influxql.ShowQueriesStatement:
				return e.TaskManager.ExecuteStatement(ctx, stmt)
			case *influxql.SelectStatement:
				// Block until the test completes so this query stays visible.
				<-blockCh
				return nil
			}
			t.Errorf("unexpected statement: %s", stmt)
			return errUnexpected
		},
	}

	// Start a "long-running" query on the forbidden database (by the admin user).
	q1, err := influxql.ParseQuery("SELECT * FROM cpu")
	require.NoError(t, err)
	go func() {
		discardOutput(e.ExecuteQuery(q1, query.ExecutionOptions{
			Database: forbiddenDB,
			UserID:   adminUser,
		}, nil))
	}()

	// Start a "long-running" query on the allowed database (by the non-admin user).
	q2, err := influxql.ParseQuery("SELECT * FROM mem")
	require.NoError(t, err)
	go func() {
		discardOutput(e.ExecuteQuery(q2, query.ExecutionOptions{
			Database: allowedDB,
			UserID:   nonAdminUser,
		}, nil))
	}()

	// Give the background queries a moment to register with the TaskManager.
	time.Sleep(50 * time.Millisecond)

	// Verify both queries are tracked (sanity check).
	allQueries := e.TaskManager.Queries()
	require.GreaterOrEqual(t, len(allQueries), 2, "expected at least 2 running queries")

	// Helper to collect database names from SHOW QUERIES result rows.
	databases := func(rows [][]interface{}) []string {
		var dbs []string
		for _, row := range rows {
			if db, ok := row[dbColumn].(string); ok && db != "" {
				dbs = append(dbs, db)
			}
		}
		return dbs
	}

	t.Run("non-admin user only sees allowed databases", func(t *testing.T) {
		showQ, err := influxql.ParseQuery("SHOW QUERIES")
		require.NoError(t, err)

		nonAdminUserInfo := &meta.UserInfo{
			Name:  nonAdminUser,
			Admin: false,
			Privileges: map[string]influxql.Privilege{
				allowedDB: influxql.AllPrivileges,
			},
		}

		results := e.ExecuteQuery(showQ, query.ExecutionOptions{
			UserID:           nonAdminUser,
			CoarseAuthorizer: nonAdminUserInfo,
		}, nil)
		result := <-results
		require.NoError(t, result.Err)
		require.Len(t, result.Series, 1)

		dbs := databases(result.Series[0].Values)
		require.Contains(t, dbs, allowedDB, "non-admin user should see queries on databases they have read access to")
		require.NotContains(t, dbs, forbiddenDB, "non-admin user should not see queries on databases they lack read access to")
	})

	t.Run("admin user sees all databases", func(t *testing.T) {
		showQ, err := influxql.ParseQuery("SHOW QUERIES")
		require.NoError(t, err)

		adminUserInfo := &meta.UserInfo{
			Name:  adminUser,
			Admin: true,
		}

		results := e.ExecuteQuery(showQ, query.ExecutionOptions{
			UserID:           adminUser,
			CoarseAuthorizer: adminUserInfo,
		}, nil)
		result := <-results
		require.NoError(t, result.Err)
		require.Len(t, result.Series, 1)

		dbs := databases(result.Series[0].Values)
		require.Contains(t, dbs, allowedDB, "admin user should see all queries")
		require.Contains(t, dbs, forbiddenDB, "admin user should see all queries")
	})
}

func TestQueryExecutor_Limit_Timeout(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			select {
			case <-ctx.Done():
				return ctx.Err()
			case <-time.After(time.Second):
				t.Errorf("timeout has not killed the query")
				return errUnexpected
			}
		},
	}
	e.TaskManager.QueryTimeout = time.Nanosecond

	results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
	result := <-results
	if result.Err == nil || !strings.Contains(result.Err.Error(), "query-timeout") {
		t.Errorf("unexpected error: %s", result.Err)
	}
}

func TestQueryExecutor_Limit_ConcurrentQueries(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	qid := make(chan uint64)

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			qid <- ctx.QueryID
			<-ctx.Done()
			return ctx.Err()
		},
	}
	e.TaskManager.MaxConcurrentQueries = 1
	defer e.Close()

	// Start first query and wait for it to be executing.
	go discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
	<-qid

	// Start second query and expect for it to fail.
	results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)

	select {
	case result := <-results:
		if len(result.Series) != 0 {
			t.Errorf("expected %d rows, got %d", 0, len(result.Series))
		}
		if result.Err == nil || !strings.Contains(result.Err.Error(), "max-concurrent-queries") {
			t.Errorf("unexpected error: %s", result.Err)
		}
	case <-qid:
		t.Errorf("unexpected statement execution for the second query")
	}
}

func TestQueryExecutor_Close(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	ch1 := make(chan struct{})
	ch2 := make(chan struct{})

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			close(ch1)
			<-ctx.Done()
			return ctx.Err()
		},
	}

	results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
	go func(results <-chan *query.Result) {
		result := <-results
		if result.Err != query.ErrQueryEngineShutdown {
			t.Errorf("unexpected error: %s", result.Err)
		}
		close(ch2)
	}(results)

	// Wait for the statement to start executing.
	<-ch1

	// Close the query executor.
	e.Close()

	// Check that the statement gets interrupted and finishes.
	select {
	case <-ch2:
	case <-time.After(100 * time.Millisecond):
		t.Fatal("closing the query manager did not kill the query after 100 milliseconds")
	}

	results = e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
	result := <-results
	if len(result.Series) != 0 {
		t.Errorf("expected %d rows, got %d", 0, len(result.Series))
	}
	if result.Err != query.ErrQueryEngineShutdown {
		t.Errorf("unexpected error: %s", result.Err)
	}
}

func TestQueryExecutor_Panic(t *testing.T) {
	q, err := influxql.ParseQuery(`SELECT count(value) FROM cpu`)
	if err != nil {
		t.Fatal(err)
	}

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			panic("test error")
		},
	}

	results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
	result := <-results
	if len(result.Series) != 0 {
		t.Errorf("expected %d rows, got %d", 0, len(result.Series))
	}
	if result.Err == nil || result.Err.Error() != "SELECT count(value) FROM cpu [panic:test error]" {
		t.Errorf("unexpected error: %s", result.Err)
	}
}

const goodStatement = `SELECT count(value) FROM cpu`

func TestQueryExecutor_NotExecuted(t *testing.T) {
	var executorFailIndex int
	var executorCallCount int
	queryStatements := []string{goodStatement, goodStatement, goodStatement, goodStatement, goodStatement}
	queryStr := strings.Join(queryStatements, ";")
	var closing chan struct{}

	q, err := influxql.ParseQuery(queryStr)
	if err != nil {
		t.Fatalf("parsing %s: %v", queryStr, err)
	}

	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			defer func() { executorCallCount++ }()
			if executorFailIndex == executorCallCount {
				closing <- struct{}{}
				close(closing)
				select {
				case <-ctx.Done():
					return nil
				}
			} else {
				return ctx.Send(&query.Result{Err: nil})
			}
		},
	}
	testFn := func(testName string, i int) {
		results := e.ExecuteQuery(q, query.ExecutionOptions{}, closing)
		checkNotExecutedResults(t, results, testName, i, len(q.Statements))
	}
	for i := 0; i < len(q.Statements); i++ {
		closing = make(chan struct{})
		executorFailIndex = i
		executorCallCount = 0
		testFn("executor", i)
	}
}

func checkNotExecutedResults(t *testing.T, results <-chan *query.Result, testName string, failIndex int, lenQuery int) {
	notExecutedIndex := failIndex + 1
	for result := range results {
		if result.Err == query.ErrNotExecuted {
			if result.StatementID != notExecutedIndex {
				t.Fatalf("StatementID for ErrNotExecuted in wrong order - expected: %d, got: %d", notExecutedIndex, result.StatementID)
			} else {
				notExecutedIndex++
			}
		}
	}
	if notExecutedIndex != lenQuery {
		t.Fatalf("wrong number of results from %s with fail index of %d - got: %d, expected: %d", testName, failIndex, notExecutedIndex-(1+failIndex), lenQuery-(1+failIndex))
	}
}

func TestQueryExecutor_InvalidSource(t *testing.T) {
	e := NewQueryExecutor()
	e.StatementExecutor = &StatementExecutor{
		ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
			return errors.New("statement executed unexpectedly")
		},
	}

	for i, tt := range []struct {
		q   string
		err string
	}{
		{
			q:   `SELECT fieldKey, fieldType FROM _fieldKeys`,
			err: `unable to use system source '_fieldKeys': use SHOW FIELD KEYS instead`,
		},
		{
			q:   `SELECT "name" FROM _measurements`,
			err: `unable to use system source '_measurements': use SHOW MEASUREMENTS instead`,
		},
		{
			q:   `SELECT "key" FROM _series`,
			err: `unable to use system source '_series': use SHOW SERIES instead`,
		},
		{
			q:   `SELECT tagKey FROM _tagKeys`,
			err: `unable to use system source '_tagKeys': use SHOW TAG KEYS instead`,
		},
		{
			q:   `SELECT "key", value FROM _tags`,
			err: `unable to use system source '_tags': use SHOW TAG VALUES instead`,
		},
	} {
		q, err := influxql.ParseQuery(tt.q)
		if err != nil {
			t.Errorf("%d. unable to parse: %s", i, tt.q)
			continue
		}

		results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)
		result := <-results
		if len(result.Series) != 0 {
			t.Errorf("%d. expected %d rows, got %d", 0, i, len(result.Series))
		}
		if result.Err == nil || result.Err.Error() != tt.err {
			t.Errorf("%d. unexpected error: %s", i, result.Err)
		}
	}
}

func discardOutput(results <-chan *query.Result) {
	for range results {
		// Read all results and discard.
	}
}

// queryExecutorStat returns the named value from the executor's queryExecutor
// statistic. It fails the test if the stat is missing or not an int64.
func queryExecutorStat(t *testing.T, e *query.Executor, key string) int64 {
	t.Helper()
	stats := e.Statistics(nil)
	require.Len(t, stats, 1)
	v, ok := stats[0].Values[key].(int64)
	require.Truef(t, ok, "stat %q missing or not int64: %v", key, stats[0].Values[key])
	return v
}

func TestQueryExecutor_Statistics_QueriesFailed(t *testing.T) {
	t.Run("success is not counted", func(t *testing.T) {
		e := NewQueryExecutor()
		e.StatementExecutor = &StatementExecutor{
			ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
				return nil
			},
		}
		q, err := influxql.ParseQuery(goodStatement)
		require.NoError(t, err)
		discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
		require.Equal(t, int64(0), queryExecutorStat(t, e, "queriesFailed"))
	})

	t.Run("single failure counts once", func(t *testing.T) {
		e := NewQueryExecutor()
		e.StatementExecutor = &StatementExecutor{
			ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
				return errUnexpected
			},
		}
		q, err := influxql.ParseQuery(goodStatement)
		require.NoError(t, err)
		discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
		require.Equal(t, int64(1), queryExecutorStat(t, e, "queriesFailed"))
	})

	t.Run("multi-statement failure counts once", func(t *testing.T) {
		e := NewQueryExecutor()
		var calls int
		e.StatementExecutor = &StatementExecutor{
			ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
				calls++
				return errUnexpected
			},
		}
		queryStr := strings.Join([]string{goodStatement, goodStatement, goodStatement}, ";")
		q, err := influxql.ParseQuery(queryStr)
		require.NoError(t, err)
		discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
		// The first statement fails; the remaining statements emit ErrNotExecuted
		// without being executed. The query counts as failed exactly once.
		require.Equal(t, 1, calls)
		require.Equal(t, int64(1), queryExecutorStat(t, e, "queriesFailed"))
	})

	t.Run("executor-streamed error counts", func(t *testing.T) {
		e := NewQueryExecutor()
		e.StatementExecutor = &StatementExecutor{
			ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
				// Mirrors real statement executors that report an error by
				// sending a Result and returning nil (see, e.g., the meta/SHOW
				// paths in coordinator/statement_executor.go).
				return ctx.Send(&query.Result{Err: errUnexpected})
			},
		}
		q, err := influxql.ParseQuery(goodStatement)
		require.NoError(t, err)
		discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
		require.Equal(t, int64(1), queryExecutorStat(t, e, "queriesFailed"))
	})

	t.Run("panic counts as failure", func(t *testing.T) {
		e := NewQueryExecutor()
		e.StatementExecutor = &StatementExecutor{
			ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
				panic("test error")
			},
		}
		q, err := influxql.ParseQuery(goodStatement)
		require.NoError(t, err)
		discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
		require.Equal(t, int64(1), queryExecutorStat(t, e, "queriesFailed"))
		require.Equal(t, int64(1), queryExecutorStat(t, e, "recoveredPanics"))
	})

	t.Run("error then panic counts once", func(t *testing.T) {
		e := NewQueryExecutor()
		var call int
		e.StatementExecutor = &StatementExecutor{
			ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
				call++
				if call == 1 {
					// First statement reports an error by sending a Result and
					// returning nil, which sets ctx.Failed().
					return ctx.Send(&query.Result{Err: errUnexpected})
				}
				// A later statement panics after the failure was already
				// recorded; the query must still be counted as failed only once.
				panic("test error")
			},
		}
		queryStr := strings.Join([]string{goodStatement, goodStatement}, ";")
		q, err := influxql.ParseQuery(queryStr)
		require.NoError(t, err)
		discardOutput(e.ExecuteQuery(q, query.ExecutionOptions{}, nil))
		require.Equal(t, int64(1), queryExecutorStat(t, e, "queriesFailed"))
		require.Equal(t, int64(1), queryExecutorStat(t, e, "recoveredPanics"))
	})
}

func TestTaskManager_SlowQueryCount(t *testing.T) {
	// newBlockingExecutor returns an executor whose statement execution blocks
	// until the returned release func is called (idempotent — safe to call more
	// than once), signalling on started once the query is attached and running.
	// Callers should defer release so a failed assertion cannot strand the
	// blocked query goroutine.
	newBlockingExecutor := func() (e *query.Executor, started <-chan struct{}, release func()) {
		startedCh := make(chan struct{})
		releaseCh := make(chan struct{})
		var once sync.Once
		e = NewQueryExecutor()
		e.StatementExecutor = &StatementExecutor{
			ExecuteStatementFn: func(stmt influxql.Statement, ctx *query.ExecutionContext) error {
				close(startedCh)
				<-releaseCh
				return nil
			},
		}
		return e, startedCh, func() { once.Do(func() { close(releaseCh) }) }
	}

	t.Run("counts in-flight queries slower than log-queries-after", func(t *testing.T) {
		e, started, release := newBlockingExecutor()
		defer e.Close()
		defer release()
		e.TaskManager.LogQueriesAfter = 20 * time.Millisecond

		q, err := influxql.ParseQuery(goodStatement)
		require.NoError(t, err)
		results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)

		<-started // query is attached and running

		// Once it has been running longer than the threshold, it is counted.
		require.Eventually(t, func() bool {
			return e.TaskManager.SlowQueryCount() == 1
		}, time.Second, 2*time.Millisecond)
		require.Equal(t, int64(1), queryExecutorStat(t, e, "queriesSlow"))

		release()
		discardOutput(results)

		// After the query finishes and detaches, the snapshot returns to zero.
		require.Equal(t, int64(0), e.TaskManager.SlowQueryCount())
	})

	t.Run("zero threshold counts nothing", func(t *testing.T) {
		e, started, release := newBlockingExecutor()
		defer e.Close()
		defer release()
		e.TaskManager.LogQueriesAfter = 0 // slow-query detection disabled

		q, err := influxql.ParseQuery(goodStatement)
		require.NoError(t, err)
		results := e.ExecuteQuery(q, query.ExecutionOptions{}, nil)

		<-started // query is attached and running

		// Even with a long-running in-flight query, a zero threshold never counts.
		require.Equal(t, int64(0), e.TaskManager.SlowQueryCount())
		require.Equal(t, int64(0), queryExecutorStat(t, e, "queriesSlow"))

		release()
		discardOutput(results)
	})
}