Repository: G-Research/fasttrackml
Pull request #821 (2024-01-22 15:59:16)

Bump go-sqlite3 library to SQLite3MultipleCiphers 1.8.2 / SQLite3 3.45.0
This PR bumps up the SQLite3 library to latest version which introduces support for JSONB. The way to use JSON paths has also changed (it looks like the way we used them before was not supposed to work in the first place) so this required changing a couple of things in the code.

Once this PR is merged, we should consider using JSONB for contexts instead of the current JSON implementation.


Patch:
diff --git a/go.mod b/go.mod
index 30b1292..d6c1f53 100644
--- a/go.mod
+++ b/go.mod
@@ -3,7 +3,7 @@ module github.com/G-Research/fasttrackml
 go 1.21.0
 
 replace (
-	github.com/mattn/go-sqlite3 v1.14.16 => github.com/jgiannuzzi/go-sqlite3 v1.14.17-0.20230327164124-765f25ea5431
+	github.com/mattn/go-sqlite3 v1.14.16 => github.com/jgiannuzzi/go-sqlite3 v1.14.17-0.20240122133042-fb824c8e339e
 	gorm.io/driver/sqlite v1.4.3 => github.com/jgiannuzzi/gorm-sqlite v1.4.4-0.20221215225833-42389ad31305
 )
 
diff --git a/go.sum b/go.sum
index ee1be2f..802f799 100644
--- a/go.sum
+++ b/go.sum
@@ -175,8 +175,8 @@ github.com/jackc/pgx/v5 v5.5.2 h1:iLlpgp4Cp/gC9Xuscl7lFL1PhhW+ZLtXZcrfCt4C3tA=
 github.com/jackc/pgx/v5 v5.5.2/go.mod h1:ez9gk+OAat140fv9ErkZDYFWmXLfV+++K0uAOiwgm1A=
 github.com/jackc/puddle/v2 v2.2.1 h1:RhxXJtFG022u4ibrCSMSiu5aOq1i77R3OHKNJj77OAk=
 github.com/jackc/puddle/v2 v2.2.1/go.mod h1:vriiEXHvEE654aYKXXjOvZM39qJ0q+azkZFrfEOc3H4=
-github.com/jgiannuzzi/go-sqlite3 v1.14.17-0.20230327164124-765f25ea5431 h1:zveJxnAiwI63z7qf5ndPbJxRIO7ZC8wFiDDzZ4WXVF4=
-github.com/jgiannuzzi/go-sqlite3 v1.14.17-0.20230327164124-765f25ea5431/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
+github.com/jgiannuzzi/go-sqlite3 v1.14.17-0.20240122133042-fb824c8e339e h1:0s+RSCZSEaZ3IM4fqMH1Nf5UDpttrn1R2QDutqiL2Zc=
+github.com/jgiannuzzi/go-sqlite3 v1.14.17-0.20240122133042-fb824c8e339e/go.mod h1:2eHXhiwb8IkHr+BDWZGa96P6+rkvnG63S2DGjv9HUNg=
 github.com/jgiannuzzi/gorm-sqlite v1.4.4-0.20221215225833-42389ad31305 h1:At4YCq3pSf9pDEkd47r7CGh3A/j6hMD8hOKSt14fIgk=
 github.com/jgiannuzzi/gorm-sqlite v1.4.4-0.20221215225833-42389ad31305/go.mod h1:0Aq3iPO+v9ZKbcdiz8gLWRw5VOPcBOPUQJFLq5e2ecI=
 github.com/jinzhu/inflection v1.0.0 h1:K317FqzuhWc8YvSVlFMCCUb36O/S9MCKRDI7QkRKD/E=
