Switched dg tweak schema to operate on struct instead of encoded (#71)

Signed-off-by: Jimmy Moore <[email protected]>
loopholelabs · Jan 7, 2025 · dedc200 · dedc200
1 parent 065dfce
commit dedc200
Show file tree

Hide file tree

Showing 3 changed files with 11 additions and 11 deletions.
diff --git a/cmd/connect.go b/cmd/connect.go
@@ -13,6 +13,7 @@ import (
 
 	"github.com/loopholelabs/logging"
 	"github.com/loopholelabs/logging/types"
+	"github.com/loopholelabs/silo/pkg/storage/config"
 	"github.com/loopholelabs/silo/pkg/storage/devicegroup"
 	"github.com/loopholelabs/silo/pkg/storage/metrics"
 	siloprom "github.com/loopholelabs/silo/pkg/storage/metrics/prometheus"
@@ -127,7 +128,7 @@ func runConnect(_ *cobra.Command, _ []string) {
 	}
 
 	// TODO: Modify schemas a bit here...
-	tweak := func(_ int, _ string, schema string) string {
+	tweak := func(_ int, _ string, schema *config.DeviceSchema) *config.DeviceSchema {
 		return schema
 	}
 

diff --git a/pkg/storage/devicegroup/device_group_from.go b/pkg/storage/devicegroup/device_group_from.go
@@ -15,7 +15,7 @@ import (
 
 func NewFromProtocol(ctx context.Context,
 	pro protocol.Protocol,
-	tweakDeviceSchema func(index int, name string, schema string) string,
+	tweakDeviceSchema func(index int, name string, schema *config.DeviceSchema) *config.DeviceSchema,
 	eventHandler func(e *packets.Event),
 	customDataHandler func(data []byte),
 	log types.Logger,
@@ -75,14 +75,13 @@ func NewFromProtocol(ctx context.Context,
 	// First create the devices we need using the schemas sent...
 	for index, di := range dgi.Devices {
 		// We may want to tweak schemas here eg autoStart = false on sync. Or modify pathnames.
-		schema := di.Schema
-		if tweakDeviceSchema != nil {
-			schema = tweakDeviceSchema(index-1, di.Name, schema)
-		}
-		ds, err := config.DecodeDeviceFromBlock(schema)
+		ds, err := config.DecodeDeviceFromBlock(di.Schema)
 		if err != nil {
 			return nil, err
 		}
+		if tweakDeviceSchema != nil {
+			ds = tweakDeviceSchema(index-1, di.Name, ds)
+		}
 		devices[index-1] = ds
 	}
 

diff --git a/pkg/storage/devicegroup/device_group_test.go b/pkg/storage/devicegroup/device_group_test.go
@@ -249,10 +249,10 @@ func TestDeviceGroupMigrate(t *testing.T) {
 	var wg sync.WaitGroup
 
 	// We will tweak schema in recv here so we have separate paths.
-	tweak := func(_ int, _ string, schema string) string {
-		s := strings.ReplaceAll(schema, "testdev_test1", "testrecv_test1")
-		s = strings.ReplaceAll(s, "testdev_test2", "testrecv_test2")
-		return s
+	tweak := func(_ int, _ string, schema *config.DeviceSchema) *config.DeviceSchema {
+		schema.Location = strings.ReplaceAll(schema.Location, "testdev_test1", "testrecv_test1")
+		schema.Location = strings.ReplaceAll(schema.Location, "testdev_test2", "testrecv_test2")
+		return schema
 	}
 
 	// TransferAuthority