Metric context query

pkg/api/mlflow/dao/repositories/helpers.go

@@ -4,6 +4,9 @@ import (
 	"fmt"
 	"strings"
 
+	"gorm.io/driver/postgres"
+	"gorm.io/gorm"
+
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/models"
 )
 
@@ -30,3 +33,22 @@ func makeParamConflictPlaceholdersAndValues(params []models.Param) (string, []in
 	}
 	return placeholders, valuesArray
 }
+
+// addJsonCondition adds condition(s) to the query to select items having the specified map of json paths
+// and values. Json path is expressed as "outerkey.nestedKey".
+func addJsonCondition(tx *gorm.DB, jsonColumnName string, jsonPathValueMap map[string]any) {
+	if len(jsonPathValueMap) == 0 {
+		return
+	}
+	switch tx.Dialector.Name() {
+	case postgres.Dialector{}.Name():
+		for k, v := range jsonPathValueMap {
+			path := strings.ReplaceAll(k, ".", ",")
+			tx.Where(fmt.Sprintf("%s#>>'{%s}' = ?", jsonColumnName, path), v)
+		}
+	default:
+		for k, v := range jsonPathValueMap {
+			tx.Where(fmt.Sprintf("%s->>'%s' = ?", jsonColumnName, k), v)
+		}
+	}
+}

pkg/api/mlflow/dao/repositories/metric.go

