unified-storage: add BatchGet support to the sqlkv implementation (#115517)

* unified-storage: add `BatchGet` support to the sqlkv implementation

* address comments

* fix linting

pkg/storage/unified/resource/sqlkv.go

@@ -34,9 +34,10 @@ func mustTemplate(filename string) *template.Template {
 
 // Templates.
 var (
-	sqlKVGet    = mustTemplate("sqlkv_get.sql")
-	sqlKVDelete = mustTemplate("sqlkv_delete.sql")
-	sqlKVKeys   = mustTemplate("sqlkv_keys.sql")
+	sqlKVKeys     = mustTemplate("sqlkv_keys.sql")
+	sqlKVGet      = mustTemplate("sqlkv_get.sql")
+	sqlKVBatchGet = mustTemplate("sqlkv_batch_get.sql")
+	sqlKVDelete   = mustTemplate("sqlkv_delete.sql")
 )
 
 // sqlKVSection can be embedded in structs used when rendering query templates
@@ -107,13 +108,23 @@ func (req sqlKVGetRequest) Results() ([]byte, error) {
 	return req.Value, nil
 }
 
-type sqlKVDeleteRequest struct {
+type sqlKVBatchGetRequest struct {
 	sqltemplate.SQLTemplate
-	sqlKVSectionKey
+	sqlKVSection
+	Keys []string
 }
 
-func (req sqlKVDeleteRequest) Validate() error {
-	return req.sqlKVSectionKey.Validate()
+func (req sqlKVBatchGetRequest) Validate() error {
+	return req.sqlKVSection.Validate()
+}
+
+func (req sqlKVBatchGetRequest) KeyPaths() []string {
+	result := make([]string, 0, len(req.Keys))
+	for _, key := range req.Keys {
+		result = append(result, req.Section+"/"+key)
+	}
+
+	return result
 }
 
 type sqlKVKeysRequest struct {
@@ -142,6 +153,15 @@ func (req sqlKVKeysRequest) SortAscending() bool {
 	return req.Options.Sort != SortOrderDesc
 }
 
+type sqlKVDeleteRequest struct {
+	sqltemplate.SQLTemplate
+	sqlKVSectionKey
+}
+
+func (req sqlKVDeleteRequest) Validate() error {
+	return req.sqlKVSectionKey.Validate()
+}
+
 var _ KV = &sqlKV{}
 
 type sqlKV struct {
@@ -188,6 +208,7 @@ func (k *sqlKV) Keys(ctx context.Context, section string, opt ListOptions) iter.
 			yield("", err)
 			return
 		}
+		defer closeRows(rows, yield)
 
 		for rows.Next() {
 			var key string
@@ -225,7 +246,41 @@ func (k *sqlKV) Get(ctx context.Context, section string, key string) (io.ReadClo
 
 func (k *sqlKV) BatchGet(ctx context.Context, section string, keys []string) iter.Seq2[KeyValue, error] {
 	return func(yield func(KeyValue, error) bool) {
-		panic("not implemented!")
+		if len(keys) == 0 {
+			return
+		}
+
+		rows, err := dbutil.QueryRows(ctx, k.db, sqlKVBatchGet, sqlKVBatchGetRequest{
+			SQLTemplate:  sqltemplate.New(k.dialect),
+			sqlKVSection: sqlKVSection{section},
+			Keys:         keys,
+		})
+		if err != nil {
+			yield(KeyValue{}, err)
+			return
+		}
+		defer closeRows(rows, yield)
+
+		for rows.Next() {
+			var key string
+			var value []byte
+			if err := rows.Scan(&key, &value); err != nil {
+				yield(KeyValue{}, fmt.Errorf("error reading row: %w", err))
+				return
+			}
+
+			kv := KeyValue{
+				Key:   strings.TrimPrefix(key, section+"/"),
+				Value: io.NopCloser(bytes.NewReader(value)),
+			}
+			if !yield(kv, nil) {
+				return
+			}
+		}
+
+		if err := rows.Err(); err != nil {
+			yield(KeyValue{}, fmt.Errorf("failed to read rows: %w", err))
+		}
 	}
 }
 
@@ -273,3 +328,10 @@ func (k *sqlKV) BatchDelete(ctx context.Context, section string, keys []string)
 func (k *sqlKV) UnixTimestamp(ctx context.Context) (int64, error) {
 	panic("not implemented!")
 }
+
+func closeRows[T any](rows db.Rows, yield func(T, error) bool) {
+	if err := rows.Close(); err != nil {
+		var zero T
+		yield(zero, fmt.Errorf("error closing rows: %w", err))
+	}
+}