snapshot.go

// Copyright 2026 Redpanda Data, Inc.
//
// Licensed as a Redpanda Enterprise file under the Redpanda Community
// License (the "License"); you may not use this file except in compliance with
// the License. You may obtain a copy of the License at
//
// https://github.com/redpanda-data/connect/blob/main/licenses/rcl.md

package replication

import (
	"context"
	"database/sql"
	"encoding/json"
	"errors"
	"fmt"
	"strings"
	"time"

	"golang.org/x/sync/errgroup"

	"github.com/redpanda-data/benthos/v4/public/service"
)

// Snapshot is responsible for creating snapshots of existing tables based on the Tables
// configuration value.
type Snapshot struct {
	dbPool                  *sql.DB
	tables                  []UserTable
	publisher               ChangePublisher
	log                     *service.Logger
	snapshotStatusMetric    *service.MetricGauge
	snapshotRowsTotalMetric *service.MetricCounter
	lobEnabled              bool
	pdbName                 string
}

// NewSnapshot creates a new instance of Snapshot capable of snapshotting provided tables.
// It does this by creating a transaction with snapshot level isolation before paging
// through rows, sending them to be batched.
func NewSnapshot(ctx context.Context,
	connectionString string,
	tables []UserTable,
	publisher ChangePublisher,
	lobEnabled bool,
	pdbName string,
	logger *service.Logger,
	metrics *service.Metrics,
) (*Snapshot, error) {
	db, err := sql.Open("oracle", connectionString)
	if err != nil {
		return nil, fmt.Errorf("connecting to oracle database for snapshotting: %w", err)
	}

	if err := ApplyNLSSettings(ctx, db); err != nil {
		db.Close()
		return nil, fmt.Errorf("configuring nls for snapshot session: %w", err)
	}

	s := &Snapshot{
		dbPool:                  db,
		tables:                  tables,
		publisher:               publisher,
		lobEnabled:              lobEnabled,
		pdbName:                 pdbName,
		snapshotStatusMetric:    metrics.NewGauge("oracledb_cdc_snapshot_status", "table"),
		snapshotRowsTotalMetric: metrics.NewCounter("oracledb_cdc_snapshot_rows_total", "table"),
		log:                     logger,
	}
	return s, nil
}

// Prepare prepares the snapshot by starting a transaction with appropriate isolation level.
// Returns the current SCN for the snapshot.
func (s *Snapshot) Prepare(ctx context.Context) (SCN, error) {
	if len(s.tables) == 0 {
		return InvalidSCN, errors.New("no tables provided")
	}

	var currentSCN SCN
	sql := `SELECT CURRENT_SCN FROM V$DATABASE`
	if err := s.dbPool.QueryRowContext(ctx, sql).Scan(&currentSCN); err != nil {
		return InvalidSCN, fmt.Errorf("getting current SCN for snapshot: %w", err)
	}

	s.log.Infof("Captured SCN before snapshot at SCN: %s", currentSCN)
	return currentSCN, nil
}

// Read launches N go routines (based on maxWorkers) and starts the process of
// iterating through each table, reading rows based on maxBatchSize, sending the row as a
// replication.MessageEvent to the configured publisher.
func (s *Snapshot) Read(ctx context.Context, maxWorkers, maxBatchSize int) error {
	s.log.Infof("Starting snapshot of %d table(s) using %d configured readers", len(s.tables), maxWorkers)

	for _, table := range s.tables {
		s.snapshotStatusMetric.Set(0, table.FullName())
	}

	wg, ctx := errgroup.WithContext(ctx)
	wg.SetLimit(maxWorkers)

	for _, table := range s.tables {
		wg.Go(s.snapshotTable(ctx, table, maxBatchSize))
	}

	if err := wg.Wait(); err != nil {
		return fmt.Errorf("processing snapshots: %w", err)
	}

	return nil
}

