Merge pull request #2502 from keboola/vb/DMD-922/llm-export-samples-security

DMD-922 - Add samples for llm export

Author: vojtabiberle

internal/pkg/llm/twinformat/fetcher.go

@@ -4,13 +4,17 @@ import (
 	"context"
 	"sort"
 	"strings"
+	"sync"
 	"time"
 
 	"github.com/keboola/keboola-sdk-go/v2/pkg/keboola"
+	"golang.org/x/sync/errgroup"
+	"golang.org/x/sync/semaphore"
 
 	"github.com/keboola/keboola-as-code/internal/pkg/llm/twinformat/configparser"
 	"github.com/keboola/keboola-as-code/internal/pkg/log"
 	"github.com/keboola/keboola-as-code/internal/pkg/telemetry"
+	"github.com/keboola/keboola-as-code/internal/pkg/utils/errors"
 )
 
 // FetcherDependencies defines the dependencies required by the Fetcher.
@@ -183,6 +187,162 @@ type componentsResult struct {
 	ComponentConfigs []*configparser.ComponentConfig
 }
 
+// TableSample represents a sample of table data.
+type TableSample struct {
+	TableID keboola.TableID
+	Columns []string
+	Rows    [][]string
+}
+
+// RowCount returns the number of rows in the sample.
+func (s *TableSample) RowCount() int {
+	return len(s.Rows)
+}
+
+// FetchTableSample fetches a sample of data from a table.
+func (f *Fetcher) FetchTableSample(ctx context.Context, tableKey keboola.TableKey, limit uint) (sample *TableSample, err error) {
+	ctx, span := f.telemetry.Tracer().Start(ctx, "keboola.go.twinformat.fetcher.FetchTableSample")
+	defer span.End(&err)
+
+	f.logger.Debugf(ctx, "Fetching sample for table %s (limit: %d)", tableKey.TableID, limit)
+
+	// Fetch table preview using the SDK.
+	// Don't pass WithLimitRows when limit is 0, as it may have different semantics than omitting.
+	var preview *keboola.TablePreview
+	if limit == 0 {
+		preview, err = f.api.PreviewTableRequest(tableKey).Send(ctx)
+	} else {
+		preview, err = f.api.PreviewTableRequest(tableKey, keboola.WithLimitRows(limit)).Send(ctx)
+	}
+	if err != nil {
+		return nil, err
+	}
+
+	sample = &TableSample{
+		TableID: tableKey.TableID,
+		Columns: preview.Columns,
+		Rows:    preview.Rows,
+	}
+
+	f.logger.Debugf(ctx, "Fetched %d rows for table %s", sample.RowCount(), tableKey.TableID)
+
+	return sample, nil
+}
+
+// indexedSample pairs a sample with its original index for ordering.
+type indexedSample struct {
+	index  int
+	sample *TableSample
+}
+
+// sampleFetchCollector collects results from concurrent sample fetching.
+type sampleFetchCollector struct {
+	mu          sync.Mutex
+	results     []indexedSample
+	failedCount int
+}
+
+// recordFailure increments the failure count.
+func (c *sampleFetchCollector) recordFailure() {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.failedCount++
+}
+
+// recordSuccess records a successfully fetched sample.
+func (c *sampleFetchCollector) recordSuccess(idx int, sample *TableSample) {
+	c.mu.Lock()
+	defer c.mu.Unlock()
+	c.results = append(c.results, indexedSample{index: idx, sample: sample})
+}
+
+// FetchTableSamples fetches samples for multiple tables concurrently.
+func (f *Fetcher) FetchTableSamples(ctx context.Context, tables []*keboola.Table, limit uint, maxTables int) (samples []*TableSample, err error) {
+	ctx, span := f.telemetry.Tracer().Start(ctx, "keboola.go.twinformat.fetcher.FetchTableSamples")
+	defer span.End(&err)
+
+	// Guard against non-positive maxTables to avoid panics from negative slice capacities.
+	if maxTables <= 0 {
+		return []*TableSample{}, nil
+	}
+
+	// Limit tables to fetch.
+	tablesToFetch := tables
+	if len(tablesToFetch) > maxTables {
+		tablesToFetch = tablesToFetch[:maxTables]
+	}
+
+	f.logger.Infof(ctx, "Fetching samples for %d tables concurrently (limit: %d rows each)", len(tablesToFetch), limit)
+
+	// Use bounded concurrency to respect API rate limits.
+	const maxConcurrency = 10
+
+	sem := semaphore.NewWeighted(maxConcurrency)
+	group, groupCtx := errgroup.WithContext(ctx)
+	collector := &sampleFetchCollector{
+		results: make([]indexedSample, 0, len(tablesToFetch)),
+	}
+
+	for i, table := range tablesToFetch {
+		idx, t := i, table
+		group.Go(func() error {
+			// Acquire semaphore - blocks until slot available or context cancelled.
+			if err := sem.Acquire(groupCtx, 1); err != nil {
+				collector.recordFailure()
+				return nil // Don't propagate - we want partial results
+			}
+			defer sem.Release(1)
+
+			f.fetchTableSample(groupCtx, collector, idx, t, limit)
+			return nil
+		})
+	}
+
+	// Wait for all goroutines to complete.
+	_ = group.Wait()
+
+	// If context was cancelled and no samples were fetched, propagate the cancellation error.
+	if len(collector.results) == 0 && ctx.Err() != nil {
+		return nil, ctx.Err()
+	}
+
+	// Sort by original index to preserve order.
+	sort.Slice(collector.results, func(i, j int) bool {
+		return collector.results[i].index < collector.results[j].index
+	})
+
+	samples = make([]*TableSample, 0, len(collector.results))
+	for _, r := range collector.results {
+		samples = append(samples, r.sample)
+	}
+
+	f.logger.Infof(ctx, "Fetched samples for %d tables (%d failed)", len(samples), collector.failedCount)
+
+	// Return error if any tables failed to fetch, but still return partial results.
+	if collector.failedCount > 0 {
+		return samples, errors.Errorf("failed to fetch samples for %d of %d tables", collector.failedCount, len(tablesToFetch))
+	}
+
+	return samples, nil
+}
+
+// fetchTableSample fetches a sample for a single table and records the result.
+func (f *Fetcher) fetchTableSample(ctx context.Context, collector *sampleFetchCollector, idx int, t *keboola.Table, limit uint) {
+	tableKey := keboola.TableKey{
+		BranchID: t.BranchID,
+		TableID:  t.TableID,
+	}
+
+	sample, err := f.FetchTableSample(ctx, tableKey, limit)
+	if err != nil {
+		f.logger.Warnf(ctx, "Failed to fetch sample for table %s: %v", t.TableID, err)
+		collector.recordFailure()
+		return
+	}
+
+	collector.recordSuccess(idx, sample)
+}
+
 // fetchAllComponents fetches all components and extracts transformation and component configs.
 // This makes a single API call and returns all data needed for processing.
 func (f *Fetcher) fetchAllComponents(ctx context.Context, branchID keboola.BranchID) (result *componentsResult, err error) {

internal/pkg/llm/twinformat/fetcher_test.go

@@ -475,3 +475,149 @@ func TestFetchTableLastImporter_NoImportEvents(t *testing.T) {
 	// No import events, should return empty
 	assert.Empty(t, componentID)
 }
+
+func TestFetchTableSample(t *testing.T) {
+	t.Parallel()
+
+	fetcher, transport := newTestFetcher(t)
+	branchID := keboola.BranchID(123)
+
+	tableKey := keboola.TableKey{
+		BranchID: branchID,
+		TableID: keboola.TableID{
+			BucketID:  keboola.BucketID{Stage: keboola.BucketStageIn, BucketName: "test"},
+			TableName: "sample_table",
+		},
+	}
+
+	// Mock table preview response - Storage API returns CSV format
+	csvData := "\"id\",\"name\",\"value\"\n\"1\",\"Alice\",\"100\"\n\"2\",\"Bob\",\"200\"\n\"3\",\"Charlie\",\"300\"\n"
+	transport.RegisterResponder(
+		http.MethodGet,
+		`=~/v2/storage/branch/123/tables/in.test.sample_table/data-preview`,
+		httpmock.NewStringResponder(200, csvData),
+	)
+
+	sample, err := fetcher.FetchTableSample(t.Context(), tableKey, 100)
+	require.NoError(t, err)
+
+	assert.Equal(t, tableKey.TableID, sample.TableID)
+	assert.Equal(t, []string{"id", "name", "value"}, sample.Columns)
+	assert.Len(t, sample.Rows, 3)
+	assert.Equal(t, 3, sample.RowCount())
+	assert.Equal(t, []string{"1", "Alice", "100"}, sample.Rows[0])
+}
+
+func TestFetchTableSamples(t *testing.T) {
+	t.Parallel()
+
+	fetcher, transport := newTestFetcher(t)
+	branchID := keboola.BranchID(123)
+
+	tables := []*keboola.Table{
+		{
+			TableKey: keboola.TableKey{
+				BranchID: branchID,
+				TableID: keboola.TableID{
+					BucketID:  keboola.BucketID{Stage: keboola.BucketStageIn, BucketName: "bucket"},
+					TableName: "table1",
+				},
+			},
+		},
+		{
+			TableKey: keboola.TableKey{
+				BranchID: branchID,
+				TableID: keboola.TableID{
+					BucketID:  keboola.BucketID{Stage: keboola.BucketStageIn, BucketName: "bucket"},
+					TableName: "table2",
+				},
+			},
+		},
+		{
+			TableKey: keboola.TableKey{
+				BranchID: branchID,
+				TableID: keboola.TableID{
+					BucketID:  keboola.BucketID{Stage: keboola.BucketStageIn, BucketName: "bucket"},
+					TableName: "table3",
+				},
+			},
+		},
+	}
+
+	// Mock preview responses for each table - Storage API returns CSV format
+	transport.RegisterResponder(
+		http.MethodGet,
+		`=~/v2/storage/branch/123/tables/in.bucket.table1/data-preview`,
+		httpmock.NewStringResponder(200, "\"col1\"\n\"val1\"\n"),
+	)
+	transport.RegisterResponder(
+		http.MethodGet,
+		`=~/v2/storage/branch/123/tables/in.bucket.table2/data-preview`,
+		httpmock.NewStringResponder(200, "\"col2\"\n\"val2\"\n"),
+	)
+	transport.RegisterResponder(
+		http.MethodGet,
+		`=~/v2/storage/branch/123/tables/in.bucket.table3/data-preview`,
+		httpmock.NewStringResponder(200, "\"col3\"\n\"val3\"\n"),
+	)
+
+	// Test with maxTables=2 (should only fetch first 2 tables)
+	samples, err := fetcher.FetchTableSamples(t.Context(), tables, 100, 2)
+	require.NoError(t, err)
+
+	assert.Len(t, samples, 2)
+	assert.Equal(t, "in.bucket.table1", samples[0].TableID.String())
+	assert.Equal(t, "in.bucket.table2", samples[1].TableID.String())
+}
+
+func TestFetchTableSamples_SkipsFailedTables(t *testing.T) {
+	t.Parallel()
+
+	fetcher, transport := newTestFetcher(t)
+	branchID := keboola.BranchID(123)
+
+	tables := []*keboola.Table{
+		{
+			TableKey: keboola.TableKey{
+				BranchID: branchID,
+				TableID: keboola.TableID{
+					BucketID:  keboola.BucketID{Stage: keboola.BucketStageIn, BucketName: "bucket"},
+					TableName: "good_table",
+				},
+			},
+		},
+		{
+			TableKey: keboola.TableKey{
+				BranchID: branchID,
+				TableID: keboola.TableID{
+					BucketID:  keboola.BucketID{Stage: keboola.BucketStageIn, BucketName: "bucket"},
+					TableName: "bad_table",
+				},
+			},
+		},
+	}
+
+	// First table succeeds - Storage API returns CSV format
+	transport.RegisterResponder(
+		http.MethodGet,
+		`=~/v2/storage/branch/123/tables/in.bucket.good_table/data-preview`,
+		httpmock.NewStringResponder(200, "\"col1\"\n\"val1\"\n"),
+	)
+
+	// Second table fails
+	transport.RegisterResponder(
+		http.MethodGet,
+		`=~/v2/storage/branch/123/tables/in.bucket.bad_table/data-preview`,
+		httpmock.NewStringResponder(500, "Internal Server Error"),
+	)
+
+	samples, err := fetcher.FetchTableSamples(t.Context(), tables, 100, 10)
+
+	// Should return error indicating partial failure, but still return partial results.
+	require.Error(t, err)
+	assert.Contains(t, err.Error(), "failed to fetch samples for 1 of 2 tables")
+
+	// Should have 1 sample (failed table is skipped)
+	assert.Len(t, samples, 1)
+	assert.Equal(t, "in.bucket.good_table", samples[0].TableID.String())
+}