This commit is contained in:
duwenxin99
2026-02-11 16:11:14 -05:00
parent e0bacad708
commit de41946669

View File

@@ -954,6 +954,7 @@ func runBigQueryExecuteSqlToolInvokeTest(t *testing.T, select1Want, invokeParamW
api: "http://127.0.0.1:5000/api/tool/my-exec-sql-tool/invoke",
requestHeader: map[string]string{},
requestBody: bytes.NewBuffer([]byte(`{}`)),
want: `"error":"parameter \"sql\" is required"`,
isErr: true,
},
{
@@ -1215,7 +1216,7 @@ func runBigQueryWriteModeProtectedTest(t *testing.T, permanentDatasetName string
name: "CREATE TABLE to permanent dataset should fail",
toolName: "my-exec-sql-tool",
requestBody: fmt.Sprintf(`{"sql": "CREATE TABLE %s.new_table (x INT64)"}`, permanentDatasetName),
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: "protected write mode only supports SELECT statements, or write operations in the anonymous dataset",
wantResult: "",
},
@@ -2578,7 +2579,7 @@ func runListTableIdsWithRestriction(t *testing.T, allowedDatasetName, disallowed
{
name: "invoke on disallowed dataset",
dataset: disallowedDatasetName,
wantStatusCode: http.StatusBadRequest, // Or the specific error code returned
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("access denied to dataset '%s'", disallowedDatasetName),
},
}
@@ -2652,7 +2653,7 @@ func runGetDatasetInfoWithRestriction(t *testing.T, allowedDatasetName, disallow
{
name: "invoke on disallowed dataset",
dataset: disallowedDatasetName,
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("access denied to dataset '%s'", disallowedDatasetName),
},
}
@@ -2704,7 +2705,7 @@ func runGetTableInfoWithRestriction(t *testing.T, allowedDatasetName, disallowed
name: "invoke on disallowed table",
dataset: disallowedDatasetName,
table: disallowedTableName,
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("access denied to dataset '%s'", disallowedDatasetName),
},
}
@@ -2759,7 +2760,7 @@ func runExecuteSqlWithRestriction(t *testing.T, allowedTableFullName, disallowed
{
name: "invoke on disallowed table",
sql: fmt.Sprintf("SELECT * FROM %s", disallowedTableFullName),
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("query accesses dataset '%s', which is not in the allowed list",
strings.Join(
strings.Split(strings.Trim(disallowedTableFullName, "`"), ".")[0:2],
@@ -2768,31 +2769,31 @@ func runExecuteSqlWithRestriction(t *testing.T, allowedTableFullName, disallowed
{
name: "disallowed create schema",
sql: "CREATE SCHEMA another_dataset",
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: "dataset-level operations like 'CREATE_SCHEMA' are not allowed",
},
{
name: "disallowed alter schema",
sql: fmt.Sprintf("ALTER SCHEMA %s SET OPTIONS(description='new one')", allowedDatasetID),
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: "dataset-level operations like 'ALTER_SCHEMA' are not allowed",
},
{
name: "disallowed create function",
sql: fmt.Sprintf("CREATE FUNCTION %s.my_func() RETURNS INT64 AS (1)", allowedDatasetID),
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: "creating stored routines ('CREATE_FUNCTION') is not allowed",
},
{
name: "disallowed create procedure",
sql: fmt.Sprintf("CREATE PROCEDURE %s.my_proc() BEGIN SELECT 1; END", allowedDatasetID),
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: "unanalyzable statements like 'CREATE PROCEDURE' are not allowed",
},
{
name: "disallowed execute immediate",
sql: "EXECUTE IMMEDIATE 'SELECT 1'",
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: "EXECUTE IMMEDIATE is not allowed when dataset restrictions are in place",
},
}
@@ -2846,7 +2847,7 @@ func runConversationalAnalyticsWithRestriction(t *testing.T, allowedDatasetName,
{
name: "invoke with disallowed table",
tableRefs: disallowedTableRefsJSON,
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("access to dataset '%s.%s' (from table '%s') is not allowed", BigqueryProject, disallowedDatasetName, disallowedTableName),
},
}
@@ -3083,7 +3084,7 @@ func runForecastWithRestriction(t *testing.T, allowedTableFullName, disallowedTa
{
name: "invoke with disallowed table name",
historyData: disallowedTableUnquoted,
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("access to dataset '%s' (from table '%s') is not allowed", disallowedDatasetFQN, disallowedTableUnquoted),
},
{
@@ -3095,7 +3096,7 @@ func runForecastWithRestriction(t *testing.T, allowedTableFullName, disallowedTa
{
name: "invoke with query on disallowed table",
historyData: fmt.Sprintf("SELECT * FROM %s", disallowedTableFullName),
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("query in history_data accesses dataset '%s', which is not in the allowed list", disallowedDatasetFQN),
},
}
@@ -3174,7 +3175,7 @@ func runAnalyzeContributionWithRestriction(t *testing.T, allowedTableFullName, d
{
name: "invoke with disallowed table name",
inputData: disallowedTableUnquoted,
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("access to dataset '%s' (from table '%s') is not allowed", disallowedDatasetFQN, disallowedTableUnquoted),
},
{
@@ -3186,7 +3187,7 @@ func runAnalyzeContributionWithRestriction(t *testing.T, allowedTableFullName, d
{
name: "invoke with query on disallowed table",
inputData: fmt.Sprintf("SELECT * FROM %s", disallowedTableFullName),
wantStatusCode: http.StatusBadRequest,
wantStatusCode: http.StatusInternalServerError,
wantInError: fmt.Sprintf("query in input_data accesses dataset '%s', which is not in the allowed list", disallowedDatasetFQN),
},
}