// snapshotTable is responsible for managing the entire process of replicating
// data from the table specified.
func (s *Snapshot) snapshotTable(ctx context.Context, table UserTable, maxBatchSize int) func() error {
	return func() error {
		var (
			err       error
			tx        *sql.Tx
			tableName = table.FullName()
		)
		l := s.log.With("src_table", tableName)
		l.Infof("Launching snapshot of table '%s'", tableName)

		switch {
		case s.pdbName != "":
			var conn *sql.Conn
			if conn, err = s.dbPool.Conn(ctx); err != nil {
				return fmt.Errorf("acquiring snapshot connection: %w", err)
			}
			defer func() {
				if err := conn.Close(); err != nil {
					// snapshot has completed at this point so logging the error is sufficient.
					s.log.Errorf("Closing snapshot connection: %v", err)
				}
			}()

			if _, err = conn.ExecContext(ctx, "ALTER SESSION SET CONTAINER = "+s.pdbName); err != nil {
				return fmt.Errorf("switching session to PDB '%s' for snapshot: %w", s.pdbName, err)
			}
			defer func() {
				if _, err := conn.ExecContext(context.Background(), "ALTER SESSION SET CONTAINER = CDB$ROOT"); err != nil {
					// logging the error is sufficient here, connection will be closed in defer call above.
					s.log.Errorf("Switching session back to root container: %v", err)
				}
			}()
			// Use context.Background() to prevent database/sql from spawning an
			// awaitDone goroutine that races with our explicit Rollback below.
			// The go-ora v2 driver has an unsynchronized field in Session that
			// causes a data race between BreakConnection (from awaitDone) and
			// IsBreak (from our Rollback). Transaction lifetime is managed
			// manually via the defer and explicit Rollback at the end.
			if tx, err = conn.BeginTx(context.Background(), nil); err != nil {
				return fmt.Errorf("beginning snapshot transaction: %w", err)
			}
		default:
			// Non-CDB mode: use db.BeginTx directly — no *Conn needed.
			// See CDB path comment above for why context.Background() is used.
			if tx, err = s.dbPool.BeginTx(context.Background(), nil); err != nil {
				return fmt.Errorf("beginning snapshot transaction: %w", err)
			}
		}

		// In Oracle, READ ONLY transactions automatically provide serializable isolation
		if _, err = tx.ExecContext(ctx, "SET TRANSACTION READ ONLY"); err != nil {
			_ = tx.Rollback()
			return fmt.Errorf("setting transaction read-only: %w", err)
		}
		defer func() {
			if err != nil {
				if rbErr := tx.Rollback(); rbErr != nil && !errors.Is(rbErr, sql.ErrTxDone) {
					l.Errorf("Failed to rollback snapshot transaction: %v", rbErr)
				}
			}
		}()

		var tablePks []string
		if tablePks, err = getTablePrimaryKeys(ctx, tx, table); err != nil {
			return err
		}

		l.Debugf("Found primary keys for table '%s': %v", table, tablePks)
		lastSeenPksValues := map[string]any{}
		for _, pk := range tablePks {
			lastSeenPksValues[pk] = nil
		}

		var (
			numRowsProcessed int
			batchCount       int
		)
		for {
			var pksForQuery map[string]any
			if numRowsProcessed > 0 {
				pksForQuery = lastSeenPksValues
			}
			batchCount, err = s.processBatch(ctx, tx, table, tablePks, pksForQuery, lastSeenPksValues, maxBatchSize, tableName)
			if err != nil {
				return fmt.Errorf("processing snapshot batch: %w", err)
			}

			numRowsProcessed += batchCount
			if batchCount < maxBatchSize {
				break
			}
		}

		if err := tx.Rollback(); err != nil {
			l.Errorf("Failed rollback snapshot transaction: %v", err)
		}
		s.snapshotStatusMetric.Set(1, tableName)
		l.Infof("Table snapshot completed, %d rows processed", numRowsProcessed)

		return nil
	}
}

// processBatch queries and processes a single page of rows from a snapshot table.
// pksForQuery is passed to querySnapshotTable for cursor-based pagination (nil on first batch).
// lastSeenPksValues is mutated in place with the PK values from the last row of the batch,
// so the caller can pass it as pksForQuery on the next iteration.
func (s *Snapshot) processBatch(ctx context.Context, tx *sql.Tx, table UserTable, tablePks []string, pksForQuery map[string]any, lastSeenPksValues map[string]any, maxBatchSize int, tableName string) (batchCount int, err error) {
	batchRows, err := querySnapshotTable(ctx, tx, table, tablePks, pksForQuery, maxBatchSize)
	if err != nil {
		return 0, fmt.Errorf("execute snapshot table query: %w", err)
	}
	defer func() {
		if closeErr := batchRows.Close(); closeErr != nil && err == nil {
			err = fmt.Errorf("closing snapshot rows: %w", closeErr)
		}
	}()

	types, err := batchRows.ColumnTypes()
	if err != nil {
		return 0, fmt.Errorf("fetch column types: %w", err)
	}

	values, mappers := prepSnapshotScannerAndMappers(types)

	columns, err := batchRows.Columns()
	if err != nil {
		return 0, fmt.Errorf("fetch columns: %w", err)
	}

	colMeta := buildColumnMeta(types)

	for batchRows.Next() {
		batchCount++

		if err := batchRows.Scan(values...); err != nil {
			return 0, err
		}

		var (
			v      any
			mapErr error
		)
		row := map[string]any{}
		for idx, value := range values {
			if v, mapErr = mappers[idx](value); mapErr != nil {
				return 0, mapErr
			}
			if !s.lobEnabled && isLOBType(types[idx].DatabaseTypeName()) {
				v = nil
			}
			row[columns[idx]] = v
			if _, ok := lastSeenPksValues[columns[idx]]; ok {
				lastSeenPksValues[columns[idx]] = value
			}
		}

		m := MessageEvent{
			Table:      table.Name,
			Schema:     table.Schema,
			Data:       row,
			Operation:  MessageOperationRead,
			SCN:        0,
			ColumnMeta: colMeta,
		}
		if err = s.publisher.Publish(ctx, &m); err != nil {
			return 0, fmt.Errorf("handling snapshot table row: %w", err)
		}
	}

	if err = batchRows.Err(); err != nil {
		return 0, fmt.Errorf("iterating snapshot table row: %w", err)
	}
	s.snapshotRowsTotalMetric.Incr(int64(batchCount), tableName)
	return batchCount, nil
}

