Added databricks_model_serving support to databricks_permissions (#2331)

* Add model serving support to `databricks_permissions`

To support permissions assignment we also need to expose a new attribute with internal ID
of the model serving endpoint.

* Update docs/resources/model_serving.md

Co-authored-by: Arpit Jasapara <[email protected]>

* Update docs/resources/model_serving.md

Co-authored-by: Arpit Jasapara <[email protected]>

---------

Co-authored-by: Arpit Jasapara <[email protected]>

Author: alexott

docs/resources/model_serving.md

@@ -68,6 +68,17 @@ The following arguments are supported:
 * `served_model_name` - (Required) The name of the served model this route configures traffic for. This needs to match the name of a `served_models` block
 * `traffic_percentage` - (Required) The percentage of endpoint traffic to send to this route. It must be an integer between 0 and 100 inclusive.
 
+## Attribute Reference
+
+In addition to all arguments above, the following attributes are exported:
+
+* `id` - Equal to the `name` argument and used to identify the serving endpoint.
+* `serving_endpoint_id` - Unique identifier of the serving endpoint primarily used to set permissions and refer to this instance for other operations.
+
+## Access Control
+
+* [databricks_permissions](permissions.md#model-serving-usage) can control which groups or individual users can *Manage*, *Query* or *View* individual serving endpoints.
+
 ## Timeouts
 
 The `timeouts` block allows you to specify `create` and `update` timeouts. The default right now is 45 minutes for both operations.

docs/resources/permissions.md

@@ -517,6 +517,52 @@ resource "databricks_permissions" "model_usage" {
 }
 ```
 
+## Model serving usage
+
+Valid permission levels for [databricks_model_serving](model_serving.md) are: `CAN_VIEW`, `CAN_QUERY`, and `CAN_MANAGE`.
+
+```hcl
+resource "databricks_model_serving" "this" {
+  name = "tf-test"
+  config {
+    served_models {
+      name                  = "prod_model"
+      model_name            = "test"
+      model_version         = "1"
+      workload_size         = "Small"
+      scale_to_zero_enabled = true
+    }
+  }
+}
+
+resource "databricks_group" "auto" {
+  display_name = "Automation"
+}
+
+resource "databricks_group" "eng" {
+  display_name = "Engineering"
+}
+
+resource "databricks_permissions" "ml_serving_usage" {
+  serving_endpoint_id = databricks_model_serving.this.serving_endpoint_id
+
+  access_control {
+    group_name       = "users"
+    permission_level = "CAN_VIEW"
+  }
+
+  access_control {
+    group_name       = databricks_group.auto.display_name
+    permission_level = "CAN_MANAGE"
+  }
+
+  access_control {
+    group_name       = databricks_group.eng.display_name
+    permission_level = "CAN_QUERY"
+  }
+}
+```
+
 ## Passwords usage
 
 By default on AWS deployments, all admin users can sign in to Databricks using either SSO or their username and password, and all API users can authenticate to the Databricks REST APIs using their username and password. As an admin, you [can limit](https://docs.databricks.com/administration-guide/users-groups/single-sign-on/index.html#optional-configure-password-access-control) admin users’ and API users’ ability to authenticate with their username and password by configuring `CAN_USE` permissions using password access control.
@@ -734,6 +780,7 @@ Exactly one of the following arguments is required:
 - `repo_path` - path of databricks repo directory(`/Repos/<username>/...`)
 - `experiment_id` - [MLflow experiment](mlflow_experiment.md) id
 - `registered_model_id` - [MLflow registered model](mlflow_model.md) id
+- `serving_endpoint_id` - [Model Serving](model_serving.md) endpoint id.
 - `authorization` - either [`tokens`](https://docs.databricks.com/administration-guide/access-control/tokens.html) or [`passwords`](https://docs.databricks.com/administration-guide/users-groups/single-sign-on/index.html#configure-password-permission).
 - `sql_endpoint_id` - [SQL warehouse](sql_endpoint.md) id
 - `sql_dashboard_id` - [SQL dashboard](sql_dashboard.md) id

internal/acceptance/model_serving_test.go

@@ -122,6 +122,15 @@ func TestAccModelServing(t *testing.T) {
 					}
 				}
 			}
+
+			resource "databricks_permissions" "ml_serving_usage" {
+				serving_endpoint_id = databricks_model_serving.endpoint.serving_endpoint_id
+			  
+				access_control {
+				  group_name       = "users"
+				  permission_level = "CAN_VIEW"
+				}
+			}
 		`, name),
 		},
 		step{

permissions/resource_permissions.go

@@ -124,7 +124,7 @@ func urlPathForObjectID(objectID string) string {
 // permissions when POSTing permissions changes through the REST API, to avoid accidentally
 // revoking the calling user's ability to manage the current object.
 func (a PermissionsAPI) shouldExplicitlyGrantCallingUserManagePermissions(objectID string) bool {
-	for _, prefix := range [...]string{"/registered-models/", "/clusters/", "/instance-pools/", "/queries/", "/sql/warehouses"} {
+	for _, prefix := range [...]string{"/registered-models/", "/clusters/", "/instance-pools/", "/serving-endpoints/", "/queries/", "/sql/warehouses"} {
 		if strings.HasPrefix(objectID, prefix) {
 			return true
 		}
@@ -309,6 +309,7 @@ func permissionsResourceIDFields() []permissionsIDFieldMapping {
 		{"experiment_id", "mlflowExperiment", "experiments", []string{"CAN_READ", "CAN_EDIT", "CAN_MANAGE"}, SIMPLE},
 		{"registered_model_id", "registered-model", "registered-models", []string{
 			"CAN_READ", "CAN_EDIT", "CAN_MANAGE_STAGING_VERSIONS", "CAN_MANAGE_PRODUCTION_VERSIONS", "CAN_MANAGE"}, SIMPLE},
+		{"serving_endpoint_id", "serving-endpoint", "serving-endpoints", []string{"CAN_VIEW", "CAN_QUERY", "CAN_MANAGE"}, SIMPLE},
 	}
 }
 

serving/resource_model_serving.go

@@ -21,6 +21,14 @@ func ResourceModelServing() *schema.Resource {
 			common.MustSchemaPath(m, "config", "served_models", "scale_to_zero_enabled").Optional = true
 			common.MustSchemaPath(m, "config", "served_models", "scale_to_zero_enabled").Default = true
 			common.MustSchemaPath(m, "config", "served_models", "name").Computed = true
+
+			common.MustSchemaPath(m, "config", "traffic_config").Computed = true
+
+			m["serving_endpoint_id"] = &schema.Schema{
+				Computed: true,
+				Type:     schema.TypeString,
+			}
+
 			return m
 		})
 
@@ -48,7 +56,12 @@ func ResourceModelServing() *schema.Resource {
 			if err != nil {
 				return err
 			}
-			return common.StructToData(*endpoint, s, d)
+			err = common.StructToData(*endpoint, s, d)
+			if err != nil {
+				return err
+			}
+			d.Set("serving_endpoint_id", endpoint.Id)
+			return nil
 		},
 		Update: func(ctx context.Context, d *schema.ResourceData, c *common.DatabricksClient) error {
 			w, err := c.WorkspaceClient()