fix: kill process manually after timeout (#493)

Co-authored-by: Ankit Sharma <[email protected]>

Author: vishalm0509

.github/workflows/performance-test.yml

@@ -95,8 +95,20 @@ jobs:
     - name: Cleanup
       if: always()
       run: |
-        aws glue delete-database --name performance_${{ matrix.driver }} || { echo "failed to delete glue database: performance_${{ matrix.driver }}"; aws glue get-database --name performance_${{ matrix.driver }} || true; }
-        aws s3 rm s3://dz-stag-github-actions/performance_${{ matrix.driver }} --recursive || { echo "failed to delete s3 bucket: performance_${{ matrix.driver }}"; aws s3 ls s3://dz-stag-github-actions/performance_${{ matrix.driver }} || true; }
+        # Delete all Glue databases matching prefix
+        for db in $(aws glue get-databases \
+          --query "DatabaseList[?starts_with(Name, 'performance_${{ matrix.driver }}')].Name" \
+          --output text); do
+          echo "Deleting Glue database: $db"
+          aws glue delete-database --name "$db" \
+          || echo "Failed to delete Glue database: $db"
+        done
+
+        # Delete corresponding S3 path
+        aws s3 rm s3://dz-stag-github-actions/performance_${{ matrix.driver }}/ --recursive \
+        || { echo "Failed to delete S3 bucket: performance_${{ matrix.driver }}"; \
+        aws s3 ls s3://dz-stag-github-actions/performance_${{ matrix.driver }} || true; }
+
         echo "Catalog cleanup completed"
 
         if [[ "${{ matrix.driver }}" != "oracle" ]]; then

destination/iceberg/iceberg.go

@@ -17,6 +17,7 @@ import (
 	"github.com/datazip-inc/olake/utils"
 	"github.com/datazip-inc/olake/utils/logger"
 	"github.com/datazip-inc/olake/utils/typeutils"
+	"github.com/spf13/viper"
 )
 
 type Iceberg struct {
@@ -261,8 +262,13 @@ func (i *Iceberg) Check(ctx context.Context) error {
 	i.options = &destination.Options{
 		ThreadID: "test_iceberg_destination",
 	}
+
+	destinationDB := "test_olake"
+	if prefix := viper.GetString(constants.DestinationDatabasePrefix); prefix != "" {
+		destinationDB = fmt.Sprintf("%s_%s", utils.Reformat(prefix), destinationDB)
+	}
 	// Create a temporary setup for checking
-	server, err := newIcebergClient(i.config, []PartitionInfo{}, i.options.ThreadID, true, false, "test_olake")
+	server, err := newIcebergClient(i.config, []PartitionInfo{}, i.options.ThreadID, true, false, destinationDB)
 	if err != nil {
 		return fmt.Errorf("failed to setup iceberg server: %s", err)
 	}
@@ -281,7 +287,7 @@ func (i *Iceberg) Check(ctx context.Context) error {
 		Type: proto.IcebergPayload_GET_OR_CREATE_TABLE,
 		Metadata: &proto.IcebergPayload_Metadata{
 			ThreadId:      server.serverID,
-			DestTableName: "test_olake",
+			DestTableName: destinationDB,
 			Schema:        icebergRawSchema(),
 		},
 	}
@@ -296,7 +302,7 @@ func (i *Iceberg) Check(ctx context.Context) error {
 	// try writing record in dest table
 	currentTime := time.Now().UTC()
 	protoSchema := icebergRawSchema()
-	record := types.CreateRawRecord("olake_test", map[string]any{"name": "olake"}, "r", &currentTime)
+	record := types.CreateRawRecord(destinationDB, map[string]any{"name": "olake"}, "r", &currentTime)
 	protoColumns, err := rawDataColumnBuffer(record, protoSchema)
 	if err != nil {
 		return fmt.Errorf("failed to create raw data column buffer: %s", err)
@@ -305,7 +311,7 @@ func (i *Iceberg) Check(ctx context.Context) error {
 		Type: proto.IcebergPayload_RECORDS,
 		Metadata: &proto.IcebergPayload_Metadata{
 			ThreadId:      server.serverID,
-			DestTableName: "test_olake",
+			DestTableName: destinationDB,
 			Schema:        protoSchema,
 		},
 		Records: []*proto.IcebergPayload_IceRecord{{

protocol/sync.go

@@ -6,12 +6,14 @@ import (
 	"sync"
 	"time"
 
+	"github.com/datazip-inc/olake/constants"
 	"github.com/datazip-inc/olake/destination"
 	"github.com/datazip-inc/olake/types"
 	"github.com/datazip-inc/olake/utils"
 	"github.com/datazip-inc/olake/utils/logger"
 	"github.com/datazip-inc/olake/utils/telemetry"
 	"github.com/spf13/cobra"
+	"github.com/spf13/viper"
 )
 
 // syncCmd represents the read command
@@ -38,6 +40,11 @@ var syncCmd = &cobra.Command{
 			return err
 		}
 
+		// to set prefix for "test_olake" db created by OLake
+		if destinationDatabasePrefix != "" {
+			viper.Set(constants.DestinationDatabasePrefix, destinationDatabasePrefix)
+		}
+
 		catalog = &types.Catalog{}
 		if err := utils.UnmarshalFile(streamsPath, catalog, false); err != nil {
 			return err

utils/testutils/test_utils.go

@@ -102,16 +102,25 @@ func GetTestConfig(driver string) *TestConfig {
 	}
 }
 
-func syncCommand(config TestConfig, useState bool) string {
+func syncCommand(config TestConfig, useState bool, flags ...string) string {
 	baseCmd := fmt.Sprintf("/test-olake/build.sh driver-%s sync --config %s --catalog %s --destination %s", config.Driver, config.SourcePath, config.CatalogPath, config.DestinationPath)
 	if useState {
 		baseCmd = fmt.Sprintf("%s --state %s", baseCmd, config.StatePath)
 	}
+
+	if len(flags) > 0 {
+		baseCmd = fmt.Sprintf("%s %s", baseCmd, strings.Join(flags, " "))
+	}
 	return baseCmd
 }
 
-func discoverCommand(config TestConfig) string {
-	return fmt.Sprintf("/test-olake/build.sh driver-%s discover --config %s", config.Driver, config.SourcePath)
+// pass flags as `--flag1, flag1 value, --flag2, flag2 value...`
+func discoverCommand(config TestConfig, flags ...string) string {
+	baseCmd := fmt.Sprintf("/test-olake/build.sh driver-%s discover --config %s", config.Driver, config.SourcePath)
+	if len(flags) > 0 {
+		baseCmd = fmt.Sprintf("%s %s", baseCmd, strings.Join(flags, " "))
+	}
+	return baseCmd
 }
 
 // TODO: check if we can remove namespace from being passed as a parameter and use a common namespace for all drivers
@@ -307,8 +316,9 @@ func (cfg *IntegrationTest) TestIntegration(t *testing.T) {
 								},
 							}
 
+							destDBPrefix := fmt.Sprintf("integration_%s", cfg.TestConfig.Driver)
 							runSync := func(c testcontainers.Container, useState bool, operation, opSymbol string, schema map[string]interface{}) error {
-								cmd := syncCommand(*cfg.TestConfig, useState)
+								cmd := syncCommand(*cfg.TestConfig, useState, "--destination-database-prefix", destDBPrefix)
 								if useState && operation != "" {
 									cfg.ExecuteQuery(ctx, t, []string{currentTestTable}, operation, false)
 								}
@@ -460,6 +470,8 @@ func (cfg *PerformanceTest) TestPerformance(t *testing.T) {
 		code, output, err := utils.ExecCommand(timedCtx, c, cmd)
 		// check if sync was canceled due to timeout (expected)
 		if timedCtx.Err() == context.DeadlineExceeded {
+			killCmd := "pkill -9 -f 'olake.*sync' || true"
+			_, _, _ = utils.ExecCommand(ctx, c, killCmd)
 			return output, nil
 		}
 		if err != nil || code != 0 {
@@ -491,11 +503,12 @@ func (cfg *PerformanceTest) TestPerformance(t *testing.T) {
 							if code, output, err := utils.ExecCommand(ctx, c, installCmd); err != nil || code != 0 {
 								return fmt.Errorf("failed to install dependencies:\n%s", string(output))
 							}
-
 							t.Logf("(backfill) running performance test for %s", cfg.TestConfig.Driver)
 
+							destDBPrefix := fmt.Sprintf("performance_%s", cfg.TestConfig.Driver)
+
 							t.Log("(backfill) discover started")
-							discoverCmd := discoverCommand(*cfg.TestConfig)
+							discoverCmd := discoverCommand(*cfg.TestConfig, "--destination-database-prefix", destDBPrefix)
 							if code, output, err := utils.ExecCommand(ctx, c, discoverCmd); err != nil || code != 0 {
 								return fmt.Errorf("failed to perform discover:\n%s", string(output))
 							}
@@ -508,7 +521,7 @@ func (cfg *PerformanceTest) TestPerformance(t *testing.T) {
 
 							t.Log("(backfill) sync started")
 							usePreChunkedState := cfg.TestConfig.Driver == string(constants.MySQL)
-							syncCmd := syncCommand(*cfg.TestConfig, usePreChunkedState)
+							syncCmd := syncCommand(*cfg.TestConfig, usePreChunkedState, "--destination-database-prefix", destDBPrefix)
 							if output, err := syncWithTimeout(ctx, c, syncCmd); err != nil {
 								return fmt.Errorf("failed to perform sync:\n%s", string(output))
 							}
@@ -529,7 +542,7 @@ func (cfg *PerformanceTest) TestPerformance(t *testing.T) {
 								t.Log("(cdc) setup cdc completed")
 
 								t.Log("(cdc) discover started")
-								discoverCmd := discoverCommand(*cfg.TestConfig)
+								discoverCmd := discoverCommand(*cfg.TestConfig, "--destination-database-prefix", destDBPrefix)
 								if code, output, err := utils.ExecCommand(ctx, c, discoverCmd); err != nil || code != 0 {
 									return fmt.Errorf("failed to perform discover:\n%s", string(output))
 								}
@@ -541,7 +554,7 @@ func (cfg *PerformanceTest) TestPerformance(t *testing.T) {
 								}
 
 								t.Log("(cdc) state creation started")
-								syncCmd := syncCommand(*cfg.TestConfig, false)
+								syncCmd := syncCommand(*cfg.TestConfig, false, "--destination-database-prefix", destDBPrefix)
 								if code, output, err := utils.ExecCommand(ctx, c, syncCmd); err != nil || code != 0 {
 									return fmt.Errorf("failed to perform initial sync:\n%s", string(output))
 								}
@@ -552,7 +565,7 @@ func (cfg *PerformanceTest) TestPerformance(t *testing.T) {
 								t.Log("(cdc) trigger cdc completed")
 
 								t.Log("(cdc) sync started")
-								syncCmd = syncCommand(*cfg.TestConfig, true)
+								syncCmd = syncCommand(*cfg.TestConfig, true, "--destination-database-prefix", destDBPrefix)
 								if output, err := syncWithTimeout(ctx, c, syncCmd); err != nil {
 									return fmt.Errorf("failed to perform CDC sync:\n%s", string(output))
 								}