googleapis--mcp-toolbox/internal/tools/firestore/firestorequerycollection/firestorequerycollection.go

// Copyright 2025 Google LLC
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
//     http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.

package firestorequerycollection

import (
	"context"
	"encoding/json"
	"fmt"
	"net/http"
	"strings"

	firestoreapi "cloud.google.com/go/firestore"
	yaml "github.com/goccy/go-yaml"
	"github.com/googleapis/mcp-toolbox/internal/tools"
	fsUtil "github.com/googleapis/mcp-toolbox/internal/tools/firestore/util"
	"github.com/googleapis/mcp-toolbox/internal/util"
	"github.com/googleapis/mcp-toolbox/internal/util/parameters"
)

// Constants for tool configuration
const (
	resourceType    = "firestore-query-collection"
	defaultLimit    = 100
	defaultAnalyze  = false
	maxFilterLength = 100 // Maximum filters to prevent abuse
)

// Parameter keys
const (
	collectionPathKey = "collectionPath"
	filtersKey        = "filters"
	orderByKey        = "orderBy"
	limitKey          = "limit"
	analyzeQueryKey   = "analyzeQuery"
)

// Firestore operators
var validOperators = map[string]bool{
	"<":                  true,
	"<=":                 true,
	">":                  true,
	">=":                 true,
	"==":                 true,
	"!=":                 true,
	"array-contains":     true,
	"array-contains-any": true,
	"in":                 true,
	"not-in":             true,
}

// Error messages
const (
	errMissingCollectionPath = "invalid or missing '%s' parameter"
	errInvalidFilters        = "invalid '%s' parameter; expected an array"
	errFilterNotString       = "filter at index %d is not a string"
	errFilterParseFailed     = "failed to parse filter at index %d: %w"
	errInvalidOperator       = "unsupported operator: %s. Valid operators are: %v"
	errMissingFilterValue    = "no value specified for filter on field '%s'"
	errOrderByParseFailed    = "failed to parse orderBy: %w"
	errTooManyFilters        = "too many filters provided: %d (maximum: %d)"
)

func init() {
	if !tools.Register(resourceType, newConfig) {
		panic(fmt.Sprintf("tool type %q already registered", resourceType))
	}
}

func newConfig(ctx context.Context, name string, decoder *yaml.Decoder) (tools.ToolConfig, error) {
	actual := Config{ConfigBase: tools.ConfigBase{Name: name}}
	if err := decoder.DecodeContext(ctx, &actual); err != nil {
		return nil, err
	}
	return actual, nil
}

// compatibleSource defines the interface for sources that can provide a Firestore client
type compatibleSource interface {
	FirestoreClient() *firestoreapi.Client
	BuildQuery(string, firestoreapi.EntityFilter, []string, string, firestoreapi.Direction, int, bool) (*firestoreapi.Query, error)
	ExecuteQuery(context.Context, *firestoreapi.Query, bool) (any, error)
}

// Config represents the configuration for the Firestore query collection tool
type Config struct {
	tools.ConfigBase `yaml:",inline"`
	Type             string                 `yaml:"type" validate:"required"`
	Source           string                 `yaml:"source" validate:"required"`
	Annotations      *tools.ToolAnnotations `yaml:"annotations,omitempty"`
}

// validate interface
var _ tools.ToolConfig = Config{}

// ToolConfigType returns the type of tool configuration
func (cfg Config) ToolConfigType() string {
	return resourceType
}

// Initialize creates a new Tool instance from the configuration
func (cfg Config) Initialize(context.Context) (tools.Tool, error) {
	if cfg.Description == "" {
		return nil, fmt.Errorf("description is required for tool %q", cfg.Name)
	}

	// Create parameters
	params := createParameters()

	return Tool{
		BaseTool: tools.NewBaseTool(
			cfg,
			tools.GetAnnotationsOrDefault(cfg.Annotations, tools.NewReadOnlyAnnotations),
			tools.Manifest{Description: cfg.Description, Parameters: params.Manifest(), AuthRequired: cfg.AuthRequired},
			params,
		),
	}, nil
}

func (t Tool) ToConfig() tools.ToolConfig {
	return t.Cfg
}

