Merge pull request #841 from datazip-inc/staging

chore: staging -> master v0.4.0

Author: vishalm0509

destination/iceberg/iceberg.go

@@ -102,7 +102,7 @@ func (i *Iceberg) Setup(ctx context.Context, stream types.StreamInterface, globa
 		logger.Infof("Creating destination table [%s] in Iceberg database [%s] for stream [%s]", i.stream.GetDestinationTable(), i.stream.GetDestinationDatabase(&i.config.IcebergDatabase), i.stream.Name())
 
 		var requestPayload proto.IcebergPayload
-		iceSchema := stream.Schema().ToIceberg(!stream.NormalizationEnabled())
+		iceSchema := stream.Schema().ToIceberg(!stream.NormalizationEnabled(), i.stream)
 		requestPayload = proto.IcebergPayload{
 			Type: proto.IcebergPayload_GET_OR_CREATE_TABLE,
 			Metadata: &proto.IcebergPayload_Metadata{

destination/parquet/parquet.go

@@ -104,9 +104,9 @@ func (p *Parquet) createNewPartitionFile(basePath string) error {
 
 	writer := func() any {
 		if p.stream.NormalizationEnabled() {
-			return pqgo.NewGenericWriter[any](pqFile, p.schema.ToTypeSchema().ToParquet(false), pqgo.Compression(&pqgo.Snappy))
+			return pqgo.NewGenericWriter[any](pqFile, p.schema.ToTypeSchema().ToParquet(false, p.stream), pqgo.Compression(&pqgo.Snappy))
 		}
-		return pqgo.NewGenericWriter[any](pqFile, p.stream.Schema().ToParquet(true), pqgo.Compression(&pqgo.Snappy))
+		return pqgo.NewGenericWriter[any](pqFile, p.stream.Schema().ToParquet(true, p.stream), pqgo.Compression(&pqgo.Snappy))
 	}()
 
 	p.partitionedFiles[basePath] = append(p.partitionedFiles[basePath], &FileMetadata{

drivers/abstract/backfill.go

@@ -40,6 +40,8 @@ func (a *AbstractDriver) Backfill(mainCtx context.Context, backfilledStreams cha
 
 	logger.Infof("Starting backfill for stream[%s] with %d chunks", stream.GetStream().Name, len(chunks))
 
+	filterDataBySelectedColumnsFn := stream.RetainSelectedColumns()
+
 	chunkProcessor := func(gCtx context.Context, _ int, chunk types.Chunk) (err error) {
 		// create backfill context, so that main context not affected if backfill retries
 		backfillCtx, backfillCtxCancel := context.WithCancel(gCtx)
@@ -66,6 +68,7 @@ func (a *AbstractDriver) Backfill(mainCtx context.Context, backfilledStreams cha
 				logger.Infof("finished chunk min[%v] and max[%v] of stream %s", chunk.Min, chunk.Max, stream.ID())
 				return nil
 			})()
+
 		return a.driver.ChunkIterator(backfillCtx, stream, chunk, func(ctx context.Context, data map[string]any) error {
 			olakeID := utils.GetKeysHash(data, stream.GetStream().SourceDefinedPrimaryKey.Array()...)
 			olakeColumns := map[string]any{
@@ -78,7 +81,10 @@ func (a *AbstractDriver) Backfill(mainCtx context.Context, backfilledStreams cha
 			if stream.GetSyncMode() == types.CDC {
 				olakeColumns[constants.CdcTimestamp] = time.Unix(0, 0)
 			}
-			return inserter.Push(ctx, types.CreateRawRecord(data, olakeColumns))
+
+			filteredData := filterDataBySelectedColumnsFn(data)
+
+			return inserter.Push(ctx, types.CreateRawRecord(filteredData, olakeColumns))
 		})
 	}
 	utils.ConcurrentInGroupWithRetry(a.GlobalConnGroup, chunks, a.driver.MaxRetries(), chunkProcessor)

drivers/abstract/cdc.go

@@ -99,6 +99,7 @@ func (a *AbstractDriver) RunChangeStream(mainCtx context.Context, pool *destinat
 //   - For Postgres: ignored (uses global replication slot)
 func (a *AbstractDriver) streamChanges(mainCtx context.Context, pool *destination.WriterPool, streamIndex int) (err error) {
 	writers := make(map[string]*destination.WriterThread)
+	filterDataBySelectedColumnsFns := make(map[string]func(map[string]interface{}) map[string]interface{})
 
 	// create cdc context, so that main context not affected if cdc retries
 	cdcCtx, cdcCtxCancel := context.WithCancel(mainCtx)
@@ -131,7 +132,15 @@ func (a *AbstractDriver) streamChanges(mainCtx context.Context, pool *destinatio
 			constants.OlakeTimestamp: time.Now().UTC(),
 		}
 		maps.Copy(olakeColumns, change.ExtraColumns)
-		return writer.Push(ctx, types.CreateRawRecord(change.Data, olakeColumns))
+
+		filterDataBySelectedColumnsFn, exists := filterDataBySelectedColumnsFns[change.Stream.ID()]
+		if !exists {
+			filterDataBySelectedColumnsFn = change.Stream.RetainSelectedColumns()
+			filterDataBySelectedColumnsFns[change.Stream.ID()] = filterDataBySelectedColumnsFn
+		}
+		filteredData := filterDataBySelectedColumnsFn(change.Data)
+
+		return writer.Push(ctx, types.CreateRawRecord(filteredData, olakeColumns))
 	})
 }
 

drivers/abstract/incremental.go

@@ -92,6 +92,8 @@ func (a *AbstractDriver) Incremental(mainCtx context.Context, pool *destination.
 					return nil
 				})()
 
+			filterDataBySelectedColumnsFn := stream.RetainSelectedColumns()
+
 			// No retry logic here - retry happens at Read level
 			return a.driver.StreamIncrementalChanges(incrementalCtx, stream, func(ctx context.Context, record map[string]any) error {
 				maxPrimaryCursorValue, maxSecondaryCursorValue = a.getMaxIncrementCursorFromData(primaryCursor, secondaryCursor, maxPrimaryCursorValue, maxSecondaryCursorValue, record)
@@ -100,7 +102,8 @@ func (a *AbstractDriver) Incremental(mainCtx context.Context, pool *destination.
 					constants.OpType:         "u",
 					constants.OlakeTimestamp: time.Now().UTC(),
 				}
-				return inserter.Push(ctx, types.CreateRawRecord(record, olakeColumns))
+				filteredData := filterDataBySelectedColumnsFn(record)
+				return inserter.Push(ctx, types.CreateRawRecord(filteredData, olakeColumns))
 			})
 		})
 		return nil

drivers/db2/internal/testdata/test_streams.json

@@ -1 +1 @@
-{"selected_streams":{"DB2INST1":[{"partition_regex":"","stream_name":"DB2_TEST_TABLE_OLAKE","normalization":true}]},"streams":[{"stream":{"name":"DB2_TEST_TABLE_OLAKE","namespace":"DB2INST1","type_schema":{"properties":{"COL_BIGINT":{"type":["integer","null"],"destination_column_name":"col_bigint"},"COL_BLOB":{"type":["string","null"],"destination_column_name":"col_blob"},"COL_BOOL":{"type":["boolean","null"],"destination_column_name":"col_bool"},"COL_CHAR":{"type":["string","null"],"destination_column_name":"col_char"},"COL_CHARACTER":{"type":["string","null"],"destination_column_name":"col_character"},"COL_CLOB":{"type":["string","null"],"destination_column_name":"col_clob"},"COL_CURSOR":{"type":["integer","null"],"destination_column_name":"col_cursor"},"COL_DATE":{"type":["timestamp","null"],"destination_column_name":"col_date"},"COL_DECIMAL":{"type":["number","null"],"destination_column_name":"col_decimal"},"COL_DOUBLE":{"type":["number","null"],"destination_column_name":"col_double"},"COL_GRAPHIC":{"type":["string","null"],"destination_column_name":"col_graphic"},"COL_INT":{"type":["null","integer_small"],"destination_column_name":"col_int"},"COL_REAL":{"type":["number_small","null"],"destination_column_name":"col_real"},"COL_SMALLINT":{"type":["integer_small","null"],"destination_column_name":"col_smallint"},"COL_TIME":{"type":["string","null"],"destination_column_name":"col_time"},"COL_TIMESTAMP":{"type":["timestamp","null"],"destination_column_name":"col_timestamp"},"COL_VARCHAR":{"type":["string","null"],"destination_column_name":"col_varchar"},"COL_VARGRAPHIC":{"type":["string","null"],"destination_column_name":"col_vargraphic"},"ID":{"type":["integer"],"destination_column_name":"id"},"_olake_id":{"type":["string","null"],"destination_column_name":"_olake_id","olake_column":true},"_olake_timestamp":{"type":["timestamp_micro","null"],"destination_column_name":"_olake_timestamp","olake_column":true},"_op_type":{"type":["null","string"],"destination_column_name":"_op_type","olake_column":true}}},"supported_sync_modes":["full_refresh","incremental"],"source_defined_primary_key":["ID"],"available_cursor_fields":["COL_CHARACTER","COL_VARCHAR","COL_TIME","COL_GRAPHIC","ID","COL_CHAR","COL_DECIMAL","COL_BLOB","COL_VARGRAPHIC","COL_CURSOR","COL_BIGINT","COL_DOUBLE","COL_SMALLINT","COL_BOOL","COL_TIMESTAMP","COL_DATE","COL_REAL","COL_INT","COL_CLOB"],"sync_mode":"incremental","destination_database":"db2_testdb:db2inst1","destination_table":"db2_test_table_olake","default_stream_properties":{"normalization":true,"append_mode":false}}}]}
+{"selected_streams":{"DB2INST1":[{"partition_regex":"","stream_name":"DB2_TEST_TABLE_OLAKE","normalization":true,"selected_columns":{"columns":["_op_type","COL_GRAPHIC","COL_BIGINT","COL_BLOB","COL_CURSOR","COL_TIMESTAMP","COL_VARGRAPHIC","COL_BOOL","COL_CHAR","COL_CLOB","COL_DOUBLE","ID","_olake_timestamp","COL_VARCHAR","COL_SMALLINT","COL_INT","COL_REAL","COL_DATE","_olake_id","COL_DECIMAL","COL_TIME","COL_CHARACTER"],"sync_new_columns":true}}]},"streams":[{"stream":{"name":"DB2_TEST_TABLE_OLAKE","namespace":"DB2INST1","type_schema":{"properties":{"COL_BIGINT":{"type":["integer","null"],"destination_column_name":"col_bigint"},"COL_BLOB":{"type":["string","null"],"destination_column_name":"col_blob"},"COL_BOOL":{"type":["boolean","null"],"destination_column_name":"col_bool"},"COL_CHAR":{"type":["string","null"],"destination_column_name":"col_char"},"COL_CHARACTER":{"type":["string","null"],"destination_column_name":"col_character"},"COL_CLOB":{"type":["string","null"],"destination_column_name":"col_clob"},"COL_CURSOR":{"type":["integer","null"],"destination_column_name":"col_cursor"},"COL_DATE":{"type":["null","timestamp"],"destination_column_name":"col_date"},"COL_DECIMAL":{"type":["number","null"],"destination_column_name":"col_decimal"},"COL_DOUBLE":{"type":["number","null"],"destination_column_name":"col_double"},"COL_GRAPHIC":{"type":["string","null"],"destination_column_name":"col_graphic"},"COL_INT":{"type":["integer_small","null"],"destination_column_name":"col_int"},"COL_REAL":{"type":["number_small","null"],"destination_column_name":"col_real"},"COL_SMALLINT":{"type":["integer_small","null"],"destination_column_name":"col_smallint"},"COL_TIME":{"type":["string","null"],"destination_column_name":"col_time"},"COL_TIMESTAMP":{"type":["timestamp","null"],"destination_column_name":"col_timestamp"},"COL_VARCHAR":{"type":["string","null"],"destination_column_name":"col_varchar"},"COL_VARGRAPHIC":{"type":["string","null"],"destination_column_name":"col_vargraphic"},"ID":{"type":["integer"],"destination_column_name":"id"},"_olake_id":{"type":["string","null"],"destination_column_name":"_olake_id","olake_column":true},"_olake_timestamp":{"type":["null","timestamp_micro"],"destination_column_name":"_olake_timestamp","olake_column":true},"_op_type":{"type":["string","null"],"destination_column_name":"_op_type","olake_column":true}}},"supported_sync_modes":["full_refresh","incremental"],"source_defined_primary_key":["ID"],"available_cursor_fields":["ID","COL_DECIMAL","COL_SMALLINT","COL_BOOL","COL_BIGINT","COL_CHAR","COL_VARCHAR","COL_DATE","COL_REAL","COL_INT","COL_BLOB","COL_CHARACTER","COL_CLOB","COL_TIMESTAMP","COL_TIME","COL_GRAPHIC","COL_VARGRAPHIC","COL_CURSOR","COL_DOUBLE"],"sync_mode":"incremental","destination_database":"db2_testdb:db2inst1","destination_table":"db2_test_table_olake","default_stream_properties":{"normalization":true,"append_mode":false}}}]}

drivers/mongodb/internal/testdata/test_streams.json

@@ -1 +1 @@
-{"selected_streams":{"olake_mongodb_test":[{"partition_regex":"","stream_name":"mongodb_test_table_olake","normalization":false},{"partition_regex":"","stream_name":"test_collection","normalization":false}]},"streams":[{"stream":{"name":"mongodb_test_table_olake","namespace":"olake_mongodb_test","type_schema":{"properties":{"_cdc_resume_token":{"type":["string","null"],"destination_column_name":"_cdc_resume_token","olake_column":true},"_cdc_timestamp":{"type":["timestamp_micro","null"],"destination_column_name":"_cdc_timestamp","olake_column":true},"_id":{"type":["string"],"destination_column_name":"_id"},"_olake_id":{"type":["string","null"],"destination_column_name":"_olake_id","olake_column":true},"_olake_timestamp":{"type":["timestamp_micro","null"],"destination_column_name":"_olake_timestamp","olake_column":true},"_op_type":{"type":["string","null"],"destination_column_name":"_op_type","olake_column":true},"created_timestamp":{"type":["integer_small"],"destination_column_name":"created_timestamp"},"id":{"type":["integer_small"],"destination_column_name":"id"},"id_bigint":{"type":["integer"],"destination_column_name":"id_bigint"},"id_bool":{"type":["boolean"],"destination_column_name":"id_bool"},"id_cursor":{"type":["integer_small"],"destination_column_name":"id_cursor"},"id_double":{"type":["number"],"destination_column_name":"id_double"},"id_int":{"type":["integer_small"],"destination_column_name":"id_int"},"id_maxkey":{"type":["unknown"],"destination_column_name":"id_maxkey"},"id_minkey":{"type":["unknown"],"destination_column_name":"id_minkey"},"id_nested":{"type":["object"],"destination_column_name":"id_nested"},"id_nil":{"type":["null"],"destination_column_name":"id_nil"},"id_regex":{"type":["unknown"],"destination_column_name":"id_regex"},"id_timestamp":{"type":["timestamp"],"destination_column_name":"id_timestamp"},"name_varchar":{"type":["string"],"destination_column_name":"name_varchar"}}},"supported_sync_modes":["strict_cdc","full_refresh","incremental","cdc"],"source_defined_primary_key":["_id"],"available_cursor_fields":["id_nested","id_int","id_bigint","id_maxkey","created_timestamp","id_double","id_timestamp","id_nil","id_regex","name_varchar","id_bool","_id","id","id_cursor","id_minkey"],"sync_mode":"cdc","destination_database":"mongodb:olake_mongodb_test","destination_table":"mongodb_test_table_olake","default_stream_properties":{"normalization":false,"append_mode":false}}},{"stream":{"name":"test_collection","namespace":"olake_mongodb_test","type_schema":{"properties":{"_cdc_resume_token":{"type":["string","null"],"destination_column_name":"_cdc_resume_token","olake_column":true},"_cdc_timestamp":{"type":["timestamp_micro","null"],"destination_column_name":"_cdc_timestamp","olake_column":true},"_olake_id":{"type":["string","null"],"destination_column_name":"_olake_id","olake_column":true},"_olake_timestamp":{"type":["timestamp_micro","null"],"destination_column_name":"_olake_timestamp","olake_column":true},"_op_type":{"type":["string","null"],"destination_column_name":"_op_type","olake_column":true}}},"supported_sync_modes":["incremental","cdc","strict_cdc","full_refresh"],"source_defined_primary_key":["_id"],"available_cursor_fields":[],"sync_mode":"cdc","destination_database":"mongodb:olake_mongodb_test","destination_table":"test_collection","default_stream_properties":{"normalization":false,"append_mode":false}}}]}
+{"selected_streams":{"olake_mongodb_test":[{"partition_regex":"","stream_name":"test_collection","normalization":false,"selected_columns":{"columns":["_olake_id","_cdc_resume_token","_olake_timestamp","_cdc_timestamp","_op_type"],"sync_new_columns":true}},{"partition_regex":"","stream_name":"mongodb_test_table_olake","normalization":false,"selected_columns":{"columns":["id_timestamp","created_timestamp","id","id_int","id_nil","id_minkey","_olake_id","_olake_timestamp","id_maxkey","id_double","id_regex","name_varchar","_id","id_bigint","_op_type","id_nested","id_cursor","id_bool","_cdc_resume_token","_cdc_timestamp"],"sync_new_columns":true}}]},"streams":[{"stream":{"name":"test_collection","namespace":"olake_mongodb_test","type_schema":{"properties":{"_cdc_resume_token":{"type":["string","null"],"destination_column_name":"_cdc_resume_token","olake_column":true},"_cdc_timestamp":{"type":["timestamp_micro","null"],"destination_column_name":"_cdc_timestamp","olake_column":true},"_olake_id":{"type":["string","null"],"destination_column_name":"_olake_id","olake_column":true},"_olake_timestamp":{"type":["timestamp_micro","null"],"destination_column_name":"_olake_timestamp","olake_column":true},"_op_type":{"type":["string","null"],"destination_column_name":"_op_type","olake_column":true}}},"supported_sync_modes":["incremental","cdc","strict_cdc","full_refresh"],"source_defined_primary_key":["_id"],"available_cursor_fields":[],"sync_mode":"cdc","destination_database":"mongodb:olake_mongodb_test","destination_table":"test_collection","default_stream_properties":{"normalization":false,"append_mode":false}}},{"stream":{"name":"mongodb_test_table_olake","namespace":"olake_mongodb_test","type_schema":{"properties":{"_cdc_resume_token":{"type":["string","null"],"destination_column_name":"_cdc_resume_token","olake_column":true},"_cdc_timestamp":{"type":["null","timestamp_micro"],"destination_column_name":"_cdc_timestamp","olake_column":true},"_id":{"type":["string"],"destination_column_name":"_id"},"_olake_id":{"type":["null","string"],"destination_column_name":"_olake_id","olake_column":true},"_olake_timestamp":{"type":["null","timestamp_micro"],"destination_column_name":"_olake_timestamp","olake_column":true},"_op_type":{"type":["string","null"],"destination_column_name":"_op_type","olake_column":true},"created_timestamp":{"type":["integer_small"],"destination_column_name":"created_timestamp"},"id":{"type":["integer_small"],"destination_column_name":"id"},"id_bigint":{"type":["integer"],"destination_column_name":"id_bigint"},"id_bool":{"type":["boolean"],"destination_column_name":"id_bool"},"id_cursor":{"type":["integer_small"],"destination_column_name":"id_cursor"},"id_double":{"type":["number"],"destination_column_name":"id_double"},"id_int":{"type":["integer_small"],"destination_column_name":"id_int"},"id_maxkey":{"type":["unknown"],"destination_column_name":"id_maxkey"},"id_minkey":{"type":["unknown"],"destination_column_name":"id_minkey"},"id_nested":{"type":["object"],"destination_column_name":"id_nested"},"id_nil":{"type":["null"],"destination_column_name":"id_nil"},"id_regex":{"type":["unknown"],"destination_column_name":"id_regex"},"id_timestamp":{"type":["timestamp"],"destination_column_name":"id_timestamp"},"name_varchar":{"type":["string"],"destination_column_name":"name_varchar"}}},"supported_sync_modes":["full_refresh","incremental","cdc","strict_cdc"],"source_defined_primary_key":["_id"],"available_cursor_fields":["id_nil","id_minkey","_id","id_cursor","id_timestamp","created_timestamp","id_double","id_nested","id","id_int","id_regex","name_varchar","id_bigint","id_bool","id_maxkey"],"sync_mode":"cdc","destination_database":"mongodb:olake_mongodb_test","destination_table":"mongodb_test_table_olake","default_stream_properties":{"normalization":false,"append_mode":false}}}]}