rdtr commented on code in PR #59425:
URL: https://github.com/apache/spark/pull/59425#discussion_r4319134491
##########
sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala:
##########
@@ -1164,4 +1164,127 @@ 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 " +
Review Comment:
Thanks, added in d4ed446e83. It fails if the recursion is replaced with a
plain `transformRelation` 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]