@@ -30,6 +30,7 @@ type MetricRepositoryProvider interface {
 		experimentIDs []string, runIDs []string, metricKeys []string,
 		viewType request.ViewType,
 		limit int32,
+		jsonPathValueMap map[string]any,
 	) (*sql.Rows, func(*sql.Rows, interface{}) error, error)
 	// GetMetricHistoryBulk returns metrics history bulk.
 	GetMetricHistoryBulk(
@@ -151,6 +152,7 @@ func (r MetricRepository) GetMetricHistories(
 	experimentIDs []string, runIDs []string, metricKeys []string,
 	viewType request.ViewType,
 	limit int32,
+	jsonPathValueMap map[string]any,
 ) (*sql.Rows, func(*sql.Rows, interface{}) error, error) {
 	// if experimentIDs has been provided then firstly get the runs by provided experimentIDs.
 	if len(experimentIDs) > 0 {
@@ -219,6 +221,11 @@ func (r MetricRepository) GetMetricHistories(
 		query.Where("metrics.key IN ?", metricKeys)
 	}
 
+	if len(jsonPathValueMap) > 0 {
+		query.Joins("LEFT JOIN contexts on metrics.context_id = contexts.id")
+		addJsonCondition(query, "contexts.json", jsonPathValueMap)
+	}
+
 	rows, err := query.Rows()
 	if err != nil {
 		return nil, nil, eris.Wrapf(

pkg/api/mlflow/service/metric/service.go

@@ -89,6 +89,7 @@ func (s Service) GetMetricHistories(
 		req.MetricKeys,
 		req.ViewType,
 		req.MaxResults,
+		map[string]any{},
 	)
 	if err != nil {
 		return nil, nil, api.NewInternalError("Unable to search runs: %s", err)

pkg/api/mlflow/service/metric/service_test.go

@@ -335,6 +335,7 @@ func TestNewService_GetMetricHistories_Ok(t *testing.T) {
 				[]string{"key1", "key2"},
 				request.ViewTypeActiveOnly,
 				int32(1),
+				map[string]any{},
 			).Return(
 				tt.expectedRows,
 				tt.expectedIter,
@@ -423,6 +424,7 @@ func TestNewService_GetMetricHistories_Error(t *testing.T) {
 					[]string{"key1", "key2"},
 					request.ViewTypeAll,
 					int32(1),
+					map[string]any{},
 				).Return(
 					nil,
 					nil,

tests/integration/golang/fixtures/base.go

@@ -22,6 +22,7 @@ func (f baseFixtures) TruncateTables() error {
 		models.Param{},
 		models.LatestMetric{},
 		models.Metric{},
+		models.Context{},
 		models.Run{},
 		models.ExperimentTag{},
 		models.Experiment{},

tests/integration/golang/fixtures/metric.go

@@ -2,12 +2,16 @@ package fixtures
 
 import (
 	"context"
+	"fmt"
+	"strings"
 
 	"github.com/rotisserie/eris"
+	"gorm.io/driver/postgres"
 	"gorm.io/gorm"
 
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/models"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/dao/repositories"
+	"github.com/G-Research/fasttrackml/pkg/database"
 )
 
 // MetricFixtures represents data fixtures object.
@@ -26,18 +30,59 @@ func NewMetricFixtures(db *gorm.DB) (*MetricFixtures, error) {
 
 // CreateMetric creates new test Metric.
 func (f MetricFixtures) CreateMetric(ctx context.Context, metric *models.Metric) (*models.Metric, error) {
+	if metric.Context != nil {
+		if err := f.baseFixtures.db.WithContext(ctx).FirstOrCreate(&metric.Context).Error; err != nil {
+			return nil, eris.Wrap(err, "error creating metric context")
+		}
+		metric.ContextID = &metric.Context.ID
+	}
 	if err := f.baseFixtures.db.WithContext(ctx).Create(metric).Error; err != nil {
-		return nil, eris.Wrap(err, "error creating test metric")
+		return nil, eris.Wrap(err, "error creating metric")
 	}
 	return metric, nil
 }
 
+// GetMetricsByRunID returns the metrics by Run ID.
+func (f MetricFixtures) GetMetricsByRunID(ctx context.Context, runID string) ([]*models.Metric, error) {
+	var metrics []*models.Metric
+	if err := f.db.WithContext(ctx).Where(
+		"run_uuid = ?", runID,
+	).Find(&metrics).Error; err != nil {
+		return nil, eris.Wrapf(err, "error getting metric by run_uuid: %v", runID)
+	}
+	return metrics, nil
+}
+
+// GetMetricsByContext returns metric by a context partial match.
+func (f MetricFixtures) GetMetricsByContext(
+	ctx context.Context,
+	metricContext map[string]any,
+) ([]*models.Metric, error) {
+	var metrics []*models.Metric
+	tx := f.db.WithContext(ctx).Model(
+		&database.Metric{},
+	).Joins(
+		"LEFT JOIN contexts on metrics.context_id = contexts.id",
+	)
+	addContextSelection(tx, "contexts.json", metricContext)
+	if err := tx.Find(&metrics).Error; err != nil {
+		return nil, eris.Wrapf(err, "error getting metric by context: %v", metricContext)
+	}
+	return metrics, nil
+}
+
 // CreateLatestMetric creates new test Latest Metric.
 func (f MetricFixtures) CreateLatestMetric(
 	ctx context.Context, metric *models.LatestMetric,
 ) (*models.LatestMetric, error) {
+	if metric.Context != nil {
+		if err := f.baseFixtures.db.WithContext(ctx).FirstOrCreate(&metric.Context).Error; err != nil {
+			return nil, eris.Wrap(err, "error creating latest metric context")
+		}
+		metric.ContextID = &metric.Context.ID
+	}
 	if err := f.baseFixtures.db.WithContext(ctx).Create(metric).Error; err != nil {
-		return nil, eris.Wrap(err, "error creating test latest metric")
+		return nil, eris.Wrap(err, "error creating latest metric")
 	}
 	return metric, nil
 }
@@ -63,3 +108,21 @@ func (f MetricFixtures) GetLatestMetricByRunID(ctx context.Context, runID string
 	}
 	return &metric, nil
 }
+
+// addContextSelection adds conditions to the query to select metrics having the provided context.
+func addContextSelection(tx *gorm.DB, columnName string, jsonPathValueMap map[string]any) {
+	if len(jsonPathValueMap) == 0 {
+		return
+	}
+	switch tx.Dialector.Name() {
+	case postgres.Dialector{}.Name():
+		for k, v := range jsonPathValueMap {
+			path := strings.ReplaceAll(k, ".", ",")
+			tx.Where(fmt.Sprintf("%s#>>'{%s}' = ?", columnName, path), v)
+		}
+	default:
+		for k, v := range jsonPathValueMap {
+			tx.Where(fmt.Sprintf("%s->>'%s' = ?", columnName, k), v)
+		}
+	}
+}

tests/integration/golang/mlflow/metric/get_histories_test.go

@@ -10,6 +10,7 @@ import (
 	"testing"
 
 	"github.com/stretchr/testify/suite"
+	"gorm.io/datatypes"
 
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow"
 	"github.com/G-Research/fasttrackml/pkg/api/mlflow/api"
@@ -45,16 +46,64 @@ func (s *GetHistoriesTestSuite) Test_Ok() {
 	})
 	s.Require().Nil(err)
 
-	_, err = s.MetricFixtures.CreateMetric(context.Background(), &models.Metric{
+	m, err := s.MetricFixtures.CreateMetric(context.Background(), &models.Metric{
 		Key:       "key1",
 		Value:     1.1,
 		Timestamp: 1234567890,
 		RunID:     run1.ID,
 		Step:      1,
-		IsNan:     false,
 		Iter:      1,
 	})
 	s.Require().Nil(err)
+	s.Require().Nil(m.ContextID)
+	s.Require().Nil(m.Context)
+
+	m, err = s.MetricFixtures.CreateMetric(context.Background(), &models.Metric{
+		Key:       "key2",
+		Value:     1.1,
+		Timestamp: 2234567890,
+		RunID:     run1.ID,
+		Step:      1,
+		Iter:      1,
+		Context: &models.Context{
+			Json: datatypes.JSON([]byte(`
+				{
+					"metrickey1": "metricvalue1",
+					"metrickey2": "metricvalue2",
+					"metricnested": { "metricnestedkey": "metricnestedvalue" }
+				}`,
+			)),
+		},
+	})
+	s.Require().Nil(err)
+	s.Require().NotNil(m.ContextID)
+	s.Require().NotNil(m.Context)
+
+	// verify metric contexts are persisting
+	metrics, err := s.MetricFixtures.GetMetricsByRunID(context.Background(), run1.ID)
+	s.Require().Nil(err)
+	s.Require().Nil(metrics[0].ContextID)
+	s.Require().NotNil(metrics[1].ContextID)
+
+	// verify metric contexts can be used for selection (toplevel key)
+	metrics, err = s.MetricFixtures.GetMetricsByContext(context.Background(), map[string]any{
+		"metrickey1": "metricvalue1",
+	})
+	s.Require().Nil(err)
+	s.Require().Len(metrics, 1)
+	s.Require().NotNil(metrics[0].ContextID)
+
+	// nested key
+	metrics, err = s.MetricFixtures.GetMetricsByContext(context.Background(), map[string]any{
+		"metricnested.metricnestedkey": "metricnestedvalue",
+	})
+	s.Require().Nil(err)
+	s.Require().Len(metrics, 1)
+	s.Require().NotNil(metrics[0].ContextID)
+
+	metrics, err = s.MetricFixtures.GetMetricsByContext(context.Background(), map[string]any{"metrickey2": "metricvalue1"})
+	s.Require().Nil(err)
+	s.Require().Len(metrics, 0)
 
 	run2, err := s.RunFixtures.CreateRun(context.Background(), &models.Run{
 		ID:             "run2",