// createParameters creates the parameter definitions for the tool
func createParameters() parameters.Parameters {
	collectionPathParameter := parameters.NewStringParameter(
		collectionPathKey,
		"The relative path to the Firestore collection to query (e.g., 'users' or 'users/userId/posts'). Note: This is a relative path, NOT an absolute path like 'projects/{project_id}/databases/{database_id}/documents/...'",
	)

	filtersDescription := `Array of filter objects to apply to the query. Each filter is a JSON string with:
- field: The field name to filter on
- op: The operator to use ("<", "<=", ">", ">=", "==", "!=", "array-contains", "array-contains-any", "in", "not-in")
- value: The value to compare against (can be string, number, boolean, or array)
Example: {"field": "age", "op": ">", "value": 18}`

	filtersParameter := parameters.NewArrayParameter(
		filtersKey,
		filtersDescription,
		parameters.NewStringParameter("item", "JSON string representation of a filter object"),
	)

	orderByParameter := parameters.NewStringParameter(
		orderByKey,
		"JSON string specifying the field and direction to order by (e.g., {\"field\": \"name\", \"direction\": \"ASCENDING\"}). Leave empty if not specified",
	)

	limitParameter := parameters.NewIntParameter(
		limitKey,
		"The maximum number of documents to return", parameters.WithIntDefault(
			defaultLimit))

	analyzeQueryParameter := parameters.NewBooleanParameter(
		analyzeQueryKey,
		"If true, returns query explain metrics including execution statistics", parameters.WithBooleanDefault(
			defaultAnalyze))

	return parameters.Parameters{
		collectionPathParameter,
		filtersParameter,
		orderByParameter,
		limitParameter,
		analyzeQueryParameter,
	}
}

// validate interface
var _ tools.Tool = Tool{}

// Tool represents the Firestore query collection tool
type Tool struct {
	tools.BaseTool[Config]
}

// FilterConfig represents a filter for the query
type FilterConfig struct {
	Field string      `json:"field"`
	Op    string      `json:"op"`
	Value interface{} `json:"value"`
}

// Validate checks if the filter configuration is valid
func (f *FilterConfig) Validate() error {
	if f.Field == "" {
		return fmt.Errorf("filter field cannot be empty")
	}

	if !validOperators[f.Op] {
		ops := make([]string, 0, len(validOperators))
		for op := range validOperators {
			ops = append(ops, op)
		}
		return fmt.Errorf(errInvalidOperator, f.Op, ops)
	}

	if f.Value == nil {
		return fmt.Errorf(errMissingFilterValue, f.Field)
	}

	return nil
}

// OrderByConfig represents ordering configuration
type OrderByConfig struct {
	Field     string `json:"field"`
	Direction string `json:"direction"`
}

// GetDirection returns the Firestore direction constant
func (o *OrderByConfig) GetDirection() firestoreapi.Direction {
	if strings.EqualFold(o.Direction, "DESCENDING") {
		return firestoreapi.Desc
	}
	return firestoreapi.Asc
}

// Invoke executes the Firestore query based on the provided parameters
func (t Tool) Invoke(ctx context.Context, resourceMgr tools.SourceProvider, params parameters.ParamValues, accessToken tools.AccessToken) (any, util.ToolboxError) {
	source, err := tools.GetCompatibleSource[compatibleSource](resourceMgr, t.Cfg.Source, t.Cfg.Name, t.Cfg.Type)
	if err != nil {
		return nil, util.NewClientServerError("source used is not compatible with the tool", http.StatusInternalServerError, err)
	}

	// Parse parameters
	queryParams, err := t.parseQueryParameters(params)
	if err != nil {
		return nil, util.NewAgentError(fmt.Sprintf("failed to parse query parameters: %v", err), err)
	}

	var filter firestoreapi.EntityFilter
	// Apply filters
	if len(queryParams.Filters) > 0 {
		filterConditions := make([]firestoreapi.EntityFilter, 0, len(queryParams.Filters))
		for _, filter := range queryParams.Filters {
			filterConditions = append(filterConditions, firestoreapi.PropertyFilter{
				Path:     filter.Field,
				Operator: filter.Op,
				Value:    filter.Value,
			})
		}

		filter = firestoreapi.AndFilter{
			Filters: filterConditions,
		}
	}

	// prevent panic incase queryParams.OrderBy is nil
	var orderByField string
	var orderByDirection firestoreapi.Direction
	if queryParams.OrderBy != nil {
		orderByField = queryParams.OrderBy.Field
		orderByDirection = queryParams.OrderBy.GetDirection()
	}

	// Build the query
	query, err := source.BuildQuery(queryParams.CollectionPath, filter, nil, orderByField, orderByDirection, queryParams.Limit, queryParams.AnalyzeQuery)
	if err != nil {
		return nil, util.ProcessGcpError(err)
	}
	resp, err := source.ExecuteQuery(ctx, query, queryParams.AnalyzeQuery)
	if err != nil {
		return nil, util.ProcessGcpError(err)
	}
	return resp, nil
}

