Generate SQL queries with comments for find operations (#3697)

Co-authored-by: noisersup <patryk@kwiatek.xyz>
FerretDB · Nov 17, 2023 · 05fb806 · 05fb806
1 parent 532787d
commit 05fb806
Show file tree

Hide file tree

Showing 7 changed files with 178 additions and 23 deletions.
diff --git a/internal/backends/collection.go b/internal/backends/collection.go
@@ -83,7 +83,7 @@ type QueryParams struct {
 	Sort          *SortField
 	Limit         int64
 	OnlyRecordIDs bool
-	Comment       string // TODO https://github.com/FerretDB/FerretDB/issues/3573
+	Comment       string
 }
 
 // QueryResult represents the results of Collection.Query method.

diff --git a/internal/backends/postgresql/collection.go b/internal/backends/postgresql/collection.go
@@ -78,7 +78,13 @@ func (c *collection) Query(ctx context.Context, params *backends.QueryParams) (*
 		}, nil
 	}
 
-	q := prepareSelectClause(c.dbName, meta.TableName, meta.Capped(), params.OnlyRecordIDs)
+	q := prepareSelectClause(&selectParams{
+		Schema:        c.dbName,
+		Table:         meta.TableName,
+		Comment:       params.Comment,
+		Capped:        meta.Capped(),
+		OnlyRecordIDs: params.OnlyRecordIDs,
+	})
 
 	var placeholder metadata.Placeholder
 
