fixup! chore: filter workspaces at router and batchrouter - like destinationIDs

Sidddddarth · Sidddddarth · commit 8957d7351320 · 2025-05-05T14:14:46.000+05:30
diff --git a/integration_test/multi_tenant_test/testdata/mtGatewayTest01.json b/integration_test/multi_tenant_test/testdata/mtGatewayTest01.json
@@ -17,6 +17,7 @@
           {
             "id": "xxxyyyzzP9kQfzOoKd1tuxchYAG",
             "name": "Dev WebHook Integration Test 1",
+            "workspaceId": "{{.workspaceId}}",
             "enabled": true,
             "isProcessorEnabled": true,
             "config": {
diff --git a/jobsdb/distinct_values_cache.go b/jobsdb/distinct_values_cache.go
@@ -26,29 +26,29 @@ type distinctValuesCache struct {
 // already cached, it returns the cached values. If not, it loads the values from the given load
 // function and caches them. The last dataset is never cached, so it is always loaded from the
 // load function. The load function is called with the missing datasets and the last dataset.
-func (dpc *distinctValuesCache) GetDistinctValues(key string, datasets []string, load func(datasets []string) (map[string][]string, error)) ([]string, error) {
+func (dvc *distinctValuesCache) GetDistinctValues(key string, datasets []string, load func(datasets []string) (map[string][]string, error)) ([]string, error) {
 	// First check if we are missing any datasets from the cache.
 	// If we are, we need to load them along with the last dataset
 	// The last dataset is never cached, so we need to load it every time
-	dpc.cacheMu.RLock()
-	missing := dpc.missing(key, datasets[:len(datasets)-1])
-	dpc.cacheMu.RUnlock()
+	dvc.cacheMu.RLock()
+	missing := dvc.missing(key, datasets[:len(datasets)-1])
+	dvc.cacheMu.RUnlock()
 
 	// If we are missing any datasets, we need to lock the key, so that
 	// we don't load the same datasets multiple times for the same key.
 	// This lock needs to be retained until the datasets are loaded into the cache.
 	if len(missing) > 0 {
-		dpc.klock.Lock(key)
+		dvc.klock.Lock(key)
 		// Check again if we are missing any datasets, to deal with race conditions
-		dpc.cacheMu.Lock()
-		missing = dpc.missing(key, datasets[:len(datasets)-1])
-		if _, ok := dpc.cache[key]; !ok {
-			dpc.cache[key] = make(map[string][]string)
+		dvc.cacheMu.Lock()
+		missing = dvc.missing(key, datasets[:len(datasets)-1])
+		if _, ok := dvc.cache[key]; !ok {
+			dvc.cache[key] = make(map[string][]string)
 		}
-		dpc.cacheMu.Unlock()
+		dvc.cacheMu.Unlock()
 		if len(missing) == 0 {
 			// If we are not missing any datasets, we need to unlock the key
-			dpc.klock.Unlock(key)
+			dvc.klock.Unlock(key)
 		}
 	}
 
@@ -59,25 +59,25 @@ func (dpc *distinctValuesCache) GetDistinctValues(key string, datasets []string,
 	}
 	// if we were missing any datasets, we need to add them to the cache and unlock the key
 	if len(missing) > 0 {
-		dpc.cacheMu.Lock()
+		dvc.cacheMu.Lock()
 		for _, ds := range missing {
-			dpc.cache[key][ds] = results[ds]
+			dvc.cache[key][ds] = results[ds]
 		}
-		dpc.cacheMu.Unlock()
-		dpc.klock.Unlock(key)
+		dvc.cacheMu.Unlock()
+		dvc.klock.Unlock(key)
 	}
 
 	// Now we need to get values for all the datasets requested so that we can calculate
 	// the distinct values.
 	// We already have some values in the results map (last dataset & missing), so we only need to fill in
 	// the rest of the datasets from the cache.
-	dpc.cacheMu.RLock()
+	dvc.cacheMu.RLock()
 	for _, ds := range datasets {
 		if _, ok := results[ds]; !ok {
-			results[ds] = dpc.cache[key][ds]
+			results[ds] = dvc.cache[key][ds]
 		}
 	}
-	dpc.cacheMu.RUnlock()
+	dvc.cacheMu.RUnlock()
 
 	// Calculating distinct values is easy, we just need to
 	// iterate over all the datasets and add them to a map
@@ -92,17 +92,17 @@ func (dpc *distinctValuesCache) GetDistinctValues(key string, datasets []string,
 }
 
 // RemoveDataset removes the dataset from the cache for all keys.
-func (dpc *distinctValuesCache) RemoveDataset(dataset string) {
-	dpc.cacheMu.Lock()
-	defer dpc.cacheMu.Unlock()
-	for key := range dpc.cache {
-		delete(dpc.cache[key], dataset)
+func (dvc *distinctValuesCache) RemoveDataset(dataset string) {
+	dvc.cacheMu.Lock()
+	defer dvc.cacheMu.Unlock()
+	for key := range dvc.cache {
+		delete(dvc.cache[key], dataset)
 	}
 }
 
-func (dpc *distinctValuesCache) missing(key string, datasets []string) []string {
+func (dvc *distinctValuesCache) missing(key string, datasets []string) []string {
 	var missing []string
-	dscache, ok := dpc.cache[key]
+	dscache, ok := dvc.cache[key]
 	if !ok {
 		return datasets
 	}
diff --git a/jobsdb/jobsdb.go b/jobsdb/jobsdb.go
@@ -1919,7 +1919,14 @@ FROM pending GROUP BY workspace_id, custom_val`
 	return g.Wait()
 }
 
-var parametersWithoutCustomval string = `SELECT '%[2]s', * FROM (
+func (jd *Handle) getDistinctValuesPerDataset(
+	ctx context.Context,
+	dsList []string,
+	param ParameterName,
+) (map[string][]string, error) {
+	var queries []string
+	for _, ds := range dsList {
+		queries = append(queries, fmt.Sprintf(`SELECT '%[2]s', * FROM (
 	WITH RECURSIVE t AS (
 		(SELECT %[1]s as parameter FROM %[2]q ORDER BY %[1]s LIMIT 1)
 		UNION ALL
@@ -1931,32 +1938,10 @@ var parametersWithoutCustomval string = `SELECT '%[2]s', * FROM (
 			)s
 		)
 	)
-SELECT * FROM t) a`
-
-// var parametersWithCustomVal string = `SELECT '%[2]s', * FROM (
-// 	WITH RECURSIVE t AS (
-// 		(SELECT %[1]s as parameter FROM %[2]q WHERE custom_val = '%[3]s' ORDER BY %[1]s LIMIT 1)
-// 		UNION ALL
-// 		(
-// 			SELECT s.* FROM t, LATERAL(
-// 				SELECT %[1]s as parameter FROM %[2]q f
-// 				WHERE custom_val = '%[3]s' AND f.%[1]s > t.parameter
-// 				ORDER BY %[1]s LIMIT 1
-// 			)s
-// 		)
-// 	)
-// SELECT * FROM t) a`
-
-func (jd *Handle) getDistinctValuesPerDataset(
-	dsList []string,
-	param ParameterName,
-) (map[string][]string, error) {
-	var queries []string
-	for _, ds := range dsList {
-		queries = append(queries, fmt.Sprintf(parametersWithoutCustomval, param.string(), ds))
+SELECT * FROM t) a`, param.string(), ds))
 	}
 	query := strings.Join(queries, " UNION ")
-	rows, err := jd.dbHandle.Query(query)
+	rows, err := jd.dbHandle.QueryContext(ctx, query)
 	if err != nil {
 		return nil, fmt.Errorf("couldn't query distinct parameter-%s: %w", param.string(), err)
 	}
@@ -1993,7 +1978,7 @@ func (jd *Handle) GetDistinctParameterValues(ctx context.Context, parameter Para
 		parameter.string(),
 		lo.Map(dsList, func(ds dataSetT, _ int) string { return ds.JobTable }),
 		func(datasets []string) (map[string][]string, error) {
-			return jd.getDistinctValuesPerDataset(datasets, parameter)
+			return jd.getDistinctValuesPerDataset(ctx, datasets, parameter)
 		},
 	)
 }
diff --git a/processor/processor_isolation_test.go b/processor/processor_isolation_test.go
@@ -48,7 +48,11 @@ func TestProcessorIsolation(t *testing.T) {
 		pipelinesPerPartition = 3
 	)
 
-	for _, isolationMode := range []isolation.Mode{isolation.ModeNone, isolation.ModeWorkspace, isolation.ModeSource} {
+	for _, isolationMode := range []isolation.Mode{
+		isolation.ModeNone,
+		isolation.ModeWorkspace,
+		isolation.ModeSource,
+	} {
 		t.Run(fmt.Sprintf("%s isolation", isolationMode), func(t *testing.T) {
 			t.Run("1 worker", func(t *testing.T) {
 				ProcIsolationScenario(t, NewProcIsolationScenarioSpec(isolationMode, workspaces, jobsPerWorkspace, 1))
@@ -212,6 +216,10 @@ func ProcIsolationScenario(t testing.TB, spec *ProcIsolationScenarioSpec) (overa
 		"workspaces":    spec.workspaces,
 		"destinationId": destinationID,
 	}
+	var destinationIDs []string
+	for i := 0; i < len(spec.workspaces); i++ {
+		destinationIDs = append(destinationIDs, destinationID+fmt.Sprintf("-%d", i))
+	}
 	configJsonPath := workspaceConfig.CreateTempFile(t, "testdata/procIsolationTestTemplate.json.tpl", templateCtx)
 	mockCBE := m.newMockConfigBackend(t, configJsonPath)
 	config.Set("CONFIG_BACKEND_URL", mockCBE.URL)
@@ -237,7 +245,7 @@ func ProcIsolationScenario(t testing.TB, spec *ProcIsolationScenarioSpec) (overa
 	config.Set("JobsDB.backup.enabled", false)
 	config.Set("JobsDB.migrateDSLoopSleepDuration", "60m")
 	config.Set("JobsDB.payloadColumnType", "text")
-	config.Set("Router.toAbortDestinationIDs", destinationID)
+	config.Set("Router.toAbortDestinationIDs", destinationIDs)
 	config.Set("archival.Enabled", false)
 	config.Set("enableStats", false)
 
diff --git a/processor/testdata/procIsolationTestTemplate.json.tpl b/processor/testdata/procIsolationTestTemplate.json.tpl
@@ -24,7 +24,7 @@
                             "webhookMethod": "POST"
                         },
                         "secretConfig": {},
-                        "id": "{{$.destinationId}}",
+                        "id": "{{$.destinationId}}-{{$index}}",
                         "name": "Des WebHook Integration Test 1",
                         "enabled": true,
                         "workspaceId": "{{$workspace}}",

Original file line number	Diff line number	Diff line change
`@@ -17,6 +17,7 @@`
`17`	`17`	`{`
`18`	`18`	`"id": "xxxyyyzzP9kQfzOoKd1tuxchYAG",`
`19`	`19`	`"name": "Dev WebHook Integration Test 1",`
	`20`	`+ "workspaceId": "{{.workspaceId}}",`
`20`	`21`	`"enabled": true,`
`21`	`22`	`"isProcessorEnabled": true,`
`22`	`23`	`"config": {`