chore: support query_data_policy (#150)

ecmadao · web-flow · commit 489218df3f78 · 2025-09-08T16:00:39.000+08:00
* chore: support query_data_policy

* fix: lint
diff --git a/VERSION b/VERSION
@@ -1 +1 @@
-3.9.3
+3.9.4
diff --git a/docs/data-sources/policy.md b/docs/data-sources/policy.md
@@ -26,6 +26,7 @@ The policy data source.
 - `disable_copy_data_policy` (Block List, Max: 1) Restrict data copying in SQL Editor (Admins/DBAs allowed) (see [below for nested schema](#nestedblock--disable_copy_data_policy))
 - `global_masking_policy` (Block List, Max: 1) (see [below for nested schema](#nestedblock--global_masking_policy))
 - `masking_exception_policy` (Block List, Max: 1) (see [below for nested schema](#nestedblock--masking_exception_policy))
+- `query_data_policy` (Block List, Max: 1) The policy for query data (see [below for nested schema](#nestedblock--query_data_policy))
 - `rollout_policy` (Block List, Max: 1) Control issue rollout. Learn more: https://docs.bytebase.com/administration/environment-policy/rollout-policy (see [below for nested schema](#nestedblock--rollout_policy))
 
 ### Read-Only
@@ -102,6 +103,20 @@ Optional:
 
 
 
+<a id="nestedblock--query_data_policy"></a>
+### Nested Schema for `query_data_policy`
+
+Required:
+
+- `disable_export` (Boolean) Disable export data in the SQL editor
+
+Optional:
+
+- `maximum_result_rows` (Number) The return rows limit. If the value <= 0, will be treated as no limit. The default value is -1.
+- `maximum_result_size` (Number) The size limit in bytes. The default value is 100MB, we will use the default value if the limit <= 0.
+- `timeout_in_seconds` (Number) The maximum time allowed for a query to run in SQL Editor. No limit when the value <= 0
+
+
 <a id="nestedblock--rollout_policy"></a>
 ### Nested Schema for `rollout_policy`
 
diff --git a/docs/data-sources/policy_list.md b/docs/data-sources/policy_list.md
@@ -36,6 +36,7 @@ Read-Only:
 - `inherit_from_parent` (Boolean)
 - `masking_exception_policy` (List of Object) (see [below for nested schema](#nestedobjatt--policies--masking_exception_policy))
 - `name` (String)
+- `query_data_policy` (List of Object) (see [below for nested schema](#nestedobjatt--policies--query_data_policy))
 - `rollout_policy` (List of Object) (see [below for nested schema](#nestedobjatt--policies--rollout_policy))
 - `type` (String)
 
@@ -100,6 +101,17 @@ Read-Only:
 
 
 
+<a id="nestedobjatt--policies--query_data_policy"></a>
+### Nested Schema for `policies.query_data_policy`
+
+Read-Only:
+
+- `disable_export` (Boolean)
+- `maximum_result_rows` (Number)
+- `maximum_result_size` (Number)
+- `timeout_in_seconds` (Number)
+
+
 <a id="nestedobjatt--policies--rollout_policy"></a>
 ### Nested Schema for `policies.rollout_policy`
 
diff --git a/docs/data-sources/setting.md b/docs/data-sources/setting.md
@@ -24,7 +24,6 @@ The setting data source.
 - `classification` (Block List, Max: 1) Classification for data masking. Require ENTERPRISE subscription. (see [below for nested schema](#nestedblock--classification))
 - `password_restriction` (Block List, Max: 1) Restrict for login password (see [below for nested schema](#nestedblock--password_restriction))
 - `semantic_types` (Block Set) Semantic types for data masking. Require ENTERPRISE subscription. (see [below for nested schema](#nestedblock--semantic_types))
-- `sql_query_restriction` (Block List, Max: 1) Restrict for SQL query result (see [below for nested schema](#nestedblock--sql_query_restriction))
 - `workspace_profile` (Block List, Max: 1) (see [below for nested schema](#nestedblock--workspace_profile))
 
 ### Read-Only
@@ -159,15 +158,6 @@ Required:
 
 
 
-<a id="nestedblock--sql_query_restriction"></a>
-### Nested Schema for `sql_query_restriction`
-
-Optional:
-
-- `maximum_result_rows` (Number) The return rows limit. If the value <= 0, will be treated as no limit. The default value is -1.
-- `maximum_result_size` (Number) The size limit in bytes. The default value is 100MB, we will use the default value if the setting not exists, or the limit <= 0.
-
-
 <a id="nestedblock--workspace_profile"></a>
 ### Nested Schema for `workspace_profile`
 
diff --git a/docs/resources/policy.md b/docs/resources/policy.md
@@ -28,6 +28,7 @@ The policy resource.
 - `global_masking_policy` (Block List, Max: 1) (see [below for nested schema](#nestedblock--global_masking_policy))
 - `inherit_from_parent` (Boolean) Decide if the policy should inherit from the parent.
 - `masking_exception_policy` (Block List, Max: 1) (see [below for nested schema](#nestedblock--masking_exception_policy))
+- `query_data_policy` (Block List, Max: 1) The policy for query data (see [below for nested schema](#nestedblock--query_data_policy))
 - `rollout_policy` (Block List, Max: 1) Control issue rollout. Learn more: https://docs.bytebase.com/administration/environment-policy/rollout-policy (see [below for nested schema](#nestedblock--rollout_policy))
 
 ### Read-Only
@@ -102,6 +103,20 @@ Optional:
 
 
 
+<a id="nestedblock--query_data_policy"></a>
+### Nested Schema for `query_data_policy`
+
+Required:
+
+- `disable_export` (Boolean) Disable export data in the SQL editor
+
+Optional:
+
+- `maximum_result_rows` (Number) The return rows limit. If the value <= 0, will be treated as no limit. The default value is -1.
+- `maximum_result_size` (Number) The size limit in bytes. The default value is 100MB, we will use the default value if the limit <= 0.
+- `timeout_in_seconds` (Number) The maximum time allowed for a query to run in SQL Editor. No limit when the value <= 0
+
+
 <a id="nestedblock--rollout_policy"></a>
 ### Nested Schema for `rollout_policy`
 
diff --git a/docs/resources/setting.md b/docs/resources/setting.md
@@ -26,7 +26,6 @@ The setting resource.
 - `environment_setting` (Block List) The environment (see [below for nested schema](#nestedblock--environment_setting))
 - `password_restriction` (Block List, Max: 1) Restrict for login password (see [below for nested schema](#nestedblock--password_restriction))
 - `semantic_types` (Block Set) Semantic types for data masking. Require ENTERPRISE subscription. (see [below for nested schema](#nestedblock--semantic_types))
-- `sql_query_restriction` (Block List, Max: 1) Restrict for SQL query result (see [below for nested schema](#nestedblock--sql_query_restriction))
 - `workspace_profile` (Block List, Max: 1) (see [below for nested schema](#nestedblock--workspace_profile))
 
 ### Read-Only
@@ -235,15 +234,6 @@ Required:
 
 
 
-<a id="nestedblock--sql_query_restriction"></a>
-### Nested Schema for `sql_query_restriction`
-
-Optional:
-
-- `maximum_result_rows` (Number) The return rows limit. If the value <= 0, will be treated as no limit. The default value is -1.
-- `maximum_result_size` (Number) The size limit in bytes. The default value is 100MB, we will use the default value if the setting not exists, or the limit <= 0.
-
-
 <a id="nestedblock--workspace_profile"></a>
 ### Nested Schema for `workspace_profile`
 
diff --git a/examples/setup/main.tf b/examples/setup/main.tf
@@ -1,7 +1,7 @@
 terraform {
   required_providers {
     bytebase = {
-      version = "3.9.2"
+      version = "3.9.4"
       # For local development, please use "terraform.local/bytebase/bytebase" instead
       source = "registry.terraform.io/bytebase/bytebase"
     }
@@ -18,7 +18,6 @@ locals {
   project_id          = "project-sample"
 }
 
-
 provider "bytebase" {
   # You need to replace the account and key with your Bytebase service account.
   service_account = local.service_account
@@ -36,3 +35,14 @@ resource "bytebase_setting" "workspace_profile" {
     domains      = ["bytebase.com"]
   }
 }
+
+resource "bytebase_policy" "query_data_policy" {
+  parent = "workspaces/-"
+  type   = "DATA_QUERY"
+  query_data_policy {
+    maximum_result_size = 200 * 1024 * 1024 # 200MB
+    maximum_result_rows = 100
+    disable_export      = false
+    timeout_in_seconds  = 60 # 60 seconds
+  }
+}
diff --git a/go.mod b/go.mod
@@ -5,8 +5,8 @@ go 1.24.4
 toolchain go1.24.5
 
 require (
-	buf.build/gen/go/bytebase/bytebase/connectrpc/go v1.18.1-20250821091751-ab434d709c89.1
-	buf.build/gen/go/bytebase/bytebase/protocolbuffers/go v1.36.8-20250821091751-ab434d709c89.1
+	buf.build/gen/go/bytebase/bytebase/connectrpc/go v1.18.1-20250908030532-58bfc338601e.1
+	buf.build/gen/go/bytebase/bytebase/protocolbuffers/go v1.36.8-20250908030532-58bfc338601e.1
 	connectrpc.com/connect v1.18.1
 	github.com/hashicorp/go-cty v1.5.0
 	github.com/hashicorp/terraform-plugin-docs v0.13.0
diff --git a/go.sum b/go.sum
@@ -1,7 +1,7 @@
-buf.build/gen/go/bytebase/bytebase/connectrpc/go v1.18.1-20250821091751-ab434d709c89.1 h1:ZmVAFhx6C+5hh2JwHIicGVXPU0P+dWWUMYsSzXJGSZM=
-buf.build/gen/go/bytebase/bytebase/connectrpc/go v1.18.1-20250821091751-ab434d709c89.1/go.mod h1:cCayz9URf5ApvJZGfIE7bfasYECFG9ECIVNNOm4Vci0=
-buf.build/gen/go/bytebase/bytebase/protocolbuffers/go v1.36.8-20250821091751-ab434d709c89.1 h1:uHBwD119t1nqozj6q/Zh/+NBfzUb/qLVRK/dhvZkbqw=
-buf.build/gen/go/bytebase/bytebase/protocolbuffers/go v1.36.8-20250821091751-ab434d709c89.1/go.mod h1:dwdKUX0jGgJ7OJe024SNHvANb1TKuBzIrZOzL/3Njtk=
+buf.build/gen/go/bytebase/bytebase/connectrpc/go v1.18.1-20250908030532-58bfc338601e.1 h1:FQEei3tpJ+OjhAAfLOTzqm+XaFboZI+UpWbXEiQJ1m4=
+buf.build/gen/go/bytebase/bytebase/connectrpc/go v1.18.1-20250908030532-58bfc338601e.1/go.mod h1:/jFGXncqA/NGlTsADOmP7o1PaKo05HUP64uWe4g7UZg=
+buf.build/gen/go/bytebase/bytebase/protocolbuffers/go v1.36.8-20250908030532-58bfc338601e.1 h1:9VCImLW8Zda6Tay+4zoS85Y8RwX94/x5XlZeeyQHZ38=
+buf.build/gen/go/bytebase/bytebase/protocolbuffers/go v1.36.8-20250908030532-58bfc338601e.1/go.mod h1:dwdKUX0jGgJ7OJe024SNHvANb1TKuBzIrZOzL/3Njtk=
 connectrpc.com/connect v1.18.1 h1:PAg7CjSAGvscaf6YZKUefjoih5Z/qYkyaTrBW8xvYPw=
 connectrpc.com/connect v1.18.1/go.mod h1:0292hj1rnx8oFrStN7cB4jjVBeqs+Yx5yDIC2prWDO8=
 dario.cat/mergo v1.0.0 h1:AGCNq9Evsj31mOgNPcLyXc+4PNABt905YmuqPYYpBWk=
diff --git a/provider/data_source_policy.go b/provider/data_source_policy.go
@@ -47,6 +47,7 @@ func dataSourcePolicy() *schema.Resource {
 					v1pb.PolicyType_DISABLE_COPY_DATA.String(),
 					v1pb.PolicyType_DATA_SOURCE_QUERY.String(),
 					v1pb.PolicyType_ROLLOUT_POLICY.String(),
+					v1pb.PolicyType_DATA_QUERY.String(),
 				}, false),
 				Description: "The policy type.",
 			},
@@ -70,6 +71,7 @@ func dataSourcePolicy() *schema.Resource {
 			"disable_copy_data_policy": getDisableCopyDataPolicySchema(true),
 			"data_source_query_policy": getDataSourceQueryPolicySchema(true),
 			"rollout_policy":           getRolloutPolicySchema(true),
+			"query_data_policy":        getDataQueryPolicySchema(true),
 		},
 	}
 }
@@ -223,6 +225,44 @@ func getGlobalMaskingPolicySchema(computed bool) *schema.Schema {
 	}
 }
 
+func getDataQueryPolicySchema(computed bool) *schema.Schema {
+	return &schema.Schema{
+		Computed:    computed,
+		Optional:    true,
+		Default:     nil,
+		Type:        schema.TypeList,
+		MaxItems:    1,
+		MinItems:    1,
+		Description: "The policy for query data",
+		Elem: &schema.Resource{
+			Schema: map[string]*schema.Schema{
+				"maximum_result_size": {
+					Type:        schema.TypeInt,
+					Optional:    true,
+					Default:     100 * 1024 * 1024,
+					Description: "The size limit in bytes. The default value is 100MB, we will use the default value if the limit <= 0.",
+				},
+				"maximum_result_rows": {
+					Type:        schema.TypeInt,
+					Optional:    true,
+					Default:     -1,
+					Description: "The return rows limit. If the value <= 0, will be treated as no limit. The default value is -1.",
+				},
+				"disable_export": {
+					Type:        schema.TypeBool,
+					Required:    true,
+					Description: "Disable export data in the SQL editor",
+				},
+				"timeout_in_seconds": {
+					Type:        schema.TypeInt,
+					Optional:    true,
+					Description: "The maximum time allowed for a query to run in SQL Editor. No limit when the value <= 0",
+				},
+			},
+		},
+	}
+}
+
 func getDisableCopyDataPolicySchema(computed bool) *schema.Schema {
 	return &schema.Schema{
 		Computed:    computed,
@@ -397,6 +437,11 @@ func flattenPolicyPayload(policy *v1pb.Policy) (string, interface{}, diag.Diagno
 			rolloutPolicy := flattenRolloutPolicy(p)
 			return "rollout_policy", rolloutPolicy, nil
 		}
+	case v1pb.PolicyType_DATA_QUERY:
+		if p := policy.GetQueryDataPolicy(); p != nil {
+			rolloutPolicy := flattenQueryDataPolicy(p)
+			return "query_data_policy", rolloutPolicy, nil
+		}
 	}
 
 	return "", nil, diag.Errorf("unsupported policy: %s", policy.Name)
@@ -429,6 +474,16 @@ func flattenDisableCopyDataPolicy(p *v1pb.DisableCopyDataPolicy) []interface{} {
 	return []interface{}{policy}
 }
 
+func flattenQueryDataPolicy(p *v1pb.QueryDataPolicy) []interface{} {
+	policy := map[string]interface{}{
+		"maximum_result_size": int(p.MaximumResultSize),
+		"maximum_result_rows": int(p.MaximumResultRows),
+		"disable_export":      p.DisableExport,
+		"timeout_in_seconds":  int(p.Timeout.Seconds),
+	}
+	return []interface{}{policy}
+}
+
 func flattenGlobalMaskingPolicy(p *v1pb.MaskingRulePolicy) ([]interface{}, error) {
 	ruleList := []interface{}{}
 
diff --git a/provider/data_source_policy_list.go b/provider/data_source_policy_list.go
@@ -23,7 +23,7 @@ func dataSourcePolicyList() *schema.Resource {
 			"parent": {
 				Type:     schema.TypeString,
 				Optional: true,
-				Default:  "",
+				Default:  internal.WorkspaceName,
 				ValidateDiagFunc: internal.ResourceNameValidation(
 					// workspace policy
 					fmt.Sprintf("^%s$", internal.WorkspaceName),
@@ -68,6 +68,7 @@ func dataSourcePolicyList() *schema.Resource {
 						"disable_copy_data_policy": getDisableCopyDataPolicySchema(true),
 						"data_source_query_policy": getDataSourceQueryPolicySchema(true),
 						"rollout_policy":           getRolloutPolicySchema(true),
+						"query_data_policy":        getDataQueryPolicySchema(true),
 					},
 				},
 			},
diff --git a/provider/data_source_setting.go b/provider/data_source_setting.go
@@ -32,16 +32,14 @@ func dataSourceSetting() *schema.Resource {
 					fmt.Sprintf("^%s%s$", internal.SettingNamePrefix, v1pb.Setting_SEMANTIC_TYPES.String()),
 					fmt.Sprintf("^%s%s$", internal.SettingNamePrefix, v1pb.Setting_ENVIRONMENT.String()),
 					fmt.Sprintf("^%s%s$", internal.SettingNamePrefix, v1pb.Setting_PASSWORD_RESTRICTION.String()),
-					fmt.Sprintf("^%s%s$", internal.SettingNamePrefix, v1pb.Setting_SQL_RESULT_SIZE_LIMIT.String()),
 				),
 			},
-			"approval_flow":         getWorkspaceApprovalSetting(true),
-			"workspace_profile":     getWorkspaceProfileSetting(true),
-			"classification":        getClassificationSetting(true),
-			"semantic_types":        getSemanticTypesSetting(true),
-			"environment_setting":   getEnvironmentSetting(true),
-			"password_restriction":  getPasswordRestrictionSetting(true),
-			"sql_query_restriction": getSQLQueryRestrictionSetting(true),
+			"approval_flow":        getWorkspaceApprovalSetting(true),
+			"workspace_profile":    getWorkspaceProfileSetting(true),
+			"classification":       getClassificationSetting(true),
+			"semantic_types":       getSemanticTypesSetting(true),
+			"environment_setting":  getEnvironmentSetting(true),
+			"password_restriction": getPasswordRestrictionSetting(true),
 		},
 	}
 }
@@ -378,34 +376,6 @@ func getWorkspaceProfileSetting(computed bool) *schema.Schema {
 	}
 }
 
-func getSQLQueryRestrictionSetting(computed bool) *schema.Schema {
-	return &schema.Schema{
-		Computed:    computed,
-		Optional:    true,
-		Default:     nil,
-		Type:        schema.TypeList,
-		MaxItems:    1,
-		MinItems:    1,
-		Description: "Restrict for SQL query result",
-		Elem: &schema.Resource{
-			Schema: map[string]*schema.Schema{
-				"maximum_result_size": {
-					Type:        schema.TypeInt,
-					Optional:    true,
-					Default:     100 * 1024 * 1024,
-					Description: "The size limit in bytes. The default value is 100MB, we will use the default value if the setting not exists, or the limit <= 0.",
-				},
-				"maximum_result_rows": {
-					Type:        schema.TypeInt,
-					Optional:    true,
-					Default:     -1,
-					Description: "The return rows limit. If the value <= 0, will be treated as no limit. The default value is -1.",
-				},
-			},
-		},
-	}
-}
-
 const minimumPasswordLength = 8
 
 func getPasswordRestrictionSetting(computed bool) *schema.Schema {
@@ -640,12 +610,6 @@ func setSettingMessage(ctx context.Context, d *schema.ResourceData, client api.C
 			return diag.Errorf("cannot set password_restriction: %s", err.Error())
 		}
 	}
-	if value := setting.GetValue().GetSqlQueryRestrictionSetting(); value != nil {
-		settingVal := flattenSQLQueryRestrictionSetting(value)
-		if err := d.Set("sql_query_restriction", settingVal); err != nil {
-			return diag.Errorf("cannot set sql_query_restriction: %s", err.Error())
-		}
-	}
 	if value := setting.GetValue().GetDataClassificationSettingValue(); value != nil {
 		settingVal := flattenClassificationSetting(value)
 		if err := d.Set("classification", settingVal); err != nil {
@@ -845,13 +809,6 @@ func flattenPasswordRestrictionSetting(setting *v1pb.PasswordRestrictionSetting)
 	return []interface{}{raw}
 }
 
-func flattenSQLQueryRestrictionSetting(setting *v1pb.SQLQueryRestrictionSetting) []interface{} {
-	raw := map[string]interface{}{}
-	raw["maximum_result_size"] = int(setting.MaximumResultSize)
-	raw["maximum_result_rows"] = int(setting.MaximumResultRows)
-	return []interface{}{raw}
-}
-
 func flattenClassificationSetting(setting *v1pb.DataClassificationSetting) []interface{} {
 	raw := map[string]interface{}{}
 
diff --git a/provider/internal/mock_client.go b/provider/internal/mock_client.go
diff --git a/provider/resource_policy.go b/provider/resource_policy.go
diff --git a/provider/resource_policy_test.go b/provider/resource_policy_test.go
diff --git a/provider/resource_setting.go b/provider/resource_setting.go
diff --git a/provider/resource_setting_test.go b/provider/resource_setting_test.go
