diff --git a/pkg/postgres/pgutils/utils.go b/pkg/postgres/pgutils/utils.go index 24dcfcff9ba12..090da0ff9a0b4 100644 --- a/pkg/postgres/pgutils/utils.go +++ b/pkg/postgres/pgutils/utils.go @@ -13,6 +13,7 @@ import ( "github.com/stackrox/rox/pkg/logging" "github.com/stackrox/rox/pkg/postgres" "github.com/stackrox/rox/pkg/utils" + "github.com/stackrox/rox/pkg/uuid" "gorm.io/gorm" "gorm.io/gorm/schema" ) @@ -62,6 +63,18 @@ func NilOrTime(t *types.Timestamp) *time.Time { return &ts } +// NilOrUUID allows for a proto string to be stored as a UUID type in Postgres +func NilOrUUID(value string) *uuid.UUID { + if value == "" { + return nil + } + id, err := uuid.FromString(value) + if err != nil { + return nil + } + return &id +} + // CreateTableFromModel executes input create statement using the input connection. func CreateTableFromModel(ctx context.Context, db *gorm.DB, createStmt *postgres.CreateStmts) { err := Retry(func() error { diff --git a/pkg/postgres/walker/types.go b/pkg/postgres/walker/types.go index 443b77ded3c5e..0fe3c12674af7 100644 --- a/pkg/postgres/walker/types.go +++ b/pkg/postgres/walker/types.go @@ -17,6 +17,7 @@ const ( Integer DataType = "integer" IntArray DataType = "intarray" BigInteger DataType = "biginteger" + UUID DataType = "uuid" ) // DataTypeToSQLType converts the internal representation to SQL diff --git a/pkg/search/postgres/common.go b/pkg/search/postgres/common.go index 540a64a8c5862..bcc990bc081df 100644 --- a/pkg/search/postgres/common.go +++ b/pkg/search/postgres/common.go @@ -122,7 +122,7 @@ func (q *query) getPortionBeforeFromClause() string { var primaryKeyPaths []string // Always select the primary keys for count. for _, pk := range q.Schema.PrimaryKeys() { - primaryKeyPaths = append(primaryKeyPaths, qualifyColumn(pk.Schema.Table, pk.ColumnName)) + primaryKeyPaths = append(primaryKeyPaths, qualifyColumn(pk.Schema.Table, pk.ColumnName, "")) } countOn = fmt.Sprintf("distinct(%s)", strings.Join(primaryKeyPaths, ", ")) } @@ -133,7 +133,11 @@ func (q *query) getPortionBeforeFromClause() string { var primaryKeyPaths []string // Always select the primary keys first. for _, pk := range q.Schema.PrimaryKeys() { - primaryKeyPaths = append(primaryKeyPaths, qualifyColumn(pk.Schema.Table, pk.ColumnName)) + var cast string + if pk.SQLType == "uuid" { + cast = "::text" + } + primaryKeyPaths = append(primaryKeyPaths, qualifyColumn(pk.Schema.Table, pk.ColumnName, cast)) } primaryKeyPortion := strings.Join(primaryKeyPaths, ", ") @@ -185,7 +189,7 @@ func (q *query) AsSQL() string { primaryKeys := q.Schema.PrimaryKeys() primaryKeyPaths := make([]string, 0, len(primaryKeys)) for _, pk := range primaryKeys { - primaryKeyPaths = append(primaryKeyPaths, qualifyColumn(pk.Schema.Table, pk.ColumnName)) + primaryKeyPaths = append(primaryKeyPaths, qualifyColumn(pk.Schema.Table, pk.ColumnName, "")) } querySB.WriteString(" group by ") querySB.WriteString(strings.Join(primaryKeyPaths, ", ")) @@ -201,8 +205,8 @@ func (q *query) AsSQL() string { return querySB.String() } -func qualifyColumn(table, column string) string { - return table + "." + column +func qualifyColumn(table, column, cast string) string { + return table + "." + column + cast } type parsedPaginationQuery struct { @@ -245,9 +249,13 @@ func populatePagination(querySoFar *query, pagination *v1.QueryPagination, schem return errors.New("search after for pagination must be defined for only the first sort option") } if so.GetField() == searchPkg.DocID.String() { + var cast string + if schema.ID().SQLType == "uuid" { + cast = "::text" + } querySoFar.Pagination.OrderBys = append(querySoFar.Pagination.OrderBys, orderByEntry{ Field: pgsearch.SelectQueryField{ - SelectPath: qualifyColumn(schema.Table, schema.ID().ColumnName), + SelectPath: qualifyColumn(schema.Table, schema.ID().ColumnName, cast), FieldType: walker.String, }, Descending: so.GetReversed(), @@ -261,10 +269,15 @@ func populatePagination(querySoFar *query, pagination *v1.QueryPagination, schem return errors.Errorf("field %s does not exist in table %s or connected tables", so.GetField(), schema.Table) } + var cast string + if dbField.SQLType == "uuid" { + cast = "::text" + } + if fieldMetadata.derivedMetadata == nil { querySoFar.Pagination.OrderBys = append(querySoFar.Pagination.OrderBys, orderByEntry{ Field: pgsearch.SelectQueryField{ - SelectPath: qualifyColumn(dbField.Schema.Table, dbField.ColumnName), + SelectPath: qualifyColumn(dbField.Schema.Table, dbField.ColumnName, cast), FieldType: dbField.DataType, }, Descending: so.GetReversed(), @@ -275,7 +288,7 @@ func populatePagination(querySoFar *query, pagination *v1.QueryPagination, schem case searchPkg.CountDerivationType: querySoFar.Pagination.OrderBys = append(querySoFar.Pagination.OrderBys, orderByEntry{ Field: pgsearch.SelectQueryField{ - SelectPath: fmt.Sprintf("count(%s)", qualifyColumn(dbField.Schema.Table, dbField.ColumnName)), + SelectPath: fmt.Sprintf("count(%s)", qualifyColumn(dbField.Schema.Table, dbField.ColumnName, "")), FieldType: dbField.DataType, }, Descending: so.GetReversed(), @@ -285,7 +298,7 @@ func populatePagination(querySoFar *query, pagination *v1.QueryPagination, schem case searchPkg.SimpleReverseSortDerivationType: querySoFar.Pagination.OrderBys = append(querySoFar.Pagination.OrderBys, orderByEntry{ Field: pgsearch.SelectQueryField{ - SelectPath: qualifyColumn(dbField.Schema.Table, dbField.ColumnName), + SelectPath: qualifyColumn(dbField.Schema.Table, dbField.ColumnName, cast), FieldType: dbField.DataType, }, Descending: !so.GetReversed(), @@ -395,6 +408,7 @@ func combineQueryEntries(entries []*pgsearch.QueryEntry, separator string) *pgse if newQE.Having != nil { newQE.Having.Query = fmt.Sprintf("(%s)", strings.Join(havingQueryStrings, separator)) } + return newQE } @@ -423,8 +437,12 @@ func compileQueryToPostgres(schema *walker.Schema, q *v1.Query, queryFields map[ case *v1.Query_BaseQuery: switch subBQ := q.GetBaseQuery().Query.(type) { case *v1.BaseQuery_DocIdQuery: + cast := "::text[]" + if schema.ID().SQLType == "uuid" { + cast = "::uuid[]" + } return &pgsearch.QueryEntry{Where: pgsearch.WhereClause{ - Query: fmt.Sprintf("%s.%s = ANY($$::text[])", schema.Table, schema.ID().ColumnName), + Query: fmt.Sprintf("%s.%s = ANY($$%s)", schema.Table, schema.ID().ColumnName, cast), Values: []interface{}{subBQ.DocIdQuery.GetIds()}, }}, nil case *v1.BaseQuery_MatchFieldQuery: diff --git a/pkg/search/postgres/query/common.go b/pkg/search/postgres/query/common.go index ba6e24d7272d5..ef74dfc09d7bc 100644 --- a/pkg/search/postgres/query/common.go +++ b/pkg/search/postgres/query/common.go @@ -99,6 +99,9 @@ func MatchFieldQuery(dbField *walker.Field, derivedMetadata *walker.DerivedSearc qualifiedColName := dbField.Schema.Table + "." + dbField.ColumnName dataType := dbField.DataType + if dbField.SQLType == "uuid" { + dataType = walker.UUID + } var goesIntoHavingClause bool if derivedMetadata != nil { switch derivedMetadata.DerivationType { diff --git a/pkg/search/postgres/query/query_functions.go b/pkg/search/postgres/query/query_functions.go index 2d26fde776668..013275a0b37ef 100644 --- a/pkg/search/postgres/query/query_functions.go +++ b/pkg/search/postgres/query/query_functions.go @@ -24,8 +24,12 @@ type queryAndFieldContext struct { func qeWithSelectFieldIfNeeded(ctx *queryAndFieldContext, whereClause *WhereClause, postTransformFunc func(interface{}) interface{}) *QueryEntry { qe := &QueryEntry{Where: *whereClause} if ctx.highlight { + var cast string + if ctx.sqlDataType == walker.UUID { + cast = "::text" + } qe.SelectedFields = []SelectQueryField{{ - SelectPath: ctx.qualifiedColumnName, + SelectPath: ctx.qualifiedColumnName + cast, FieldType: ctx.sqlDataType, FieldPath: ctx.field.FieldPath, PostTransform: postTransformFunc, @@ -47,6 +51,7 @@ var datatypeToQueryFunc = map[walker.DataType]queryFunction{ walker.Numeric: newNumericQuery, walker.EnumArray: queryOnArray(newEnumQuery, getEnumArrayPostTransformFunc), walker.IntArray: queryOnArray(newNumericQuery, getIntArrayPostTransformFunc), + walker.UUID: newUUIDQuery, // Map is handled separately. } diff --git a/pkg/search/postgres/query/uuid_query.go b/pkg/search/postgres/query/uuid_query.go new file mode 100644 index 0000000000000..6c0f06369fd7f --- /dev/null +++ b/pkg/search/postgres/query/uuid_query.go @@ -0,0 +1,108 @@ +package pgsearch + +import ( + "fmt" + "regexp" + "strings" + + "github.com/pkg/errors" + pkgSearch "github.com/stackrox/rox/pkg/search" + "github.com/stackrox/rox/pkg/utils" + "github.com/stackrox/rox/pkg/uuid" +) + +func newUUIDQuery(ctx *queryAndFieldContext) (*QueryEntry, error) { + whereClause, err := newUUIDQueryWhereClause(ctx.qualifiedColumnName, ctx.value, ctx.queryModifiers...) + if err != nil { + return nil, err + } + return qeWithSelectFieldIfNeeded(ctx, &whereClause, nil), nil +} + +func newUUIDQueryWhereClause(columnName string, value string, queryModifiers ...pkgSearch.QueryModifier) (WhereClause, error) { + if len(value) == 0 { + return WhereClause{}, errors.New("value in search query cannot be empty") + } + + if value == pkgSearch.WildcardString { + return WhereClause{ + Query: fmt.Sprintf("%s is not null", columnName), + Values: []interface{}{}, + equivalentGoFunc: func(foundValue interface{}) bool { + foundVal := strings.ToLower(foundValue.(string)) + return foundVal != "" + }, + }, nil + } + + if len(queryModifiers) == 0 { + uuidVal, err := uuid.FromString(value) + if err != nil { + return WhereClause{}, errors.Wrapf(err, "value %q in search query must be valid UUID", value) + } + return WhereClause{ + Query: fmt.Sprintf("%s = $$", columnName), + Values: []interface{}{uuidVal}, + equivalentGoFunc: func(foundValue interface{}) bool { + return strings.EqualFold(foundValue.(string), value) + }, + }, nil + } + if queryModifiers[0] == pkgSearch.AtLeastOne { + panic("I dont think this is used") + } + var negationString string + negated := queryModifiers[0] == pkgSearch.Negation + if negated { + negationString = "!" + if len(queryModifiers) == 1 { + uuidVal, err := uuid.FromString(value) + if err != nil { + return WhereClause{}, errors.Wrapf(err, "value %q in search query must be valid UUID", value) + } + return WhereClause{ + Query: fmt.Sprintf("%s != $$", columnName), + Values: []interface{}{uuidVal}, + equivalentGoFunc: func(foundValue interface{}) bool { + return strings.EqualFold(foundValue.(string), value) != negated + }, + }, nil + } + queryModifiers = queryModifiers[1:] + } + + switch queryModifiers[0] { + case pkgSearch.Regex: + re, err := regexp.Compile(value) + if err != nil { + return WhereClause{}, fmt.Errorf("invalid regexp %s: %w", value, err) + } + + if value != "*" { + return WhereClause{}, fmt.Errorf("invalid regexp %s for UUID field", value) + } + return WhereClause{ + Query: fmt.Sprintf("%s is not null", columnName), + Values: []interface{}{value}, + equivalentGoFunc: func(foundValue interface{}) bool { + foundVal := strings.ToLower(foundValue.(string)) + return re.MatchString(foundVal) != negated + }, + }, nil + case pkgSearch.Equality: + uuidVal, err := uuid.FromString(value) + if err != nil { + return WhereClause{}, errors.Wrapf(err, "value %q in search query must be valid UUID", value) + } + return WhereClause{ + Query: fmt.Sprintf("%s %s= $$", columnName, negationString), + Values: []interface{}{uuidVal}, + equivalentGoFunc: func(foundValue interface{}) bool { + return strings.EqualFold(foundValue.(string), value) != negated + }, + }, nil + } + err := fmt.Errorf("unknown query modifier: %s", queryModifiers[0]) + utils.Should(err) + return WhereClause{}, err +} diff --git a/pkg/search/postgres/query_metadata.go b/pkg/search/postgres/query_metadata.go index 0ca00b16876e7..4734d12d78820 100644 --- a/pkg/search/postgres/query_metadata.go +++ b/pkg/search/postgres/query_metadata.go @@ -139,6 +139,14 @@ var ( return val.([]string) }, }, + walker.UUID: { + alloc: func() interface{} { + return pointers.String("") + }, + printer: func(val interface{}) []string { + return []string{*(val.(*string))} + }, + }, } )