Merge pull request #934 from oom-ai/jinghan/export_stream

Jinghan/implement API `ChannelExportStream`

Author: jinghancc

internal/database/offline/types.go

@@ -8,6 +8,8 @@ import (
 
 type ExportOpt struct {
 	SnapshotTable string
+	CdcTable      *string
+	UnixMilli     *int64
 	EntityName    string
 	Features      types.FeatureList
 	Limit         *uint64

pkg/oomstore/export.go

@@ -6,6 +6,8 @@ import (
 	"fmt"
 	"os"
 
+	"github.com/oom-ai/oomstore/internal/database/dbutil"
+
 	"github.com/pkg/errors"
 	"github.com/spf13/cast"
 
@@ -14,7 +16,7 @@ import (
 )
 
 /*
-Export feature values of a particular revision.
+ChannelExport exports batch feature values of a particular revision.
 Usage Example:
 	exportResult, err := store.Export(ctx, opt)
 	if err != nil {
@@ -93,3 +95,48 @@ func (s *OomStore) Export(ctx context.Context, opt types.ExportOpt) error {
 	}
 	return exportResult.CheckStreamError()
 }
+
+// ChannelExportStream exports the latest streaming feature values up to the given timestamp.
+// Currently, this API can only export features in one feature group.
+func (s *OomStore) ChannelExportStream(ctx context.Context, opt types.ChannelExportStreamOpt) (*types.ExportResult, error) {
+	if err := validateFeatureFullNames(opt.FeatureFullNames); err != nil {
+		return nil, errors.WithStack(err)
+	}
+	features, err := s.ListFeature(ctx, types.ListFeatureOpt{
+		FeatureFullNames: &opt.FeatureFullNames,
+	})
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	if len(features.GroupIDs()) != 1 {
+		return nil, fmt.Errorf("expected 1 group, got %d groups", len(features.GroupIDs()))
+	}
+	group := features[0].Group
+	revisions, err := s.ListRevision(ctx, &group.ID)
+	if err != nil {
+		return nil, errors.WithStack(err)
+	}
+	revision := revisions.Before(opt.UnixMilli)
+	if revision == nil {
+		return nil, fmt.Errorf("no feature values up to %d, use a later timestamp", opt.UnixMilli)
+	}
+	if revision.SnapshotTable == "" {
+		if err = s.Snapshot(ctx, group.Name); err != nil {
+			return nil, errors.WithStack(err)
+		}
+	}
+
+	snapshotTable := dbutil.OfflineStreamSnapshotTableName(group.ID, revision.Revision)
+	cdcTable := dbutil.OfflineStreamCdcTableName(group.ID, revision.Revision)
+	stream, exportErr := s.offline.Export(ctx, offline.ExportOpt{
+		SnapshotTable: snapshotTable,
+		CdcTable:      &cdcTable,
+		UnixMilli:     &opt.UnixMilli,
+		EntityName:    group.Entity.Name,
+		Features:      features,
+		Limit:         opt.Limit,
+	})
+
+	header := append([]string{group.Entity.Name}, features.Names()...)
+	return types.NewExportResult(header, stream, exportErr), nil
+}

pkg/oomstore/types/options.go

@@ -36,6 +36,12 @@ type ChannelExportOpt struct {
 	Limit        *uint64
 }
 
+type ChannelExportStreamOpt struct {
+	UnixMilli        int64
+	FeatureFullNames []string
+	Limit            *uint64
+}
+
 type ExportOpt struct {
 	RevisionID     int
 	FeatureNames   []string

pkg/oomstore/types/revision.go

@@ -1,6 +1,7 @@
 package types
 
 import (
+	"sort"
 	"time"
 )
 
@@ -62,6 +63,25 @@ func (l RevisionList) Filter(filter func(*Revision) bool) (rs RevisionList) {
 	return
 }
 
+func (l RevisionList) Before(unixMilli int64) *Revision {
+	if len(l) == 0 {
+		return nil
+	}
+	sort.Slice(l, func(i, j int) bool {
+		return l[i].Revision < l[j].Revision
+	})
+	if l[0].Revision > unixMilli {
+		return nil
+	}
+	var i int
+	for i = range l {
+		if l[i].Revision > unixMilli {
+			break
+		}
+	}
+	return l[i-1]
+}
+
 func (l RevisionList) GroupIDs() []int {
 	groupIDMap := make(map[int]struct{})
 	for _, r := range l {