Storage: Support continue at specified resource version (#84868)

* support continue at specified resource version * detect whether list continue pages need to use entity_history, remove BatchRead, expand selectQuery helper * refactor continue token handling * fix tests, increase history chunk size * lint fix
2024-04-08 09:57:02 -04:00
parent d0312af571
commit 9c46e71d5a
6 changed files with 427 additions and 530 deletions
@@ -21,57 +21,115 @@ func (d Direction) String() string {
 	return "ASC"
 }

+type joinQuery struct {
+	query string
+	args  []any
+}
+
+type whereClause struct {
+	query string
+	args  []any
+}
+
 type selectQuery struct {
 	dialect  migrator.Dialect
-	fields   []string // SELECT xyz
-	from     string   // FROM object
+	fields   []string    // SELECT xyz
+	from     string      // FROM object
+	joins    []joinQuery // JOIN object
 	offset   int64
 	limit    int64
 	oneExtra bool

-	where []string
-	args  []any
+	where []whereClause
+
+	groupBy []string

 	orderBy   []string
 	direction []Direction
 }

-func (q *selectQuery) addWhere(f string, val ...any) {
-	q.args = append(q.args, val...)
-	// if the field contains a question mark, we assume it's a raw where clause
-	if strings.Contains(f, "?") {
-		q.where = append(q.where, f)
-		// otherwise we assume it's a field name
-	} else {
-		q.where = append(q.where, q.dialect.Quote(f)+"=?")
+func NewSelectQuery(dialect migrator.Dialect, from string) *selectQuery {
+	return &selectQuery{
+		dialect: dialect,
+		from:    from,
 	}
 }

-func (q *selectQuery) addWhereInSubquery(f string, subquery string, subqueryArgs []any) {
-	q.args = append(q.args, subqueryArgs...)
-	q.where = append(q.where, q.dialect.Quote(f)+" IN ("+subquery+")")
+func (q *selectQuery) From(from string) {
+	q.from = from
 }

-func (q *selectQuery) addWhereIn(f string, vals []string) {
+func (q *selectQuery) SetLimit(limit int64) {
+	q.limit = limit
+}
+
+func (q *selectQuery) SetOffset(offset int64) {
+	q.offset = offset
+}
+
+func (q *selectQuery) SetOneExtra() {
+	q.oneExtra = true
+}
+
+func (q *selectQuery) UnsetOneExtra() {
+	q.oneExtra = false
+}
+
+func (q *selectQuery) AddFields(f ...string) {
+	for _, field := range f {
+		q.fields = append(q.fields, "t."+q.dialect.Quote(field))
+	}
+}
+
+func (q *selectQuery) AddRawFields(f ...string) {
+	q.fields = append(q.fields, f...)
+}
+
+func (q *selectQuery) AddJoin(j string, args ...any) {
+	q.joins = append(q.joins, joinQuery{query: j, args: args})
+}
+
+func (q *selectQuery) AddWhere(f string, val ...any) {
+	// if the field contains a question mark, we assume it's a raw where clause
+	if strings.Contains(f, "?") {
+		q.where = append(q.where, whereClause{f, val})
+		// otherwise we assume it's a field name
+	} else {
+		q.where = append(q.where, whereClause{"t." + q.dialect.Quote(f) + "=?", val})
+	}
+}
+
+func (q *selectQuery) AddWhereInSubquery(f string, subquery string, subqueryArgs []any) {
+	q.where = append(q.where, whereClause{"t." + q.dialect.Quote(f) + " IN (" + subquery + ")", subqueryArgs})
+}
+
+func (q *selectQuery) AddWhereIn(f string, vals []any) {
 	count := len(vals)
 	if count > 1 {
 		sb := strings.Builder{}
-		sb.WriteString(q.dialect.Quote(f))
+		sb.WriteString("t." + q.dialect.Quote(f))
 		sb.WriteString(" IN (")
 		for i := 0; i < count; i++ {
 			if i > 0 {
 				sb.WriteString(",")
 			}
 			sb.WriteString("?")
-			q.args = append(q.args, vals[i])
 		}
 		sb.WriteString(") ")
-		q.where = append(q.where, sb.String())
+		q.where = append(q.where, whereClause{sb.String(), vals})
 	} else if count == 1 {
-		q.addWhere(f, vals[0])
+		q.AddWhere(f, vals[0])
 	}
 }

+func ToAnyList[T any](input []T) []any {
+	list := make([]any, len(input))
+	for i, v := range input {
+		list[i] = v
+	}
+	return list
+}
+
 const sqlLikeEscape = "#"

 var sqlLikeEscapeReplacer = strings.NewReplacer(
@@ -85,39 +143,57 @@ func escapeJSONStringSQLLike(s string) string {
 	return sqlLikeEscapeReplacer.Replace(string(b))
 }

-func (q *selectQuery) addWhereJsonContainsKV(field string, key string, value string) {
+func (q *selectQuery) AddWhereJsonContainsKV(field string, key string, value string) {
 	escapedKey := escapeJSONStringSQLLike(key)
 	escapedValue := escapeJSONStringSQLLike(value)
-	q.where = append(q.where, q.dialect.Quote(field)+" LIKE ? ESCAPE ?")
-	q.args = append(q.args, "{%"+escapedKey+":"+escapedValue+"%}", sqlLikeEscape)
+	q.where = append(q.where, whereClause{
+		"t." + q.dialect.Quote(field) + " LIKE ? ESCAPE ?",
+		[]any{"{%\"" + escapedKey + "\":\"" + escapedValue + "\"%}", sqlLikeEscape},
+	})
 }

-func (q *selectQuery) addOrderBy(field string, direction Direction) {
+func (q *selectQuery) AddGroupBy(f string) {
+	q.groupBy = append(q.groupBy, f)
+}
+
+func (q *selectQuery) AddOrderBy(field string, direction Direction) {
 	q.orderBy = append(q.orderBy, field)
 	q.direction = append(q.direction, direction)
 }

-func (q *selectQuery) toQuery() (string, []any) {
-	args := q.args
+func (q *selectQuery) ToQuery() (string, []any) {
+	args := []any{}
 	sb := strings.Builder{}
 	sb.WriteString("SELECT ")
-	quotedFields := make([]string, len(q.fields))
-	for i, f := range q.fields {
-		quotedFields[i] = q.dialect.Quote(f)
-	}
-	sb.WriteString(strings.Join(quotedFields, ","))
+	sb.WriteString(strings.Join(q.fields, ","))
 	sb.WriteString(" FROM ")
 	sb.WriteString(q.from)
+	sb.WriteString(" AS t")
+
+	for _, j := range q.joins {
+		sb.WriteString(" " + j.query)
+		args = append(args, j.args...)
+	}

 	// Templated where string
-	where := len(q.where)
-	if where > 0 {
+	if len(q.where) > 0 {
 		sb.WriteString(" WHERE ")
-		for i := 0; i < where; i++ {
+		for i, w := range q.where {
 			if i > 0 {
 				sb.WriteString(" AND ")
 			}
-			sb.WriteString(q.where[i])
+			sb.WriteString(w.query)
+			args = append(args, w.args...)
+		}
+	}
+
+	if len(q.groupBy) > 0 {
+		sb.WriteString(" GROUP BY ")
+		for i, f := range q.groupBy {
+			if i > 0 {
+				sb.WriteString(",")
+			}
+			sb.WriteString("t." + q.dialect.Quote(f))
 		}
 	}

@@ -127,21 +203,19 @@ func (q *selectQuery) toQuery() (string, []any) {
 			if i > 0 {
 				sb.WriteString(",")
 			}
-			sb.WriteString(q.dialect.Quote(f))
+			sb.WriteString("t." + q.dialect.Quote(f))
 			sb.WriteString(" ")
 			sb.WriteString(q.direction[i].String())
 		}
 	}

 	limit := q.limit
-	if limit < 1 {
-		limit = 20
-		q.limit = limit
+	if limit > 0 {
+		if q.oneExtra {
+			limit = limit + 1
+		}
+		sb.WriteString(q.dialect.LimitOffset(limit, q.offset))
 	}
-	if q.oneExtra {
-		limit = limit + 1
-	}
-	sb.WriteString(q.dialect.LimitOffset(limit, q.offset))

 	return sb.String(), args
 }