Support model in GenerateMetricsViewFile

proto/gen/rill/runtime/v1/runtime.swagger.yaml

@@ -677,18 +677,39 @@ paths:
           schema:
             type: object
             properties:
+              model:
+                type: string
+                description: |-
+                  Model to base the metrics view on.
+                  If you set this, do NOT set connector, database, database_schema or table.
               connector:
                 type: string
+                description: |-
+                  Connector for the table.
+                  See "table" for more details.
               database:
                 type: string
+                description: |-
+                  Database for the table.
+                  See "table" for more details.
               databaseSchema:
                 type: string
+                description: |-
+                  Database schema for the table.
+                  See "table" for more details.
               table:
                 type: string
+                description: |-
+                  Table to base the metrics view on.
+                  If you set this, do NOT set model.
               path:
                 type: string
+                description: Path to save the metrics view file to.
               useAi:
                 type: boolean
+                description: |-
+                  If true, the AI will be used to generate the metrics view file.
+                  Otherwise, it falls back to a simpler heuristic approach.
             title: Request message for RuntimeService.GenerateMetricsViewFile
       tags:
         - RuntimeService
@@ -4439,6 +4460,7 @@ definitions:
     properties:
       aiSucceeded:
         type: boolean
+        description: Indicates if AI-based generation succeeded. If it failed, it falls back to the simpler heuristic approach.
     title: Response message for RuntimeService.GenerateMetricsViewFile
   v1GenerateRendererResponse:
     type: object

proto/rill/runtime/v1/api.proto

@@ -532,16 +532,31 @@ message UnpackEmptyResponse {}
 // Request message for RuntimeService.GenerateMetricsViewFile
 message GenerateMetricsViewFileRequest {
   string instance_id = 1 [(validate.rules).string = {pattern: "^[_\\-a-zA-Z0-9]+$"}];
+  // Model to base the metrics view on.
+  // If you set this, do NOT set connector, database, database_schema or table.
+  string model = 8;
+  // Connector for the table.
+  // See "table" for more details.
   string connector = 2;
+  // Database for the table.
+  // See "table" for more details.
   string database = 6;
+  // Database schema for the table.
+  // See "table" for more details.
   string database_schema = 7;
+  // Table to base the metrics view on.
+  // If you set this, do NOT set model.
   string table = 3;
+  // Path to save the metrics view file to.
   string path = 4;
+  // If true, the AI will be used to generate the metrics view file.
+  // Otherwise, it falls back to a simpler heuristic approach.
   bool use_ai = 5;
 }
 
 // Response message for RuntimeService.GenerateMetricsViewFile
 message GenerateMetricsViewFileResponse {
+  // Indicates if AI-based generation succeeded. If it failed, it falls back to the simpler heuristic approach.
   bool ai_succeeded = 1;
 }
 

runtime/server/generate_metrics_view.go