diff --git a/pkg/api/aim/query/clause.go b/pkg/api/aim/query/clause.go
index 945fb4f..9dcdce6 100644
--- a/pkg/api/aim/query/clause.go
+++ b/pkg/api/aim/query/clause.go
@@ -2,6 +2,7 @@ package query
 
 import (
 	"database/sql/driver"
+	"fmt"
 	"reflect"
 	"strings"
 
@@ -67,33 +68,29 @@ type Json struct {
 // Build builds positive statement.
 func (json Json) Build(builder clause.Builder) {
 	json.writeColumn(builder)
-	jsonPath := json.JsonPath
 	switch json.Dialector {
 	case postgres.Dialector{}.Name():
 		//nolint:errcheck,gosec
 		builder.WriteString("#>>")
-		jsonPath = "{" + strings.ReplaceAll(jsonPath, ",", ".") + "}"
 	default:
 		//nolint:errcheck,gosec
 		builder.WriteString("->>")
 	}
-	builder.AddVar(builder, jsonPath)
+	builder.AddVar(builder, json.jsonPathForDialect())
 }
 
 // NegationBuild builds negative statement.
 func (json Json) NegationBuild(builder clause.Builder) {
 	json.writeColumn(builder)
-	jsonPath := json.JsonPath
 	switch json.Dialector {
 	case postgres.Dialector{}.Name():
 		//nolint:errcheck,gosec
 		builder.WriteString("#>>")
-		jsonPath = "{" + strings.ReplaceAll(jsonPath, ",", ".") + "}"
 	default:
 		//nolint:errcheck,gosec
 		builder.WriteString("->>")
 	}
-	builder.AddVar(builder, jsonPath)
+	builder.AddVar(builder, json.jsonPathForDialect())
 }
 
 func (json Json) writeColumn(builder clause.Builder) {
@@ -109,6 +106,16 @@ func (json Json) writeColumn(builder clause.Builder) {
 	}
 }
 
+func (json Json) jsonPathForDialect() string {
+	switch json.Dialector {
+	case postgres.Dialector{}.Name():
+		jsonPath := removePrefix(json.JsonPath)
+		return "{" + strings.ReplaceAll(jsonPath, ",", ".") + "}"
+	default:
+		return addPrefix(json.JsonPath)
+	}
+}
+
 type JsonEq struct {
 	Left  Json
 	Value any
@@ -186,6 +193,19 @@ func (neq JsonNeq) NegationBuild(builder clause.Builder) {
 	JsonEq(neq).Build(builder)
 }
 
+// addPrefix adds leading $. to a jsonPath if needed
+func addPrefix(jsonPath string) string {
+	if strings.HasPrefix(jsonPath, "$.") {
+		return jsonPath
+	}
+	return fmt.Sprintf("$.%s", jsonPath)
+}
+
+// removePrefix removes a leading $. for a jsonPath
+func removePrefix(jsonPath string) string {
+	return strings.TrimPrefix(jsonPath, "$.")
+}
+
 func eqNil(value interface{}) bool {
 	if valuer, ok := value.(driver.Valuer); ok && !eqNilReflect(valuer) {
 		//nolint:errcheck
diff --git a/pkg/api/mlflow/dao/repositories/helpers.go b/pkg/api/mlflow/dao/repositories/helpers.go
index 095b224..10d3d5e 100644
--- a/pkg/api/mlflow/dao/repositories/helpers.go
+++ b/pkg/api/mlflow/dao/repositories/helpers.go
@@ -59,7 +59,7 @@ func BuildJsonCondition(
 		conditionTemplate = "%s->>? = ?"
 		idx := 0
 		for k, v := range jsonPathValueMap {
-			args[idx] = k
+			args[idx] = fmt.Sprintf("$.%s", k)
 			args[idx+1] = v
 			idx = idx + 2
 		}


Test patch:
diff --git a/pkg/api/aim/query/query_test.go b/pkg/api/aim/query/query_test.go
index b5b9306..54ba6e4 100644
--- a/pkg/api/aim/query/query_test.go
+++ b/pkg/api/aim/query/query_test.go
@@ -148,6 +148,16 @@ func (s *QueryTestSuite) TestPostgresDialector_Ok() {
 				`AND ("metrics_0"."value" < $4 AND "runs"."lifecycle_stage" <> $5)`,
 			expectedVars: []interface{}{"my_metric", "{key1}", "value1", -1, models.LifecycleStageDeleted},
 		},
+		{
+			name:  "TestMetricContextSliceTupleWithPrefix",
+			query: `run.metrics["my_metric", {"$.key1": "value1"}].last < -1`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
+				`LEFT JOIN latest_metrics metrics_0 ON runs.run_uuid = metrics_0.run_uuid AND metrics_0.key = $1 ` +
+				`LEFT JOIN contexts contexts_1 ON metrics_0.context_id = contexts_1.id ` +
+				`WHERE "contexts_1"."json"#>>$2 = $3 ` +
+				`AND ("metrics_0"."value" < $4 AND "runs"."lifecycle_stage" <> $5)`,
+			expectedVars: []interface{}{"my_metric", "{key1}", "value1", -1, models.LifecycleStageDeleted},
+		},
 		{
 			name:  "TestMetricCompound",
 			query: `run.metrics["my_metric"].last < -1 and metric.context.key1 == "value1"`,
@@ -285,8 +295,9 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 			query:         `metric.context.key1 == 'value1'`,
 			selectMetrics: true,
 			expectedSQL: `SELECT ID FROM "metrics" ` +
+
 				`WHERE IFNULL("contexts"."json", JSON('{}'))->>$1 = $2 AND "runs"."lifecycle_stage" <> $3`,
-			expectedVars: []interface{}{"key1", "value1", models.LifecycleStageDeleted},
+			expectedVars: []interface{}{"$.key1", "value1", models.LifecycleStageDeleted},
 		},
 		{
 			name:          "TestMetricContextNegative",
@@ -294,7 +305,7 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 			selectMetrics: true,
 			expectedSQL: `SELECT ID FROM "metrics" ` +
 				`WHERE IFNULL("contexts"."json", JSON('{}'))->>$1 <> $2 AND "runs"."lifecycle_stage" <> $3`,
-			expectedVars: []interface{}{"key1", "value1", models.LifecycleStageDeleted},
+			expectedVars: []interface{}{"$.key1", "value1", models.LifecycleStageDeleted},
 		},
 		{
 			name:  "TestMetricKeySlice",
@@ -312,7 +323,17 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 				`LEFT JOIN contexts contexts_1 ON metrics_0.context_id = contexts_1.id ` +
 				`WHERE IFNULL("contexts_1"."json", JSON('{}'))->>$2 = $3 ` +
 				`AND ("metrics_0"."value" < $4 AND "runs"."lifecycle_stage" <> $5)`,
-			expectedVars: []interface{}{"my_metric", "key1", "value1", -1, models.LifecycleStageDeleted},
+			expectedVars: []interface{}{"my_metric", "$.key1", "value1", -1, models.LifecycleStageDeleted},
+		},
+		{
+			name:  "TestMetricContextSliceTupleWithPrefix",
+			query: `run.metrics["my_metric", {"$.key1": "value1"}].last < -1`,
+			expectedSQL: `SELECT "run_uuid" FROM "runs" ` +
+				`LEFT JOIN latest_metrics metrics_0 ON runs.run_uuid = metrics_0.run_uuid AND metrics_0.key = $1 ` +
+				`LEFT JOIN contexts contexts_1 ON metrics_0.context_id = contexts_1.id ` +
+				`WHERE IFNULL("contexts_1"."json", JSON('{}'))->>$2 = $3 ` +
+				`AND ("metrics_0"."value" < $4 AND "runs"."lifecycle_stage" <> $5)`,
+			expectedVars: []interface{}{"my_metric", "$.key1", "value1", -1, models.LifecycleStageDeleted},
 		},
 		{
 			name:  "TestMetricCompound",
@@ -321,7 +342,7 @@ func (s *QueryTestSuite) TestSqliteDialector_Ok() {
 				`LEFT JOIN latest_metrics metrics_0 ON runs.run_uuid = metrics_0.run_uuid AND metrics_0.key = $1 ` +
 				`WHERE ("metrics_0"."value" < $2 AND IFNULL("contexts"."json", JSON('{}'))->>$3 = $4) ` +
 				`AND "runs"."lifecycle_stage" <> $5`,
-			expectedVars: []interface{}{"my_metric", -1, "key1", "value1", models.LifecycleStageDeleted},
+			expectedVars: []interface{}{"my_metric", -1, "$.key1", "value1", models.LifecycleStageDeleted},
 		},
 	}
 
diff --git a/pkg/api/mlflow/dao/repositories/helpers_test.go b/pkg/api/mlflow/dao/repositories/helpers_test.go
index e099cfe..9cb46f7 100644
--- a/pkg/api/mlflow/dao/repositories/helpers_test.go
+++ b/pkg/api/mlflow/dao/repositories/helpers_test.go
@@ -84,7 +84,7 @@ func TestBuildJsonCondition(t *testing.T) {
 				"key2.nested": "value2",
 			},
 			expectedSQL:  "contexts.json->>? = ? AND contexts.json->>? = ?",
-			expectedArgs: []interface{}{"key1", "value1", "key2.nested", "value2"},
+			expectedArgs: []interface{}{"$.key1", "value1", "$.key2.nested", "value2"},
 		},
 		{
 			name:             "SqliteEmptyMap",