Create Separate Accessors For Trace And Dependency Stores

Signed-off-by: Mahad Zaryab <[email protected]>

Author: mahadzaryab1

cmd/jaeger/internal/exporters/storageexporter/exporter.go

@@ -34,7 +34,7 @@ func newExporter(config *Config, otel component.TelemetrySettings) *storageExpor
 }
 
 func (exp *storageExporter) start(_ context.Context, host component.Host) error {
-	f, err := jaegerstorage.GetStorageFactoryV2(exp.config.TraceStorage, host)
+	f, err := jaegerstorage.GetTraceStoreFactory(exp.config.TraceStorage, host)
 	if err != nil {
 		return fmt.Errorf("cannot find storage factory: %w", err)
 	}

cmd/jaeger/internal/extension/jaegerquery/server.go

@@ -71,17 +71,20 @@ func (s *server) Start(ctx context.Context, host component.Host) error {
 	telset.Metrics = telset.Metrics.
 		Namespace(metrics.NSOptions{Name: "jaeger"}).
 		Namespace(metrics.NSOptions{Name: "query"})
-	f, err := jaegerstorage.GetStorageFactoryV2(s.config.Storage.TracesPrimary, host)
+	tf, err := jaegerstorage.GetTraceStoreFactory(s.config.Storage.TracesPrimary, host)
 	if err != nil {
-		return fmt.Errorf("cannot find factory for primary storage %s: %w", s.config.Storage.TracesPrimary, err)
+		return fmt.Errorf("cannot find factory for trace storage %s: %w", s.config.Storage.TracesPrimary, err)
 	}
-
-	traceReader, err := f.CreateTraceReader()
+	traceReader, err := tf.CreateTraceReader()
 	if err != nil {
 		return fmt.Errorf("cannot create trace reader: %w", err)
 	}
 
-	depReader, err := f.CreateDependencyReader()
+	df, err := jaegerstorage.GetDependencyStoreFactory(s.config.Storage.TracesPrimary, host)
+	if err != nil {
+		return fmt.Errorf("cannot find factory for dependency storage %s: %w", s.config.Storage.TracesPrimary, err)
+	}
+	depReader, err := df.CreateDependencyReader()
 	if err != nil {
 		return fmt.Errorf("cannot create dependencies reader: %w", err)
 	}

cmd/jaeger/internal/extension/jaegerquery/server_test.go

@@ -150,7 +150,7 @@ func TestServerStart(t *testing.T) {
 					TracesPrimary: "need-factory-error",
 				},
 			},
-			expectedErr: "cannot find factory for primary storage",
+			expectedErr: "cannot find factory for trace storage",
 		},
 		{
 			name: "span reader error",

cmd/jaeger/internal/extension/jaegerstorage/extension.go

@@ -21,7 +21,9 @@ import (
 	"github.com/jaegertracing/jaeger/plugin/storage/grpc"
 	"github.com/jaegertracing/jaeger/plugin/storage/memory"
 	"github.com/jaegertracing/jaeger/storage"
+	"github.com/jaegertracing/jaeger/storage_v2/depstore"
 	"github.com/jaegertracing/jaeger/storage_v2/factoryadapter"
+	"github.com/jaegertracing/jaeger/storage_v2/tracestore"
 )
 
 var _ Extension = (*storageExt)(nil)
@@ -73,7 +75,15 @@ func GetMetricStorageFactory(name string, host component.Host) (storage.MetricSt
 	return mf, nil
 }
 
-func GetStorageFactoryV2(name string, host component.Host) (*factoryadapter.Factory, error) {
+func GetTraceStoreFactory(name string, host component.Host) (tracestore.Factory, error) {
+	return getV2FactoryAdapter(name, host)
+}
+
+func GetDependencyStoreFactory(name string, host component.Host) (depstore.Factory, error) {
+	return getV2FactoryAdapter(name, host)
+}
+
+func getV2FactoryAdapter(name string, host component.Host) (*factoryadapter.Factory, error) {
 	f, err := GetStorageFactory(name, host)
 	if err != nil {
 		return nil, err

cmd/jaeger/internal/extension/jaegerstorage/extension_test.go

@@ -100,7 +100,7 @@ func TestStorageFactoryBadShutdownError(t *testing.T) {
 
 func TestGetFactoryV2Error(t *testing.T) {
 	host := componenttest.NewNopHost()
-	_, err := GetStorageFactoryV2("something", host)
+	_, err := GetTraceStoreFactory("something", host)
 	require.ErrorContains(t, err, "cannot find extension")
 }
 
@@ -112,7 +112,7 @@ func TestGetFactory(t *testing.T) {
 	require.NoError(t, err)
 	require.NotNil(t, f)
 
-	f2, err := GetStorageFactoryV2(name, host)
+	f2, err := GetTraceStoreFactory(name, host)
 	require.NoError(t, err)
 	require.NotNil(t, f2)