func getTablePrimaryKeys(ctx context.Context, tx *sql.Tx, table UserTable) ([]string, error) {
	// ALL_CONSTRAINTS/ALL_CONS_COLUMNS work in any container context after ALTER SESSION SET CONTAINER.
	pkSQL := `
	SELECT acc.column_name
	FROM all_constraints ac
	JOIN all_cons_columns acc
		ON ac.constraint_name = acc.constraint_name
		AND ac.owner = acc.owner
	WHERE ac.constraint_type = 'P'
		AND UPPER(ac.table_name) = UPPER(:1)
		AND UPPER(ac.owner) = UPPER(:2)
	ORDER BY acc.position`
	pkArgs := []any{table.Name, table.Schema}

	rows, err := tx.QueryContext(ctx, pkSQL, pkArgs...)
	if err != nil {
		return nil, fmt.Errorf("get primary key: %w", err)
	}
	defer rows.Close()

	var pks []string
	for rows.Next() {
		var pk string
		if err := rows.Scan(&pk); err != nil {
			return nil, err
		}
		pks = append(pks, pk)
	}
	if err := rows.Err(); err != nil {
		return nil, fmt.Errorf("discovering primary keys for table '%s': %w", table.FullName(), err)
	}
	if len(pks) == 0 {
		return nil, fmt.Errorf("can't find a primary key for table '%s', does it exist and have one set?", table.FullName())
	}

	return pks, nil
}

func querySnapshotTable(ctx context.Context, tx *sql.Tx, table UserTable, pk []string, lastSeenPkVal map[string]any, limit int) (*sql.Rows, error) {
	// Oracle uses FETCH FIRST instead of TOP, and it comes at the end
	snapshotQueryParts := []string{
		fmt.Sprintf(`SELECT * FROM "%s"."%s"`, table.Schema, table.Name),
	}

	if lastSeenPkVal == nil {
		snapshotQueryParts = append(snapshotQueryParts, buildOrderByClause(pk))
		snapshotQueryParts = append(snapshotQueryParts, fmt.Sprintf("FETCH FIRST %d ROWS ONLY", limit))

		q := strings.Join(snapshotQueryParts, " ")
		return tx.QueryContext(ctx, q)
	}

	// Build lexicographic comparison for composite keys
	// For pk [col1, col2, col3], generates:
	// WHERE (col1 > ?) OR (col1 = ? AND col2 > ?) OR (col1 = ? AND col2 = ? AND col3 > ?)
	// Oracle uses positional parameters (:1, :2, etc.) or named parameters
	var (
		lastSeenPkVals []any
		paramIdx       int
		where          strings.Builder
	)

	where.WriteString("WHERE ")
	for i := range pk {
		if i > 0 {
			where.WriteString(" OR ")
		}
		where.WriteString("(")
		// Add equality conditions for all previous columns
		for j := range i {
			if j > 0 {
				where.WriteString(" AND ")
			}
			paramIdx++
			fmt.Fprintf(&where, `"%s" = :%d`, pk[j], paramIdx)
			lastSeenPkVals = append(lastSeenPkVals, lastSeenPkVal[pk[j]])
		}
		// Add greater-than condition for current column
		if i > 0 {
			where.WriteString(" AND ")
		}
		paramIdx++
		fmt.Fprintf(&where, `"%s" > :%d`, pk[i], paramIdx)
		lastSeenPkVals = append(lastSeenPkVals, lastSeenPkVal[pk[i]])
		where.WriteString(")")
	}

	snapshotQueryParts = append(snapshotQueryParts, where.String())
	snapshotQueryParts = append(snapshotQueryParts, buildOrderByClause(pk))
	snapshotQueryParts = append(snapshotQueryParts, fmt.Sprintf("FETCH FIRST %d ROWS ONLY", limit))
	q := strings.Join(snapshotQueryParts, " ")
	return tx.QueryContext(ctx, q, lastSeenPkVals...)
}

