sarutak commented on code in PR #59425:
URL: https://github.com/apache/spark/pull/59425#discussion_r4256050203
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala:
##########
@@ -1164,4 +1164,125 @@ class SparkDeclarativePipelinesServerSuite
relation.getCommand.getSqlCommand.getInput)
}
}
+
+ // Helper to build a SqlCommand ExecutePlanRequest with
PipelineAnalysisContext.
+ private def buildSqlCommandRequestWithPipelineContext(
+ graphId: String,
+ input: proto.Relation): proto.ExecutePlanRequest = {
+ val context = proto.PipelineAnalysisContext
+ .newBuilder()
+ .setDataflowGraphId(graphId)
+ .build()
+ val userContext = proto.UserContext
+ .newBuilder()
+ .addExtensions(com.google.protobuf.Any.pack(context))
+ .setUserId("test_user")
+ .build()
+ val plan = proto.Plan
+ .newBuilder()
+ .setCommand(
+ proto.Command
+ .newBuilder()
+ .setSqlCommand(proto.SqlCommand.newBuilder().setInput(input).build())
+ .build())
+ .build()
+ proto.ExecutePlanRequest
+ .newBuilder()
+ .setUserContext(userContext)
+ .setPlan(plan)
+ .setSessionId(UUID.randomUUID().toString)
+ .build()
+ }
+
+ // Helper to build a SQL Relation.
+ private def sqlRelation(query: String): proto.Relation =
+ proto.Relation
+ .newBuilder()
+ .setSql(proto.SQL.newBuilder().setQuery(query))
+ .build()
+
+ // Helper to build a WITH_RELATIONS Relation with the given root SQL and
reference SQLs.
+ private def withRelationsRelation(
+ rootSql: String,
+ refs: Seq[(String, String)]): proto.Relation = {
+ val withRelationsBuilder =
proto.WithRelations.newBuilder().setRoot(sqlRelation(rootSql))
+ refs.foreach { case (alias, refSql) =>
+ withRelationsBuilder.addReferences(
+ proto.Relation
+ .newBuilder()
+ .setSubqueryAlias(
+ proto.SubqueryAlias
+ .newBuilder()
+ .setAlias(alias)
+ .setInput(sqlRelation(refSql))
+ .build())
+ .build())
+ }
+
proto.Relation.newBuilder().setWithRelations(withRelationsBuilder.build()).build()
+ }
+
+ test(
+ "SPARK-57189: blockUnsupportedSqlCommand catches Command at root of
WITH_RELATIONS " +
+ "without executing it") {
+ withRawBlockingStub { implicit stub =>
+ val graphId = createDataflowGraph
+ val tableName = "spark_catalog.default.test_with_relations_root_table"
+ try {
+ sql(s"DROP TABLE IF EXISTS $tableName")
+ val request = buildSqlCommandRequestWithPipelineContext(
+ graphId,
+ withRelationsRelation(
+ rootSql = s"CREATE TABLE $tableName (id INT) USING parquet",
+ refs = Seq("ref0" -> "SELECT 1 AS id")))
+ val ex = intercept[RuntimeException] {
+ stub.executePlan(request).next()
+ }
+
assert(ex.getMessage.contains("UNSUPPORTED_PIPELINE_SPARK_SQL_COMMAND"))
+ // Verify the guard caught the Command before any side effect could
materialize.
+ assert(!spark.catalog.tableExists(tableName))
+ } finally {
+ sql(s"DROP TABLE IF EXISTS $tableName")
+ }
+ }
+ }
+
+ test(
+ "SPARK-57189: blockUnsupportedSqlCommand catches Command inside reference
of " +
+ "WITH_RELATIONS without executing it") {
+ withRawBlockingStub { implicit stub =>
+ val graphId = createDataflowGraph
+ val tableName = "spark_catalog.default.test_with_relations_ref_table"
+ try {
+ sql(s"DROP TABLE IF EXISTS $tableName")
+ val request = buildSqlCommandRequestWithPipelineContext(
+ graphId,
+ withRelationsRelation(
+ rootSql = "SELECT * FROM ref0",
+ refs = Seq("ref0" -> s"CREATE TABLE $tableName (id INT) USING
parquet")))
+ val ex = intercept[RuntimeException] {
+ stub.executePlan(request).next()
+ }
+
assert(ex.getMessage.contains("UNSUPPORTED_PIPELINE_SPARK_SQL_COMMAND"))
+ // Verify the Command in the reference did not execute as a side
effect.
+ assert(!spark.catalog.tableExists(tableName))
+ } finally {
+ sql(s"DROP TABLE IF EXISTS $tableName")
+ }
+ }
+ }
+
+ test(
+ "SPARK-57189: blockUnsupportedSqlCommand allows valid SELECT
WITH_RELATIONS in pipeline " +
+ "analysis context") {
+ withRawBlockingStub { implicit stub =>
+ val graphId = createDataflowGraph
+ val request = buildSqlCommandRequestWithPipelineContext(
+ graphId,
+ withRelationsRelation(
+ rootSql = "SELECT * FROM ref0",
+ refs = Seq("ref0" -> "SELECT 1 AS id")))
+ val response = stub.executePlan(request).next()
+ assert(response.hasSqlCommandResult)
Review Comment:
This positive test passes both before and after the fix: reverting the
production change does not turn it red (the old guard did notblock a plain
SELECT either, so there was no over-block to begin with). It is therefore a
forward-looking guard that protects the new recursive
`blockUnsupportedSqlCommandIn` from a future regression where it would wrongly
block a valid SELECT WITH_RELATIONS. That is worth protecting, but it is worth
a one line comment (or a name that reflects "does not over-block a valid
SELECT") so a future reader does not mistake it for the test that proves the
double-execution/bypass fix. The two negative tests above are the fix-pinning
ones: without the fix, `transformRelation` on the WITH_RELATIONS executes the
CREATE TABLE and returns a `CommandResult` that is not on the blocklist, so
both `intercept[RuntimeException]` and `assert(!tableExists)` fail. Those are
non-vacuous. Only this positive case is a forward-looking non-regression check.
--
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.
To unsubscribe, e-mail: [email protected]
For queries about this service, please contact Infrastructure at:
[email protected]
---------------------------------------------------------------------
To unsubscribe, e-mail: [email protected]
For additional commands, e-mail: [email protected]
