[SPARK-53593][SDP] Add response field for DefineDataset and DefineFlow RPC

sql/connect/common/src/main/protobuf/spark/connect/pipelines.proto

@@ -146,11 +146,19 @@ message PipelineCommand {
 message PipelineCommandResult {
   oneof result_type {
     CreateDataflowGraphResult create_dataflow_graph_result = 1;
+    DefineDatasetResult define_dataset_result = 2;
+    DefineFlowResult define_flow_result = 3;
   }
   message CreateDataflowGraphResult {
     // The ID of the created graph.
     optional string dataflow_graph_id = 1;
   }
+  message DefineDatasetResult {
+    optional string resolved_data_name = 1;
+  }
+  message DefineFlowResult {
+    optional string resolved_flow_name = 1;
+  }
 }
 
 // The type of dataset.

sql/connect/server/src/main/scala/org/apache/spark/sql/connect/pipelines/PipelinesHandler.scala

@@ -81,12 +81,25 @@ private[connect] object PipelinesHandler extends Logging {
         defaultResponse
       case proto.PipelineCommand.CommandTypeCase.DEFINE_DATASET =>
         logInfo(s"Define pipelines dataset cmd received: $cmd")
-        defineDataset(cmd.getDefineDataset, sessionHolder)
-        defaultResponse
+        val qualifiedIdentifier = defineDataset(cmd.getDefineDataset, sessionHolder)
+        PipelineCommandResult.newBuilder()
+          .setDefineDatasetResult(
+            PipelineCommandResult.DefineDatasetResult.newBuilder()
+            .setResolvedDataName(qualifiedIdentifier.quotedString)
+            .build()
+          )
+        .build()
       case proto.PipelineCommand.CommandTypeCase.DEFINE_FLOW =>
         logInfo(s"Define pipelines flow cmd received: $cmd")
-        defineFlow(cmd.getDefineFlow, transformRelationFunc, sessionHolder)
-        defaultResponse
+        val qualifiedIdentifier =
+          defineFlow(cmd.getDefineFlow, transformRelationFunc, sessionHolder)
+        PipelineCommandResult.newBuilder()
+          .setDefineFlowResult(
+            PipelineCommandResult.DefineFlowResult.newBuilder()
+            .setResolvedFlowName(qualifiedIdentifier.quotedString)
+            .build()
+          )
+        .build()
       case proto.PipelineCommand.CommandTypeCase.START_RUN =>
         logInfo(s"Start pipeline cmd received: $cmd")
         startRun(cmd.getStartRun, responseObserver, sessionHolder)
@@ -140,20 +153,22 @@ private[connect] object PipelinesHandler extends Logging {
 
   private def defineDataset(
       dataset: proto.PipelineCommand.DefineDataset,
-      sessionHolder: SessionHolder): Unit = {
+      sessionHolder: SessionHolder): TableIdentifier = {
     val dataflowGraphId = dataset.getDataflowGraphId
     val graphElementRegistry =
       sessionHolder.dataflowGraphRegistry.getDataflowGraphOrThrow(dataflowGraphId)
 
     dataset.getDatasetType match {
       case proto.DatasetType.MATERIALIZED_VIEW | proto.DatasetType.TABLE =>
-        val tableIdentifier =
-          GraphIdentifierManager.parseTableIdentifier(
-            dataset.getDatasetName,
-            sessionHolder.session)
+        val qualifiedIdentifier = GraphIdentifierManager.parseAndQualifyTableIdentifier(
+          rawTableIdentifier = GraphIdentifierManager.parseTableIdentifier
+          (dataset.getDatasetName, sessionHolder.session),
+          currentCatalog = Some(graphElementRegistry.defaultCatalog),
+          currentDatabase = Some(graphElementRegistry.defaultDatabase)
+        ).identifier
         graphElementRegistry.registerTable(
           Table(
-            identifier = tableIdentifier,
+            identifier = qualifiedIdentifier,
             comment = Option(dataset.getComment),
             specifiedSchema = Option.when(dataset.hasSchema)(
               DataTypeProtoConverter
@@ -164,26 +179,30 @@ private[connect] object PipelinesHandler extends Logging {
             properties = dataset.getTablePropertiesMap.asScala.toMap,
             baseOrigin = QueryOrigin(
               objectType = Option(QueryOriginType.Table.toString),
-              objectName = Option(tableIdentifier.unquotedString),
+              objectName = Option(qualifiedIdentifier.quotedString),
               language = Option(Python())),
             format = Option.when(dataset.hasFormat)(dataset.getFormat),
             normalizedPath = None,
             isStreamingTable = dataset.getDatasetType == proto.DatasetType.TABLE))
+        qualifiedIdentifier
       case proto.DatasetType.TEMPORARY_VIEW =>
-        val viewIdentifier =
-          GraphIdentifierManager.parseTableIdentifier(
-            dataset.getDatasetName,
-            sessionHolder.session)
-
+        val viewIdentifier = GraphIdentifierManager
+          .parseAndValidateTemporaryViewIdentifier(
+            rawViewIdentifier =
+              GraphIdentifierManager.parseTableIdentifier(
+                dataset.getDatasetName,
+                sessionHolder.session)
+          )
         graphElementRegistry.registerView(
           TemporaryView(
             identifier = viewIdentifier,
             comment = Option(dataset.getComment),
             origin = QueryOrigin(
               objectType = Option(QueryOriginType.View.toString),
-              objectName = Option(viewIdentifier.unquotedString),
+              objectName = Option(viewIdentifier.quotedString),
               language = Option(Python())),
             properties = Map.empty))
+        viewIdentifier
       case _ =>
         throw new IllegalArgumentException(s"Unknown dataset type: ${dataset.getDatasetType}")
     }
@@ -192,40 +211,65 @@ private[connect] object PipelinesHandler extends Logging {
   private def defineFlow(
       flow: proto.PipelineCommand.DefineFlow,
       transformRelationFunc: Relation => LogicalPlan,
-      sessionHolder: SessionHolder): Unit = {
+      sessionHolder: SessionHolder): TableIdentifier = {
     val dataflowGraphId = flow.getDataflowGraphId
     val graphElementRegistry =
       sessionHolder.dataflowGraphRegistry.getDataflowGraphOrThrow(dataflowGraphId)
+    val defaultCatalog = graphElementRegistry.defaultCatalog
+    val defaultDatabase = graphElementRegistry.defaultDatabase
 
     val isImplicitFlow = flow.getFlowName == flow.getTargetDatasetName
-
-    val flowIdentifier = GraphIdentifierManager
+    val rawFlowIdentifier = GraphIdentifierManager
       .parseTableIdentifier(name = flow.getFlowName, spark = sessionHolder.session)
 
     // If the flow is not an implicit flow (i.e. one defined as part of dataset creation), then
     // it must be a single-part identifier.
-    if (!isImplicitFlow && !IdentifierHelper.isSinglePartIdentifier(flowIdentifier)) {
+    if (!isImplicitFlow && !IdentifierHelper.isSinglePartIdentifier(rawFlowIdentifier)) {
       throw new AnalysisException(
         "MULTIPART_FLOW_NAME_NOT_SUPPORTED",
         Map("flowName" -> flow.getFlowName))
     }
 
+    val rawDestinationIdentifier = GraphIdentifierManager
+      .parseTableIdentifier(name = flow.getTargetDatasetName, spark = sessionHolder.session)
+    val isImplicitFlowForTempView =
+      graphElementRegistry.getViews()
+      .filter(_.isInstanceOf[TemporaryView])
+      .exists(_.identifier == rawDestinationIdentifier)
+
+    val Seq(flowIdentifier, destinationIdentifier) = Seq(
+      rawFlowIdentifier, rawDestinationIdentifier
+    ).map { rawIdentifier =>
+      // If the flow is created implicitly as part of defining a view, then we do not
+      // qualify the flow identifier and the flow destination. This is because views are
+      // not permitted to have multipart
+      if (isImplicitFlow && isImplicitFlowForTempView) {
+        rawIdentifier
+      } else {
+        GraphIdentifierManager.parseAndQualifyFlowIdentifier(
+          rawFlowIdentifier = rawIdentifier,
+          currentCatalog = Some(defaultCatalog),
+          currentDatabase = Some(defaultDatabase)
+        ).identifier
+      }
+    }
+
     graphElementRegistry.registerFlow(
       new UnresolvedFlow(
         identifier = flowIdentifier,
-        destinationIdentifier = GraphIdentifierManager
-          .parseTableIdentifier(name = flow.getTargetDatasetName, spark = sessionHolder.session),
+        destinationIdentifier = destinationIdentifier,
         func =
           FlowAnalysis.createFlowFunctionFromLogicalPlan(transformRelationFunc(flow.getRelation)),
         sqlConf = flow.getSqlConfMap.asScala.toMap,
         once = false,
         queryContext = QueryContext(
-          Option(graphElementRegistry.defaultCatalog),
-          Option(graphElementRegistry.defaultDatabase)),
+          Option(defaultCatalog),
+          Option(defaultDatabase)),
         origin = QueryOrigin(
           objectType = Option(QueryOriginType.Flow.toString),
-          objectName = Option(flowIdentifier.unquotedString),
+          objectName = Option(flowIdentifier.quotedString),
           language = Option(Python()))))
+    flowIdentifier
   }
 
   private def startRun(

sql/connect/server/src/test/scala/org/apache/spark/sql/connect/pipelines/SparkDeclarativePipelinesServerSuite.scala

@@ -20,7 +20,7 @@ package org.apache.spark.sql.connect.pipelines
 import java.util.UUID
 
 import org.apache.spark.connect.proto
-import org.apache.spark.connect.proto.{DatasetType, Expression, PipelineCommand, Relation, UnresolvedTableValuedFunction}
+import org.apache.spark.connect.proto.{DatasetType, Expression, PipelineCommand, PipelineCommandResult, Relation, UnresolvedTableValuedFunction}
 import org.apache.spark.connect.proto.PipelineCommand.{DefineDataset, DefineFlow}
 import org.apache.spark.internal.Logging
 import org.apache.spark.sql.connect.service.{SessionKey, SparkConnectService}
@@ -486,4 +486,138 @@ class SparkDeclarativePipelinesServerSuite
       assert(graphsAfter.isEmpty, "Graph should be removed after drop")
     }
   }
+
+  gridTest("DefineDataset returns resolved data name for default catalog/schema")(
+    Seq(
+      ("TEMPORARY_VIEW", DatasetType.TEMPORARY_VIEW, "tv"),
+      ("TABLE", DatasetType.TABLE, "tb"),
+      ("MV", DatasetType.MATERIALIZED_VIEW, "mv")
+    )
+  ) { case (_, datasetType, datasetName) =>
+    withRawBlockingStub { implicit stub =>
+      val graphId = createDataflowGraph
+      assert(Option(graphId).isDefined)
+
+      val defineDataset = DefineDataset
+        .newBuilder()
+        .setDataflowGraphId(graphId)
+        .setDatasetName(datasetName)
+        .setDatasetType(datasetType)
+
+      val pipelineCmd = PipelineCommand.newBuilder()
+        .setDefineDataset(defineDataset)
+        .build()
+
+      val res = sendPlan(buildPlanFromPipelineCommand(pipelineCmd)).getPipelineCommandResult
+      assert(res !== PipelineCommandResult.getDefaultInstance)
+      assert(res.hasDefineDatasetResult)
+
+      val graphResult = res.getDefineDatasetResult
+      if (datasetName.contains("tv")) {
+        assert(graphResult.getResolvedDataName == s"`${datasetName}`")
+      } else {
+        assert(graphResult.getResolvedDataName == s"`spark_catalog`.`default`.`${datasetName}`")
+      }
+    }
+  }
+
+  gridTest("DefineDataset returns resolved data name for custom catalog/schema")(
+    Seq(
+      ("TEMPORARY_VIEW", DatasetType.TEMPORARY_VIEW,
+        "tv", "custom_catalog", "custom_db", "`tv`"),
+      ("TABLE", DatasetType.TABLE,
+        "tb", "my_catalog", "my_db", "`my_catalog`.`my_db`.`tb`"),
+      ("MV", DatasetType.MATERIALIZED_VIEW,
+        "mv", "another_catalog", "another_db", "`another_catalog`.`another_db`.`mv`")
+    )
+  ) {
+    case (_, datasetType, datasetName, defaultCatalog, defaultDatabase, expectedResolvedName) =>
+      withRawBlockingStub { implicit stub =>
+        // Build and send the CreateDataflowGraph command with custom catalog/db
+        val graphId = sendPlan(
+          buildCreateDataflowGraphPlan(
+            proto.PipelineCommand.CreateDataflowGraph
+              .newBuilder()
+              .setDefaultCatalog(defaultCatalog)
+              .setDefaultDatabase(defaultDatabase)
+              .build()
+          )
+        ).getPipelineCommandResult.getCreateDataflowGraphResult.getDataflowGraphId
+
+        assert(graphId.nonEmpty)
+
+        // Build DefineDataset with the created graphId and dataset info
+        val defineDataset = DefineDataset
+          .newBuilder()
+          .setDataflowGraphId(graphId)
+          .setDatasetName(datasetName)
+          .setDatasetType(datasetType)
+
+        val pipelineCmd = PipelineCommand.newBuilder()
+          .setDefineDataset(defineDataset)
+          .build()
+
+        val res = sendPlan(buildPlanFromPipelineCommand(pipelineCmd)).getPipelineCommandResult
+
+        assert(res !== PipelineCommandResult.getDefaultInstance)
+        assert(res.hasDefineDatasetResult)
+
+        val graphResult = res.getDefineDatasetResult
+        assert(graphResult.getResolvedDataName == expectedResolvedName)
+      }
+  }
+
+  gridTest("DefineFlow returns resolved flow name")(
+    Seq(
+      ("MV default catalog/schema",
+      "spark_catalog", "default", "`spark_catalog`.`default`.`mv`"),
+      ("MV custom catalog/schema",
+      "custom_catalog", "custom_db", "`custom_catalog`.`custom_db`.`mv`")
+    )
+  ) { case (_, defaultCatalog, defaultDatabase, expectedResolvedFlowName) =>
+    withRawBlockingStub { implicit stub =>
+      val graphId =
+      if (expectedResolvedFlowName.contains("custom")) {
+        sendPlan(
+        buildCreateDataflowGraphPlan(
+          proto.PipelineCommand.CreateDataflowGraph
+            .newBuilder()
+            .setDefaultCatalog(defaultCatalog)
+            .setDefaultDatabase(defaultDatabase)
+            .build()
+        )
+      ).getPipelineCommandResult.getCreateDataflowGraphResult.getDataflowGraphId
+      } else {
+        createDataflowGraph
+      }
+
+      assert(graphId.nonEmpty)
+      val defineFlow = DefineFlow
+        .newBuilder()
+        .setDataflowGraphId(graphId)
+        .setFlowName("mv")
+        .setTargetDatasetName("mv")
+        .setRelation(Relation
+                    .newBuilder()
+                    .setUnresolvedTableValuedFunction(
+                      UnresolvedTableValuedFunction
+                        .newBuilder()
+                        .setFunctionName("range")
+                        .addArguments(Expression
+                          .newBuilder()
+                          .setLiteral(Expression.Literal.newBuilder().setInteger(5).build())
+                          .build())
+                        .build())
+                    .build())
+        .build()
+      val pipelineCmd = PipelineCommand.newBuilder()
+        .setDefineFlow(defineFlow)
+        .build()
+      val res = sendPlan(buildPlanFromPipelineCommand(pipelineCmd)).getPipelineCommandResult
+      assert(res !== PipelineCommandResult.getDefaultInstance)
+      assert(res.hasDefineFlowResult)
+      val graphResult = res.getDefineFlowResult
+      assert(graphResult.getResolvedFlowName == expectedResolvedFlowName)
+    }
+  }
 }