// Close safely closes all open connections opened for the snapshotting process.
// It should be called after a non-recoverale error or once the snapshot process has completed.
func (s *Snapshot) Close() error {
	if s.dbPool != nil {
		if err := s.dbPool.Close(); err != nil {
			return fmt.Errorf("closing database connection: %w", err)
		}
	}
	return nil
}

func prepSnapshotScannerAndMappers(cols []*sql.ColumnType) (values []any, mappers []func(any) (any, error)) {
	stringMapping := func(mapper func(s string) (any, error)) func(any) (any, error) {
		return func(v any) (any, error) {
			s, ok := v.(*sql.NullString)
			if !ok {
				return nil, fmt.Errorf("expected %T got %T", "", v)
			}
			if !s.Valid {
				return nil, nil
			}
			return mapper(s.String)
		}
	}
	for _, col := range cols {
		var val any
		var mapper func(any) (any, error)

		// Oracle database type names
		switch col.DatabaseTypeName() {
		case "RAW", "LONG RAW", "BLOB", "LongRaw":
			val = new(sql.Null[[]byte])
			mapper = snapshotValueMapper[[]byte]
		case "DATE", "TIMESTAMP", "TIMESTAMP WITH TIME ZONE", "TIMESTAMP WITH LOCAL TIME ZONE",
			"TimeStampTZ", "TimeStampDTY", "TimeStampTZ_DTY", "TimeStampLTZ_DTY", "TimeStampeLTZ", "TIMESTAMPTZ":
			val = new(sql.NullTime)
			mapper = func(v any) (any, error) {
				s, ok := v.(*sql.NullTime)
				if !ok {
					return nil, fmt.Errorf("expected %T got %T", time.Time{}, v)
				}
				if !s.Valid {
					return nil, nil
				}
				return s.Time, nil
			}
		case "NUMBER", "INTEGER", "INT", "SMALLINT", "FLOAT":
			// Oracle NUMBER type can represent both integers and decimals.
			// For integer-width columns (scale=0, precision<=18), scan as int64
			// to match the streaming path's ParseInt behavior.
			// For all others, scan as json.Number to preserve arbitrary precision.
			precision, scale, ok := col.DecimalSize()
			if ok && scale == 0 && precision > 0 && precision <= MaxInt64DecimalPrecision {
				val = new(sql.Null[int64])
				mapper = snapshotValueMapper[int64]
			} else {
				val = new(sql.NullString)
				mapper = stringMapping(func(s string) (any, error) {
					return json.Number(s), nil
				})
			}
		case "BINARY_FLOAT", "IBFloat", "BFloat", "BINARY_DOUBLE", "IBDouble", "BDouble":
			val = new(sql.Null[float64])
			mapper = snapshotValueMapper[float64]
		case "CLOB", "NCLOB", "LONG", "LongVarChar":
			// Character large objects - handle as string
			val = new(sql.NullString)
			mapper = stringMapping(func(s string) (any, error) {
				return s, nil
			})
		case "JSON":
			// Oracle 21c+ native JSON type
			val = new(sql.NullString)
			mapper = stringMapping(func(s string) (v any, err error) {
				err = json.Unmarshal([]byte(s), &v)
				return
			})
		default:
			// Default to string for VARCHAR2, CHAR, NVARCHAR2, NCHAR, etc.
			val = new(sql.Null[string])
			mapper = snapshotValueMapper[string]
		}
		values = append(values, val)
		mappers = append(mappers, mapper)
	}
	return
}

func buildOrderByClause(pk []string) string {
	quoted := make([]string, len(pk))
	for i, col := range pk {
		quoted[i] = `"` + col + `"`
	}
	return "ORDER BY " + strings.Join(quoted, ", ")
}

// buildColumnMeta extracts lightweight type metadata from sql.ColumnType values
// for carrying through MessageEvent to the schema cache.
func buildColumnMeta(types []*sql.ColumnType) []ColumnMeta {
	meta := make([]ColumnMeta, len(types))
	for i, ct := range types {
		meta[i] = ColumnMeta{
			Name:     ct.Name(),
			TypeName: ct.DatabaseTypeName(),
		}
		if precision, scale, ok := ct.DecimalSize(); ok {
			meta[i].Precision = precision
			meta[i].Scale = scale
			meta[i].HasDecimalSize = true
		}
	}
	return meta
}

func isLOBType(dbType string) bool {
	switch dbType {
	case "CLOB", "NCLOB", "BLOB", "LONG", "LONG RAW",
		"LongVarChar", "LongRaw": // go-ora driver-level names for CLOB/NCLOB/LONG and BLOB/LONG RAW
		return true
	}
	return false
}

func snapshotValueMapper[T any](v any) (any, error) {
	s, ok := v.(*sql.Null[T])
	if !ok {
		var e T
		return nil, fmt.Errorf("expected %T got %T", e, v)
	}
	if !s.Valid {
		return nil, nil
	}
	return s.V, nil
}