@@ -34,6 +34,7 @@ const aiGenerateTimeout = 30 * time.Second
 func (s *Server) GenerateMetricsViewFile(ctx context.Context, req *runtimev1.GenerateMetricsViewFileRequest) (*runtimev1.GenerateMetricsViewFileResponse, error) {
 	observability.AddRequestAttributes(ctx,
 		attribute.String("args.instance_id", req.InstanceId),
+		attribute.String("args.model", req.Model),
 		attribute.String("args.connector", req.Connector),
 		attribute.String("args.database", req.Database),
 		attribute.String("args.database_schema", req.DatabaseSchema),
@@ -54,54 +55,83 @@ func (s *Server) GenerateMetricsViewFile(ctx context.Context, req *runtimev1.Gen
 		return nil, err
 	}
 
-	// If a connector is not provided, default to the instance's OLAP connector
-	if req.Connector == "" {
-		req.Connector = inst.ResolveOLAPConnector()
+	// Check that only one of model or table is provided
+	if req.Model != "" && req.Table != "" {
+		return nil, status.Error(codes.InvalidArgument, "only one of model or table can be provided")
+	} else if req.Model == "" && req.Table == "" {
+		return nil, status.Error(codes.InvalidArgument, "either model or table must be provided")
 	}
-	isDefaultConnector := req.Connector == inst.ResolveOLAPConnector()
 
-	// Connect to connector and check it's an OLAP db
-	olap, release, err := s.runtime.OLAP(ctx, req.InstanceId, req.Connector)
-	if err != nil {
-		return nil, err
+	// The `model:` field has fuzzy logic, supporting either models, sources, or external table names.
+	// So we need some similarly fuzzy logic here to determine if there's a matching model.
+	var modelFound bool
+	var modelConnector, modelTable string
+	searchName := req.Table
+	if req.Model != "" {
+		searchName = req.Model
 	}
-	defer release()
-
-	// Get table info
-	tbl, err := olap.InformationSchema().Lookup(ctx, req.Database, req.DatabaseSchema, req.Table)
-	if err != nil {
-		return nil, status.Errorf(codes.InvalidArgument, "table not found: %s", err)
-	}
-
-	// The table may have been created by a model. Search for a model with the same name in the same connector.
-	// NOTE: If it's a source, we will also mark it as a model. The metrics view YAML supports that, and we'll anyway deprecate sources soon.
-	var isModel bool
 	ctrl, err := s.runtime.Controller(ctx, req.InstanceId)
 	if err != nil {
 		return nil, err
 	}
-	model, err := ctrl.Get(ctx, &runtimev1.ResourceName{Kind: runtime.ResourceKindModel, Name: tbl.Name}, false)
+	model, err := ctrl.Get(ctx, &runtimev1.ResourceName{Kind: runtime.ResourceKindModel, Name: searchName}, false)
 	if err != nil && !errors.Is(err, drivers.ErrResourceNotFound) {
 		return nil, err
 	}
 	if model != nil {
-		// Check model is for this table.
-		modelState := model.GetModel().State
-		if modelState.ResultConnector == req.Connector && strings.EqualFold(modelState.ResultTable, tbl.Name) {
-			isModel = true
-		}
+		modelFound = true
+		modelConnector = model.GetModel().State.ResultConnector
+		modelTable = model.GetModel().State.ResultTable
 	} else {
 		// Check if it's a source
-		source, err := ctrl.Get(ctx, &runtimev1.ResourceName{Kind: runtime.ResourceKindSource, Name: tbl.Name}, false)
+		source, err := ctrl.Get(ctx, &runtimev1.ResourceName{Kind: runtime.ResourceKindSource, Name: searchName}, false)
 		if err != nil && !errors.Is(err, drivers.ErrResourceNotFound) {
 			return nil, err
 		}
 		if source != nil {
-			sourceState := source.GetSource().State
-			if sourceState.Connector == req.Connector && strings.EqualFold(sourceState.Table, tbl.Name) {
-				isModel = true
-			}
+			modelFound = true
+			modelConnector = source.GetSource().State.Connector
+			modelTable = source.GetSource().State.Table
+		}
+	}
+
+	// Depending on the result, we populate req.Connector and req.Table which we use subsequently.
+	if modelFound {
+		// We found a matching model.
+		if req.Model != "" { // We found req.Model
+			req.Connector = modelConnector
+			req.Table = modelTable
+		} else if (req.Connector == "" || req.Connector == modelConnector) && strings.EqualFold(req.Table, modelTable) { // We found a model with the same name as req.Table
+			req.Connector = modelConnector
+			req.Table = modelTable
+		} else { // We found a model that doesn't match the request.
+			modelFound = false
 		}
+	} else {
+		// We did not find a model. We proceed to check if "model" references an external table in the default OLAP.
+		if req.Model != "" {
+			req.Connector = ""
+			req.Table = req.Model
+		}
+	}
+
+	// If a connector is not provided, default to the instance's OLAP connector
+	if req.Connector == "" {
+		req.Connector = inst.ResolveOLAPConnector()
+	}
+	isDefaultConnector := req.Connector == inst.ResolveOLAPConnector()
+
+	// Connect to connector and check it's an OLAP db
+	olap, release, err := s.runtime.OLAP(ctx, req.InstanceId, req.Connector)
+	if err != nil {
+		return nil, err
+	}
+	defer release()
+
+	// Get table info
+	tbl, err := olap.InformationSchema().Lookup(ctx, req.Database, req.DatabaseSchema, req.Table)
+	if err != nil {
+		return nil, status.Errorf(codes.InvalidArgument, "table not found: %s", err)
 	}
 
 	// Try to generate the YAML with AI
@@ -110,7 +140,7 @@ func (s *Server) GenerateMetricsViewFile(ctx context.Context, req *runtimev1.Gen
 	if req.UseAi {
 		// Generate
 		start := time.Now()
-		res, err := s.generateMetricsViewYAMLWithAI(ctx, req.InstanceId, olap.Dialect().String(), req.Connector, tbl, isDefaultConnector, isModel)
+		res, err := s.generateMetricsViewYAMLWithAI(ctx, req.InstanceId, olap.Dialect().String(), req.Connector, tbl, isDefaultConnector, modelFound)
 		if err != nil {
 			s.logger.Warn("failed to generate metrics view YAML using AI", zap.Error(err))
 		} else {
@@ -138,7 +168,7 @@ func (s *Server) GenerateMetricsViewFile(ctx context.Context, req *runtimev1.Gen
 
 	// If we didn't manage to generate the YAML using AI, we fall back to the simple generator
 	if data == "" {
-		data, err = generateMetricsViewYAMLSimple(req.Connector, tbl, isDefaultConnector, isModel)
+		data, err = generateMetricsViewYAMLSimple(req.Connector, tbl, isDefaultConnector, modelFound)
 		if err != nil {
 			return nil, err
 		}

runtime/server/generate_metrics_view_test.go

@@ -6,6 +6,8 @@ import (
 	"time"
 
 	runtimev1 "github.com/rilldata/rill/proto/gen/rill/runtime/v1"
+	"github.com/rilldata/rill/runtime"
+	"github.com/rilldata/rill/runtime/drivers"
 	"github.com/rilldata/rill/runtime/pkg/activity"
 	"github.com/rilldata/rill/runtime/pkg/ratelimit"
 	"github.com/rilldata/rill/runtime/server"
@@ -15,30 +17,102 @@ import (
 )
 
 func TestGenerateMetricsView(t *testing.T) {
-	rt, instanceID := testruntime.NewInstanceForProject(t, "ad_bids")
+	rt, instanceID := testruntime.NewInstanceWithOptions(t, testruntime.InstanceOptions{
+		Files: map[string]string{
+			"rill.yaml": ``,
+			// Normal model
+			"ad_bids.sql": `SELECT now() AS time, 'DA' AS country, 3.141 as price`,
+			// Create a non-default duckdb connector
+			"custom_duckdb.yaml": `
+type: connector
+driver: duckdb
+`,
+		},
+	})
 
-	server, err := server.NewServer(context.Background(), &server.Options{}, rt, zap.NewNop(), ratelimit.NewNoop(), activity.NewNoopClient())
-	require.NoError(t, err)
+	// Create some externally managed tables
+	olapExecAdhoc(t, rt, instanceID, "duckdb", "CREATE TABLE IF NOT EXISTS foo AS SELECT now() AS time, 'DA' AS country, 3.141 as price")
+	olapExecAdhoc(t, rt, instanceID, "custom_duckdb", "CREATE TABLE IF NOT EXISTS foo AS SELECT now() AS time, 'DA' AS country, 3.141 as price")
 
 	ctx, cancel := context.WithTimeout(testCtx(), 25*time.Second)
 	defer cancel()
 
-	_, err = server.GenerateMetricsViewFile(ctx, &runtimev1.GenerateMetricsViewFileRequest{
-		InstanceId: instanceID,
-		Table:      "ad_bids",
-		Path:       "/metrics/ad_bids_metrics_view.yaml",
-		UseAi:      false,
-	})
-	require.NoError(t, err)
-
 	repo, release, err := rt.Repo(ctx, instanceID)
 	require.NoError(t, err)
 	defer release()
 
-	data, err := repo.Get(ctx, "/metrics/ad_bids_metrics_view.yaml")
+	server, err := server.NewServer(ctx, &server.Options{}, rt, zap.NewNop(), ratelimit.NewNoop(), activity.NewNoopClient())
 	require.NoError(t, err)
 
-	require.Contains(t, data, "model: ad_bids")
-	require.Contains(t, data, "measures:")
-	require.Contains(t, data, "format_preset: humanize")
+	tt := []struct {
+		name      string
+		model     string
+		connector string
+		table     string
+		contains  []string
+	}{
+		{
+			name:  "model passed in request",
+			model: "ad_bids",
+			contains: []string{
+				"model: ad_bids",
+				"measures:",
+				"format_preset: humanize",
+			},
+		},
+		{
+			name:     "model passed in request that matches a table",
+			model:    "foo",
+			contains: []string{"model: foo"},
+		},
+		{
+			name:     "table passed in request that matches a model",
+			table:    "ad_bids",
+			contains: []string{"model: ad_bids"},
+		},
+		{
+			name:     "table passed in request that does not match a model",
+			table:    "foo",
+			contains: []string{"model: foo"},
+		},
+		{
+			name:      "table in non-default connector passed in request",
+			table:     "foo",
+			connector: "custom_duckdb",
+			contains:  []string{"connector: custom_duckdb", "model: foo"},
+		},
+	}
+
+	for _, tc := range tt {
+		t.Run(tc.name, func(t *testing.T) {
+			err = repo.Delete(ctx, "/metrics/generated_metrics_view.yaml", true)
+			require.NoError(t, err)
+
+			_, err = server.GenerateMetricsViewFile(ctx, &runtimev1.GenerateMetricsViewFileRequest{
+				InstanceId: instanceID,
+				Model:      tc.model,
+				Connector:  tc.connector,
+				Table:      tc.table,
+				Path:       "/metrics/generated_metrics_view.yaml",
+				UseAi:      false,
+			})
+			require.NoError(t, err)
+
+			data, err := repo.Get(ctx, "/metrics/generated_metrics_view.yaml")
+			require.NoError(t, err)
+
+			for _, c := range tc.contains {
+				require.Contains(t, data, c)
+			}
+		})
+	}
+}
+
+func olapExecAdhoc(t *testing.T, rt *runtime.Runtime, instanceID, connector, query string) {
+	h, release, err := rt.AcquireHandle(context.Background(), instanceID, connector)
+	require.NoError(t, err)
+	defer release()
+	olap, _ := h.AsOLAP(instanceID)
+	err = olap.Exec(context.Background(), &drivers.Statement{Query: query})
+	require.NoError(t, err)
 }