// queryParameters holds all parsed query parameters
type queryParameters struct {
	CollectionPath string
	Filters        []FilterConfig
	OrderBy        *OrderByConfig
	Limit          int
	AnalyzeQuery   bool
}

// parseQueryParameters extracts and validates parameters from the input
func (t Tool) parseQueryParameters(params parameters.ParamValues) (*queryParameters, error) {
	mapParams := params.AsMap()

	// Get collection path
	collectionPath, ok := mapParams[collectionPathKey].(string)
	if !ok || collectionPath == "" {
		return nil, fmt.Errorf(errMissingCollectionPath, collectionPathKey)
	}

	// Validate collection path
	if err := fsUtil.ValidateCollectionPath(collectionPath); err != nil {
		return nil, fmt.Errorf("invalid collection path: %w", err)
	}

	result := &queryParameters{
		CollectionPath: collectionPath,
		Limit:          defaultLimit,
		AnalyzeQuery:   defaultAnalyze,
	}

	// Parse filters
	if filtersRaw, ok := mapParams[filtersKey]; ok && filtersRaw != nil {
		filters, err := t.parseFilters(filtersRaw)
		if err != nil {
			return nil, err
		}
		result.Filters = filters
	}

	// Parse orderBy
	if orderByRaw, ok := mapParams[orderByKey]; ok && orderByRaw != nil {
		orderBy, err := t.parseOrderBy(orderByRaw)
		if err != nil {
			return nil, err
		}
		result.OrderBy = orderBy
	}

	// Parse limit
	if limit, ok := mapParams[limitKey].(int); ok {
		result.Limit = limit
	}

	// Parse analyze
	if analyze, ok := mapParams[analyzeQueryKey].(bool); ok {
		result.AnalyzeQuery = analyze
	}

	return result, nil
}

// parseFilters parses and validates filter configurations
func (t Tool) parseFilters(filtersRaw interface{}) ([]FilterConfig, error) {
	filters, ok := filtersRaw.([]any)
	if !ok {
		return nil, fmt.Errorf(errInvalidFilters, filtersKey)
	}

	if len(filters) > maxFilterLength {
		return nil, fmt.Errorf(errTooManyFilters, len(filters), maxFilterLength)
	}

	result := make([]FilterConfig, 0, len(filters))
	for i, filterRaw := range filters {
		filterJSON, ok := filterRaw.(string)
		if !ok {
			return nil, fmt.Errorf(errFilterNotString, i)
		}

		var filter FilterConfig
		if err := json.Unmarshal([]byte(filterJSON), &filter); err != nil {
			return nil, fmt.Errorf(errFilterParseFailed, i, err)
		}

		if err := filter.Validate(); err != nil {
			return nil, fmt.Errorf("filter at index %d is invalid: %w", i, err)
		}

		result = append(result, filter)
	}

	return result, nil
}

// parseOrderBy parses the orderBy configuration
func (t Tool) parseOrderBy(orderByRaw interface{}) (*OrderByConfig, error) {
	orderByJSON, ok := orderByRaw.(string)
	if !ok || orderByJSON == "" {
		return nil, nil
	}

	var orderBy OrderByConfig
	if err := json.Unmarshal([]byte(orderByJSON), &orderBy); err != nil {
		return nil, fmt.Errorf(errOrderByParseFailed, err)
	}

	if orderBy.Field == "" {
		return nil, nil
	}

	return &orderBy, nil
}