From d4b1cbdfbf52ddfd6d08b50f8a92319bf273ad90 Mon Sep 17 00:00:00 2001 From: maruware Date: Thu, 28 May 2026 02:46:29 +0900 Subject: [PATCH] feat: Support collection group --- completer.go | 39 +++++++++++++++++++++++++++++++++------ completer_test.go | 30 ++++++++++++++++++++++++++++++ docs/operations.md | 16 ++++++++++++++++ executor.go | 36 ++++++++++++++++++++++++++---------- executor_test.go | 19 +++++++++++++++++++ operation.go | 32 +++++++++++++++++++++++++------- parser.go | 46 ++++++++++++++++++++++++++++++++++++++++------ parser_test.go | 14 ++++++++++++++ token.go | 36 +++++++++++++++++++----------------- 9 files changed, 222 insertions(+), 46 deletions(-) diff --git a/completer.go b/completer.go index 4c9b91b..9cdc87a 100644 --- a/completer.go +++ b/completer.go @@ -20,10 +20,11 @@ var rootSuggestions = []prompt.Suggest{ } var ( - selectSuggestion = prompt.Suggest{Text: "SELECT", Description: "SELECT [field...]"} - whereSuggestion = prompt.Suggest{Text: "WHERE", Description: "WHERE [field] [operator] [value]"} - orderBySuggestion = prompt.Suggest{Text: "ORDER BY", Description: "ORDER BY [field] [ASC/DESC]"} - limitSuggestion = prompt.Suggest{Text: "LIMIT", Description: "LIMIT [count]"} + selectSuggestion = prompt.Suggest{Text: "SELECT", Description: "SELECT [field...]"} + whereSuggestion = prompt.Suggest{Text: "WHERE", Description: "WHERE [field] [operator] [value]"} + orderBySuggestion = prompt.Suggest{Text: "ORDER BY", Description: "ORDER BY [field] [ASC/DESC]"} + limitSuggestion = prompt.Suggest{Text: "LIMIT", Description: "LIMIT [count]"} + collectionGroupSuggestion = prompt.Suggest{Text: "COLLECTION_GROUP", Description: "COLLECTION_GROUP [collection]"} ) var querySuggestions = []prompt.Suggest{ @@ -83,12 +84,24 @@ func (c *Completer) Parse() ([]prompt.Suggest, error) { } func (c *Completer) parseQueryOperation() ([]prompt.Suggest, error) { - if !c.expectPeek(IDENT) { + if !c.peekTokenIs(IDENT) && !c.peekTokenIs(COLLECTION_GROUP) { return []prompt.Suggest{}, nil } + c.nextToken() + + collectionGroupMode := false + if c.curTokenIs(COLLECTION_GROUP) { + collectionGroupMode = true + if !c.expectPeek(IDENT) { + return []prompt.Suggest{}, nil + } + } // collection := c.curToken.Literal if c.peekTokenIs(EOF) { + if collectionGroupMode { + return querySuggestions, nil + } collection := normalizeFirestorePath(c.curToken.Literal) parts := strings.Split(collection, "/") if len(parts)%2 == 0 { @@ -109,6 +122,7 @@ func (c *Completer) parseQueryOperation() ([]prompt.Suggest, error) { for _, col := range collections { suggestions = append(suggestions, newCollectionSuggestion(baseDoc, col)) } + suggestions = append(suggestions, collectionGroupSuggestion) return prompt.FilterHasPrefix(suggestions, c.curToken.Literal, false), nil } @@ -292,11 +306,23 @@ func (c *Completer) parseGetOperation() ([]prompt.Suggest, error) { } func (c *Completer) parseCountOperation() ([]prompt.Suggest, error) { - if !c.expectPeek(IDENT) { + if !c.peekTokenIs(IDENT) && !c.peekTokenIs(COLLECTION_GROUP) { return []prompt.Suggest{}, nil } + c.nextToken() + + collectionGroupMode := false + if c.curTokenIs(COLLECTION_GROUP) { + collectionGroupMode = true + if !c.expectPeek(IDENT) { + return []prompt.Suggest{}, nil + } + } if c.peekTokenIs(EOF) { + if collectionGroupMode { + return []prompt.Suggest{whereSuggestion}, nil + } collection := normalizeFirestorePath(c.curToken.Literal) parts := strings.Split(collection, "/") if len(parts)%2 == 0 { @@ -317,6 +343,7 @@ func (c *Completer) parseCountOperation() ([]prompt.Suggest, error) { for _, col := range collections { suggestions = append(suggestions, newCollectionSuggestion(baseDoc, col)) } + suggestions = append(suggestions, collectionGroupSuggestion) return prompt.FilterHasPrefix(suggestions, c.curToken.Literal, false), nil } diff --git a/completer_test.go b/completer_test.go index 7316d06..e67458d 100644 --- a/completer_test.go +++ b/completer_test.go @@ -43,6 +43,21 @@ func TestCompleter(t *testing.T) { input: `QUERY us`, want: []prompt.Suggest{newCollectionSuggestion("", "user")}, }, + { + desc: "middle of query with collection_group keyword", + input: `QUERY C`, + want: []prompt.Suggest{collectionGroupSuggestion}, + }, + { + desc: "query collection_group with name and where", + input: `QUERY COLLECTION_GROUP posts WH`, + want: []prompt.Suggest{whereSuggestion}, + }, + { + desc: "query collection_group with name", + input: `QUERY COLLECTION_GROUP posts`, + want: querySuggestions, + }, { desc: "middle of query with sub collection", input: `QUERY user/1/p`, @@ -174,6 +189,21 @@ func TestCompleter(t *testing.T) { input: `COUNT us`, want: []prompt.Suggest{newCollectionSuggestion("", "user")}, }, + { + desc: "middle of count with collection_group keyword", + input: `COUNT C`, + want: []prompt.Suggest{collectionGroupSuggestion}, + }, + { + desc: "count collection_group with name and where", + input: `COUNT COLLECTION_GROUP posts WH`, + want: []prompt.Suggest{whereSuggestion}, + }, + { + desc: "count collection_group with name", + input: `COUNT COLLECTION_GROUP posts`, + want: []prompt.Suggest{whereSuggestion}, + }, { desc: "middle of count with sub collection", input: `COUNT user/1/p`, diff --git a/docs/operations.md b/docs/operations.md index e1160f3..e2a0f4a 100644 --- a/docs/operations.md +++ b/docs/operations.md @@ -8,6 +8,7 @@ Query documents in a collection. Supports `SELECT`, `WHERE`, `ORDER BY`, and `LI ``` QUERY [SELECT ...] [WHERE ...] [ORDER BY ...] [LIMIT ...] +QUERY COLLECTION_GROUP [SELECT ...] [WHERE ...] [ORDER BY ...] [LIMIT ...] ``` ### Examples @@ -21,6 +22,9 @@ QUERY users WHERE age = 20 -- Subcollection query QUERY users/abc123/posts + +-- Collection group query +QUERY COLLECTION_GROUP posts WHERE title = "post-1-1" ``` ## GET @@ -49,6 +53,7 @@ Count documents in a collection. Supports `WHERE` clause for filtering. ``` COUNT [WHERE ...] +COUNT COLLECTION_GROUP [WHERE ...] ``` Returns a single integer. @@ -61,6 +66,9 @@ COUNT users -- Count with filter COUNT users WHERE name = "takashi" + +-- Collection group count +COUNT COLLECTION_GROUP posts WHERE title = "post-1-1" ``` ## Collection Path @@ -75,3 +83,11 @@ collection/docId/subcollection/docId/subcollection/... - **Document path** (odd number of segments): `users/abc123`, `users/abc123/posts/xyz789` Leading slashes are automatically stripped. + +## Collection Group + +`COLLECTION_GROUP` targets all collections with the same collection ID across the database hierarchy. + +- Use a single collection ID (for example, `posts`) +- Slash-separated paths are not allowed (for example, `users/posts` is invalid) +- Supported with `QUERY` and `COUNT` diff --git a/executor.go b/executor.go index 975044c..f0cd6b0 100644 --- a/executor.go +++ b/executor.go @@ -20,17 +20,25 @@ func NewExecutor(ctx context.Context, fs *firestore.Client) *Executor { } func (exe *Executor) ExecuteQuery(ctx context.Context, op *QueryOperation) ([]*firestore.DocumentSnapshot, error) { - collection := exe.fs.Collection(op.Collection()) - if collection == nil { - return nil, ErrInvalidCollection + var q firestore.Query + if op.IsCollectionGroup() { + q = exe.fs.CollectionGroup(op.Collection()).Query + } else { + collection := exe.fs.Collection(op.Collection()) + if collection == nil { + return nil, ErrInvalidCollection + } + q = collection.Query } - q := collection.Query + for _, filter := range op.filters { fieldName := filter.FieldName() value := filter.Value() if fieldName == FieldDocumentID { fieldName = firestore.DocumentID - value = toDocRefValue(collection, value) + if !op.IsCollectionGroup() { + value = toDocRefValue(exe.fs.Collection(op.Collection()), value) + } } q = q.Where(fieldName, string(filter.Operator()), value) } @@ -76,17 +84,25 @@ func (exe *Executor) ExecuteGet(ctx context.Context, op *GetOperation) (*firesto } func (exe *Executor) ExecuteCount(ctx context.Context, op *CountOperation) (int64, error) { - collection := exe.fs.Collection(op.Collection()) - if collection == nil { - return 0, ErrInvalidCollection + var q firestore.Query + if op.IsCollectionGroup() { + q = exe.fs.CollectionGroup(op.Collection()).Query + } else { + collection := exe.fs.Collection(op.Collection()) + if collection == nil { + return 0, ErrInvalidCollection + } + q = collection.Query } - q := collection.Query + for _, filter := range op.filters { fieldName := filter.FieldName() value := filter.Value() if fieldName == FieldDocumentID { fieldName = firestore.DocumentID - value = toDocRefValue(collection, value) + if !op.IsCollectionGroup() { + value = toDocRefValue(exe.fs.Collection(op.Collection()), value) + } } q = q.Where(fieldName, string(filter.Operator()), value) } diff --git a/executor_test.go b/executor_test.go index dd95368..1eb8a46 100644 --- a/executor_test.go +++ b/executor_test.go @@ -146,6 +146,18 @@ func TestQuery(t *testing.T) { }, }, }, + { + desc: "query with collection group", + input: NewCollectionGroupQueryOperation("posts", nil, []Filter{ + NewStringFilter("title", "==", "post-1-1"), + }, nil, 0), + want: []map[string]any{ + { + "title": "post-1-1", + "publishedAt": time.Date(2025, 1, 2, 0, 0, 0, 0, time.UTC), + }, + }, + }, { desc: "query with not equal", input: &QueryOperation{collection: "users", filters: []Filter{ @@ -448,6 +460,13 @@ func TestCount(t *testing.T) { }), want: 3, }, + { + desc: "count with collection group", + input: NewCollectionGroupCountOperation("posts", []Filter{ + NewStringFilter("title", "==", "post-1-1"), + }), + want: 1, + }, } for _, tt := range tests { diff --git a/operation.go b/operation.go index cb76fb3..3c02b3a 100644 --- a/operation.go +++ b/operation.go @@ -133,17 +133,22 @@ type OrderBy struct { type QueryOperation struct { BaseOperation - collection string - selects []string - filters []Filter - orderBys []OrderBy - limit int + collection string + collectionGroup bool + selects []string + filters []Filter + orderBys []OrderBy + limit int } func NewQueryOperation(collection string, selects []string, filters []Filter, orderBys []OrderBy, limit int) *QueryOperation { return &QueryOperation{collection: collection, selects: selects, filters: filters, orderBys: orderBys, limit: limit} } +func NewCollectionGroupQueryOperation(collectionGroup string, selects []string, filters []Filter, orderBys []OrderBy, limit int) *QueryOperation { + return &QueryOperation{collection: collectionGroup, collectionGroup: true, selects: selects, filters: filters, orderBys: orderBys, limit: limit} +} + func (op *QueryOperation) OperationType() OperationType { return OPERATION_TYPE_QUERY } @@ -152,6 +157,10 @@ func (op *QueryOperation) Collection() string { return op.collection } +func (op *QueryOperation) IsCollectionGroup() bool { + return op.collectionGroup +} + type GetOperation struct { BaseOperation collection string @@ -181,14 +190,19 @@ func (op *GetOperation) Selects() []string { type CountOperation struct { BaseOperation - collection string - filters []Filter + collection string + collectionGroup bool + filters []Filter } func NewCountOperation(collection string, filters []Filter) *CountOperation { return &CountOperation{collection: collection, filters: filters} } +func NewCollectionGroupCountOperation(collectionGroup string, filters []Filter) *CountOperation { + return &CountOperation{collection: collectionGroup, collectionGroup: true, filters: filters} +} + func (op *CountOperation) OperationType() OperationType { return OPERATION_TYPE_COUNT } @@ -196,3 +210,7 @@ func (op *CountOperation) OperationType() OperationType { func (op *CountOperation) Collection() string { return op.collection } + +func (op *CountOperation) IsCollectionGroup() bool { + return op.collectionGroup +} diff --git a/parser.go b/parser.go index 848dd5d..6cc5d56 100644 --- a/parser.go +++ b/parser.go @@ -54,11 +54,23 @@ func (p *Parser) Parse() (ParseResult, error) { func (p *Parser) parseQueryOperation() (*QueryOperation, error) { op := &QueryOperation{} - if !p.expectPeek(IDENT) { - return nil, fmt.Errorf("invalid: expected collection but got %s", p.curToken.Literal) + if !p.peekTokenIs(IDENT) && !p.peekTokenIs(COLLECTION_GROUP) { + return nil, fmt.Errorf("invalid: expected collection but got %s", p.peekToken.Literal) } + p.nextToken() - op.collection = normalizeFirestorePath(p.curToken.Literal) + if p.curTokenIs(COLLECTION_GROUP) { + op.collectionGroup = true + if !p.expectPeek(IDENT) { + return nil, fmt.Errorf("invalid: expected collection group name but got %s", p.peekToken.Literal) + } + op.collection = p.curToken.Literal + if err := validateCollectionGroupName(op.collection); err != nil { + return nil, err + } + } else { + op.collection = normalizeFirestorePath(p.curToken.Literal) + } p.nextToken() @@ -170,11 +182,23 @@ func (p *Parser) parseGetOperation() (*GetOperation, error) { func (p *Parser) parseCountOperation() (*CountOperation, error) { op := &CountOperation{} - if !p.expectPeek(IDENT) { - return nil, fmt.Errorf("invalid: expected collection but got %s", p.curToken.Literal) + if !p.peekTokenIs(IDENT) && !p.peekTokenIs(COLLECTION_GROUP) { + return nil, fmt.Errorf("invalid: expected collection but got %s", p.peekToken.Literal) } + p.nextToken() - op.collection = normalizeFirestorePath(p.curToken.Literal) + if p.curTokenIs(COLLECTION_GROUP) { + op.collectionGroup = true + if !p.expectPeek(IDENT) { + return nil, fmt.Errorf("invalid: expected collection group name but got %s", p.peekToken.Literal) + } + op.collection = p.curToken.Literal + if err := validateCollectionGroupName(op.collection); err != nil { + return nil, err + } + } else { + op.collection = normalizeFirestorePath(p.curToken.Literal) + } p.nextToken() @@ -445,3 +469,13 @@ func parseTime(timeStr string) (time.Time, error) { } return time.Time{}, fmt.Errorf("invalid timestamp format: %s", timeStr) } + +func validateCollectionGroupName(name string) error { + if name == "" { + return fmt.Errorf("invalid: collection group name is required") + } + if strings.Contains(name, "/") { + return fmt.Errorf("invalid: collection group name must not contain '/': %s", name) + } + return nil +} diff --git a/parser_test.go b/parser_test.go index 53d7f2b..2a101b2 100644 --- a/parser_test.go +++ b/parser_test.go @@ -141,6 +141,13 @@ func TestParse(t *testing.T) { input: `QUERY /user`, want: &QueryOperation{collection: "user"}, }, + { + desc: "query with collection group", + input: `QUERY COLLECTION_GROUP posts WHERE title = "post-1-1"`, + want: &QueryOperation{collection: "posts", collectionGroup: true, filters: []Filter{ + NewStringFilter("title", OPERATOR_EQ, "post-1-1"), + }}, + }, { desc: "query with select", input: `QUERY user SELECT name, age`, @@ -208,6 +215,13 @@ func TestParse(t *testing.T) { NewIntFilter("age", OPERATOR_EQ, 20), }}, }, + { + desc: "count with collection group", + input: `COUNT COLLECTION_GROUP posts WHERE title = "post-1-1"`, + want: &CountOperation{collection: "posts", collectionGroup: true, filters: []Filter{ + NewStringFilter("title", OPERATOR_EQ, "post-1-1"), + }}, + }, { desc: "get", input: `GET user/1`, diff --git a/token.go b/token.go index 2b39fad..89dc4c9 100644 --- a/token.go +++ b/token.go @@ -6,10 +6,11 @@ const ( EOF = "EOF" ILLEGAL = "ILLEGAL" - GET = "GET" - QUERY = "QUERY" - COUNT = "COUNT" - SELECT = "SELECT" + GET = "GET" + QUERY = "QUERY" + COUNT = "COUNT" + SELECT = "SELECT" + COLLECTION_GROUP = "COLLECTION_GROUP" WHERE = "WHERE" EQ = "=" @@ -55,18 +56,19 @@ type Token struct { Literal string } -var keywards = map[string]TokenType{ - "GET": GET, - "QUERY": QUERY, - "COUNT": COUNT, - "SELECT": SELECT, - "WHERE": WHERE, - "AND": AND, - "ORDER": ORDER, - "BY": BY, - "ASC": ASC, - "DESC": DESC, - "LIMIT": LIMIT, +var keywords = map[string]TokenType{ + "GET": GET, + "QUERY": QUERY, + "COUNT": COUNT, + "SELECT": SELECT, + "COLLECTION_GROUP": COLLECTION_GROUP, + "WHERE": WHERE, + "AND": AND, + "ORDER": ORDER, + "BY": BY, + "ASC": ASC, + "DESC": DESC, + "LIMIT": LIMIT, } var operators = map[string]TokenType{ @@ -88,7 +90,7 @@ var metacommands = map[string]TokenType{ func LookupIdent(ident string) TokenType { u := strings.ToUpper(ident) - if tok, ok := keywards[u]; ok { + if tok, ok := keywords[u]; ok { return tok } if tok, ok := operators[u]; ok {