@@ -297,7 +303,13 @@ func (c *collection) Explain(ctx context.Context, params *backends.ExplainParams
 
 	res := new(backends.ExplainResult)
 
-	q := `EXPLAIN (VERBOSE true, FORMAT JSON) ` + prepareSelectClause(c.dbName, meta.TableName, meta.Capped(), false)
+	opts := &selectParams{
+		Schema: c.dbName,
+		Table:  meta.TableName,
+		Capped: meta.Capped(),
+	}
+
+	q := `EXPLAIN (VERBOSE true, FORMAT JSON) ` + prepareSelectClause(opts)
 
 	var placeholder metadata.Placeholder
 

diff --git a/internal/backends/postgresql/query.go b/internal/backends/postgresql/query.go
@@ -31,34 +31,57 @@ import (
 	"github.com/FerretDB/FerretDB/internal/util/must"
 )
 
+// selectParams contains params that specify how prepareSelectClause function will
+// build the SELECT SQL query.
+type selectParams struct {
+	Schema  string
+	Table   string
+	Comment string
+
+	Capped        bool
+	OnlyRecordIDs bool
+}
+
 // prepareSelectClause returns SELECT clause for default column of provided schema and table name.
 //
 // For capped collection with onlyRecordIDs, it returns select clause for recordID column.
 //
 // For capped collection, it returns select clause for recordID column and default column.
-func prepareSelectClause(schema, table string, capped, onlyRecordIDs bool) string {
-	if capped && onlyRecordIDs {
+func prepareSelectClause(params *selectParams) string {
+	if params == nil {
+		params = new(selectParams)
+	}
+
+	if params.Comment != "" {
+		params.Comment = strings.ReplaceAll(params.Comment, "/*", "/ *")
+		params.Comment = strings.ReplaceAll(params.Comment, "*/", "* /")
+		params.Comment = `/* ` + params.Comment + ` */`
+	}
+
+	if params.Capped && params.OnlyRecordIDs {
 		return fmt.Sprintf(
-			`SELECT %s FROM %s`,
+			`SELECT %s %s FROM %s`,
+			params.Comment,
 			metadata.RecordIDColumn,
-			pgx.Identifier{schema, table}.Sanitize(),
+			pgx.Identifier{params.Schema, params.Table}.Sanitize(),
 		)
 	}
 
-	if capped {
+	if params.Capped {
 		return fmt.Sprintf(
-			`SELECT %s, %s FROM %s`,
+			`SELECT %s %s, %s FROM %s`,
+			params.Comment,
 			metadata.RecordIDColumn,
 			metadata.DefaultColumn,
-			pgx.Identifier{schema, table}.Sanitize(),
+			pgx.Identifier{params.Schema, params.Table}.Sanitize(),
 		)
 	}
 
-	// TODO https://github.com/FerretDB/FerretDB/issues/3573
 	return fmt.Sprintf(
-		`SELECT %s FROM %s`,
+		`SELECT %s %s FROM %s`,
+		params.Comment,
 		metadata.DefaultColumn,
-		pgx.Identifier{schema, table}.Sanitize(),
+		pgx.Identifier{params.Schema, params.Table}.Sanitize(),
 	)
 }
 
@@ -81,10 +104,9 @@ func prepareWhereClause(p *metadata.Placeholder, sqlFilters *types.Document) (st
 			return "", nil, lazyerrors.Error(err)
 		}
 
-		// Is the comment below correct? Does it also skip things like $or?
-		// TODO https://github.com/FerretDB/FerretDB/issues/3573
-
-		// don't pushdown $comment, it's attached to query in handlers
+		// don't pushdown $comment, as it's attached to query with select clause
+		//
+		// all of the other top-level operators such as `$or` do not support pushdown yet
 		if strings.HasPrefix(rootKey, "$") {
 			continue
 		}

diff --git a/internal/backends/postgresql/query_test.go b/internal/backends/postgresql/query_test.go
@@ -15,6 +15,7 @@
 package postgresql
 
 import (
+	"fmt"
 	"math"
 	"testing"
 	"time"
@@ -28,6 +29,67 @@ import (
 	"github.com/FerretDB/FerretDB/internal/util/must"
 )
 
+func TestPrepareSelectClause(t *testing.T) {
+	t.Parallel()
+	schema := "schema"
+	table := "table"
+	comment := "*/ 1; DROP SCHEMA " + schema + " CASCADE -- "
+
+	for name, tc := range map[string]struct { //nolint:vet // used for test only
+		capped        bool
+		onlyRecordIDs bool
+
+		expectQuery string
+	}{
+		"CappedRecordID": {
+			capped:        true,
+			onlyRecordIDs: true,
+			expectQuery: fmt.Sprintf(
+				`SELECT %s %s FROM "%s"."%s"`,
+				"/* * / 1; DROP SCHEMA "+schema+" CASCADE --  */",
+				metadata.RecordIDColumn,
+				schema,
+				table,
+			),
+		},
+		"Capped": {
+			capped: true,
+			expectQuery: fmt.Sprintf(
+				`SELECT %s %s, %s FROM "%s"."%s"`,
+				"/* * / 1; DROP SCHEMA "+schema+" CASCADE --  */",
+				metadata.RecordIDColumn,
+				metadata.DefaultColumn,
+				schema,
+				table,
+			),
+		},
+		"FullRecord": {
+			expectQuery: fmt.Sprintf(
+				`SELECT %s %s FROM "%s"."%s"`,
+				"/* * / 1; DROP SCHEMA "+schema+" CASCADE --  */",
+				metadata.DefaultColumn,
+				schema,
+				table,
+			),
+		},
+	} {
+		name, tc := name, tc
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			query := prepareSelectClause(&selectParams{
+				Schema:        schema,
+				Table:         table,
+				Comment:       comment,
+				Capped:        tc.capped,
+				OnlyRecordIDs: tc.onlyRecordIDs,
+			})
+
+			assert.Equal(t, tc.expectQuery, query)
+		})
+	}
+}
+
 func TestPrepareWhereClause(t *testing.T) {
 	t.Parallel()
 	objectID := types.ObjectID{0x62, 0x56, 0xc5, 0xba, 0x0b, 0xad, 0xc0, 0xff, 0xee, 0xff, 0xff, 0xff}

diff --git a/internal/backends/sqlite/collection.go b/internal/backends/sqlite/collection.go
@@ -83,7 +83,7 @@ func (c *collection) Query(ctx context.Context, params *backends.QueryParams) (*
 		}
 	}
 
-	q := prepareSelectClause(meta.TableName, meta.Capped(), params.OnlyRecordIDs) + whereClause
+	q := prepareSelectClause(meta.TableName, params.Comment, meta.Capped(), params.OnlyRecordIDs) + whereClause
 
 	q += prepareOrderByClause(params.Sort, meta.Capped())
 
@@ -253,7 +253,7 @@ func (c *collection) Explain(ctx context.Context, params *backends.ExplainParams
 		params = new(backends.ExplainParams)
 	}
 
-	selectClause := prepareSelectClause(meta.TableName, meta.Capped(), false)
+	selectClause := prepareSelectClause(meta.TableName, "", meta.Capped(), false)
 
 	var queryPushdown bool
 	var whereClause string

diff --git a/internal/backends/sqlite/query.go b/internal/backends/sqlite/query.go
@@ -16,6 +16,7 @@ package sqlite
 
 import (
 	"fmt"
+	"strings"
 
 	"github.com/FerretDB/FerretDB/internal/backends"
 	"github.com/FerretDB/FerretDB/internal/backends/sqlite/metadata"
@@ -26,16 +27,22 @@ import (
 // For capped collection with onlyRecordIDs, it returns select clause for recordID column.
 //
 // For capped collection, it returns select clause for recordID column and default column.
-func prepareSelectClause(table string, capped, onlyRecordIDs bool) string {
+func prepareSelectClause(table, comment string, capped, onlyRecordIDs bool) string {
+	if comment != "" {
+		comment = strings.ReplaceAll(comment, "/*", "/ *")
+		comment = strings.ReplaceAll(comment, "*/", "* /")
+		comment = `/* ` + comment + ` */`
+	}
+
 	if capped && onlyRecordIDs {
-		return fmt.Sprintf(`SELECT %s FROM %q`, metadata.RecordIDColumn, table)
+		return fmt.Sprintf(`SELECT %s %s FROM %q`, comment, metadata.RecordIDColumn, table)
 	}
 
 	if capped {
-		return fmt.Sprintf(`SELECT %s, %s FROM %q`, metadata.RecordIDColumn, metadata.DefaultColumn, table)
+		return fmt.Sprintf(`SELECT %s %s, %s FROM %q`, comment, metadata.RecordIDColumn, metadata.DefaultColumn, table)
 	}
 
-	return fmt.Sprintf(`SELECT %s FROM %q`, metadata.DefaultColumn, table)
+	return fmt.Sprintf(`SELECT %s %s FROM %q`, comment, metadata.DefaultColumn, table)
 }
 
 // prepareOrderByClause returns ORDER BY clause.

diff --git a/internal/backends/sqlite/query_test.go b/internal/backends/sqlite/query_test.go
@@ -15,13 +15,65 @@
 package sqlite
 
 import (
+	"fmt"
 	"testing"
 
 	"github.com/stretchr/testify/assert"
 
 	"github.com/FerretDB/FerretDB/internal/backends"
+	"github.com/FerretDB/FerretDB/internal/backends/sqlite/metadata"
 )
 
+func TestPrepareSelectClause(t *testing.T) {
+	t.Parallel()
+	table := "table"
+	comment := "*/ 1; DROP TABLE " + table + " CASCADE -- "
+
+	for name, tc := range map[string]struct { //nolint:vet // used for test only
+		capped        bool
+		onlyRecordIDs bool
+
+		expectQuery string
+	}{
+		"CappedRecordID": {
+			capped:        true,
+			onlyRecordIDs: true,
+			expectQuery: fmt.Sprintf(
+				`SELECT %s %s FROM %q`,
+				"/* * / 1; DROP TABLE "+table+" CASCADE --  */",
+				metadata.RecordIDColumn,
+				table,
+			),
+		},
+		"Capped": {
+			capped: true,
+			expectQuery: fmt.Sprintf(
+				`SELECT %s %s, %s FROM %q`,
+				"/* * / 1; DROP TABLE "+table+" CASCADE --  */",
+				metadata.RecordIDColumn,
+				metadata.DefaultColumn,
+				table,
+			),
+		},
+		"FullRecord": {
+			expectQuery: fmt.Sprintf(
+				`SELECT %s %s FROM %q`,
+				"/* * / 1; DROP TABLE "+table+" CASCADE --  */",
+				metadata.DefaultColumn,
+				table,
+			),
+		},
+	} {
+		name, tc := name, tc
+		t.Run(name, func(t *testing.T) {
+			t.Parallel()
+
+			query := prepareSelectClause(table, comment, tc.capped, tc.onlyRecordIDs)
+			assert.Equal(t, tc.expectQuery, query)
+		})
+	}
+}
+
 func TestPrepareOrderByClause(t *testing.T) {
 	t.Parallel()