Expose StreamingQueryManager (#690)

Author: suhsteve

src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/SparkSessionTests.cs

@@ -2,7 +2,6 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
-using System;
 using System.Collections.Generic;
 using System.Linq;
 using Microsoft.Spark.E2ETest.Utils;
@@ -42,6 +41,8 @@ public void TestSignaturesV2_3_X()
 
             Assert.IsType<RuntimeConfig>(_spark.Conf());
 
+            Assert.IsType<StreamingQueryManager>(_spark.Streams());
+
             Assert.IsType<SparkSession>(_spark.NewSession());
 
             Assert.IsType<DataFrameReader>(_spark.Read());

src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/Streaming/StreamingQueryManagerTests.cs

@@ -0,0 +1,57 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Linq;
+using Microsoft.Spark.E2ETest.Utils;
+using Microsoft.Spark.Sql;
+using Microsoft.Spark.Sql.Streaming;
+using Xunit;
+
+namespace Microsoft.Spark.E2ETest.IpcTests
+{
+    [Collection("Spark E2E Tests")]
+    public class StreamingQueryManagerTests
+    {
+        private readonly SparkSession _spark;
+
+        public StreamingQueryManagerTests(SparkFixture fixture)
+        {
+            _spark = fixture.Spark;
+        }
+
+        /// <summary>
+        /// Test signatures for APIs up to Spark 2.3.*.
+        /// The purpose of this test is to ensure that JVM calls can be successfully made.
+        /// Note that this is not testing functionality of each function.
+        /// </summary>
+        [Fact]
+        public void TestSignaturesV2_3_X()
+        {
+            var intMemoryStream = new MemoryStream<int>(_spark);
+            StreamingQuery sq1 = intMemoryStream
+                .ToDF().WriteStream().QueryName("intQuery").Format("console").Start();
+            string id1 = sq1.Id;
+
+            var stringMemoryStream = new MemoryStream<string>(_spark);
+            StreamingQuery sq2 = stringMemoryStream
+                .ToDF().WriteStream().QueryName("stringQuery").Format("console").Start();
+            string id2 = sq2.Id;
+
+            StreamingQueryManager sqm = _spark.Streams();
+
+            StreamingQuery[] streamingQueries = sqm.Active().ToArray();
+            Assert.Equal(2, streamingQueries.Length);
+
+            Assert.IsType<StreamingQuery>(sqm.Get(id1));
+            Assert.IsType<StreamingQuery>(sqm.Get(id2));
+
+            sqm.ResetTerminated();
+
+            sqm.AwaitAnyTermination(10);
+
+            sq1.Stop();
+            sq2.Stop();
+        }
+    }
+}

src/csharp/Microsoft.Spark.E2ETest/IpcTests/Sql/Streaming/StreamingQueryTests.cs

@@ -0,0 +1,52 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Linq;
+using Microsoft.Spark.E2ETest.Utils;
+using Microsoft.Spark.Sql;
+using Microsoft.Spark.Sql.Streaming;
+using Xunit;
+
+namespace Microsoft.Spark.E2ETest.IpcTests
+{
+    [Collection("Spark E2E Tests")]
+    public class StreamingQueryTests
+    {
+        private readonly SparkSession _spark;
+
+        public StreamingQueryTests(SparkFixture fixture)
+        {
+            _spark = fixture.Spark;
+        }
+
+        /// <summary>
+        /// Test signatures for APIs up to Spark 2.3.*.
+        /// The purpose of this test is to ensure that JVM calls can be successfully made.
+        /// Note that this is not testing functionality of each function.
+        /// </summary>
+        [Fact]
+        public void TestSignaturesV2_3_X()
+        {
+            var intMemoryStream = new MemoryStream<int>(_spark);
+            StreamingQuery sq = intMemoryStream
+                .ToDF().WriteStream().QueryName("testQuery").Format("console").Start();
+
+            Assert.IsType<string>(sq.Name);
+
+            Assert.IsType<string>(sq.Id);
+
+            Assert.IsType<string>(sq.RunId);
+
+            Assert.IsType<bool>(sq.IsActive());
+
+            Assert.IsType<bool>(sq.AwaitTermination(10));
+
+            sq.Explain();
+
+            Assert.Null(sq.Exception());
+
+            sq.Stop();
+        }
+    }
+}

src/csharp/Microsoft.Spark.E2ETest/Utils/MemoryStream.cs

@@ -0,0 +1,40 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.Spark.Interop.Ipc;
+using Microsoft.Spark.Sql;
+
+namespace Microsoft.Spark.E2ETest.Utils
+{
+    /// <summary>
+    /// A source of continually arriving data for a streaming query.
+    /// Produces value stored in memory as they are added by the user.
+    /// </summary>
+    /// <typeparam name="T">
+    /// Specifies the type of the elements contained in the MemoryStream.
+    /// </typeparam>
+    internal class MemoryStream<T> : IJvmObjectReferenceProvider
+    {
+        private readonly JvmObjectReference _jvmObject;
+
+        internal MemoryStream(SparkSession sparkSession)
+        {
+            JvmObjectReference sparkSessionRef =
+                ((IJvmObjectReferenceProvider)sparkSession).Reference;
+
+            _jvmObject = (JvmObjectReference)sparkSessionRef.Jvm.CallStaticJavaMethod(
+                "org.apache.spark.sql.test.TestUtils",
+                "createMemoryStream",
+                sparkSessionRef.Invoke("sqlContext"),
+                typeof(T).Name);
+        }
+
+        JvmObjectReference IJvmObjectReferenceProvider.Reference => _jvmObject;
+
+        internal DataFrame ToDF() => new DataFrame((JvmObjectReference)_jvmObject.Invoke("toDF"));
+
+        // TODO: "addData" returns an Offset. Expose class if needed.	
+        internal void AddData(T[] data) => _jvmObject.Invoke("addData", data);
+    }
+}

src/csharp/Microsoft.Spark/Sql/SparkSession.cs

@@ -128,6 +128,14 @@ public void Dispose()
         public RuntimeConfig Conf() =>
             new RuntimeConfig((JvmObjectReference)_jvmObject.Invoke("conf"));
 
+        /// <summary>
+        /// Returns a <see cref="StreamingQueryManager"/> that allows managing all the
+        /// <see cref="StreamingQuery"/> instances active on <c>this</c> context.
+        /// </summary>
+        /// <returns><see cref="StreamingQueryManager"/> object</returns>
+        public StreamingQueryManager Streams() =>
+            new StreamingQueryManager((JvmObjectReference)_jvmObject.Invoke("streams"));
+
         /// <summary>
         /// Start a new session with isolated SQL configurations, temporary tables, registered
         /// functions are isolated, but sharing the underlying SparkContext and cached data.

src/csharp/Microsoft.Spark/Sql/Streaming/StreamingQuery.cs

@@ -2,6 +2,7 @@
 // The .NET Foundation licenses this file to you under the MIT license.
 // See the LICENSE file in the project root for more information.
 
+using Microsoft.Spark.Interop.Internal.Scala;
 using Microsoft.Spark.Interop.Ipc;
 
 namespace Microsoft.Spark.Sql.Streaming
@@ -22,6 +23,24 @@ public sealed class StreamingQuery : IJvmObjectReferenceProvider
         /// </summary>
         public string Name => (string)_jvmObject.Invoke("name");
 
+        /// <summary>
+        /// Returns the unique id of this query that persists across restarts from checkpoint data.
+        /// That is, this id is generated when a query is started for the first time, and
+        /// will be the same every time it is restarted from checkpoint data. Also see
+        /// <see cref="RunId"/>.
+        /// </summary>
+        public string Id =>
+            (string)((JvmObjectReference)_jvmObject.Invoke("id")).Invoke("toString");
+
+        /// <summary>
+        /// Returns the unique id of this run of the query. That is, every start/restart of
+        /// a query will generated a unique runId. Therefore, every time a query is restarted
+        /// from checkpoint, it will have the same <see cref="Id"/> but different
+        /// <see cref="RunId"/>s.
+        /// </summary>
+        public string RunId =>
+            (string)((JvmObjectReference)_jvmObject.Invoke("runId")).Invoke("toString");
+
         /// <summary>
         /// Returns true if this query is actively running.
         /// </summary>
@@ -67,5 +86,21 @@ public bool AwaitTermination(long timeoutMs) =>
         /// </summary>
         /// <param name="extended">Whether to do extended explain or not</param>
         public void Explain(bool extended = false) => _jvmObject.Invoke("explain", extended);
+
+        /// <summary>
+        /// The <see cref="StreamingQueryException"/> if the query was terminated by an exception,
+        /// null otherwise.
+        /// </summary>
+        public StreamingQueryException Exception()
+        {
+            var optionalException = new Option((JvmObjectReference)_jvmObject.Invoke("exception"));
+            if (optionalException.IsDefined())
+            {
+                var exception = (JvmObjectReference)optionalException.Get();
+                return new StreamingQueryException((string)exception.Invoke("toString"));
+            }
+
+            return null;
+        }
     }
 }

src/csharp/Microsoft.Spark/Sql/Streaming/StreamingQueryException.cs

@@ -0,0 +1,19 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using Microsoft.Spark.Interop.Ipc;
+
+namespace Microsoft.Spark.Sql.Streaming
+{
+    /// <summary>
+    /// Exception that stopped a <see cref="StreamingQuery"/>.
+    /// </summary>
+    public class StreamingQueryException : JvmException
+    {
+        public StreamingQueryException(string message)
+            : base(message)
+        {
+        }
+    }
+}

src/csharp/Microsoft.Spark/Sql/Streaming/StreamingQueryManager.cs

@@ -0,0 +1,99 @@
+// Licensed to the .NET Foundation under one or more agreements.
+// The .NET Foundation licenses this file to you under the MIT license.
+// See the LICENSE file in the project root for more information.
+
+using System.Collections.Generic;
+using System.Linq;
+using Microsoft.Spark.Interop.Ipc;
+
+namespace Microsoft.Spark.Sql.Streaming
+{
+    /// <summary>
+    /// A class to manage all the <see cref="StreamingQuery"/> active
+    /// in a <see cref="SparkSession"/>.
+    /// </summary>
+    public sealed class StreamingQueryManager : IJvmObjectReferenceProvider
+    {
+        private readonly JvmObjectReference _jvmObject;
+
+        internal StreamingQueryManager(JvmObjectReference jvmObject) => _jvmObject = jvmObject;
+
+        JvmObjectReference IJvmObjectReferenceProvider.Reference => _jvmObject;
+
+        /// <summary>
+        /// Returns a list of active queries associated with this SQLContext.
+        /// </summary>
+        /// <returns>Active queries associated with this SQLContext.</returns>
+        public IEnumerable<StreamingQuery> Active() =>
+            ((JvmObjectReference[])_jvmObject.Invoke("active"))
+                .Select(sq => new StreamingQuery(sq));
+
+        /// <summary>
+        /// Returns an active query from this SQLContext or throws exception if an active
+        /// query with this name doesn't exist.
+        /// </summary>
+        /// <param name="id">Id of the <see cref="StreamingQuery"/>.</param>
+        /// <returns>
+        /// <see cref="StreamingQuery"/> if there is an active query with the given id.
+        /// </returns>
+        public StreamingQuery Get(string id) =>
+            new StreamingQuery((JvmObjectReference)_jvmObject.Invoke("get", id));
+
+        /// <summary>
+        /// Wait until any of the queries on the associated SQLContext has terminated since the
+        /// creation of the context, or since <see cref="ResetTerminated"/> was called. If any
+        /// query was terminated with an exception, then the exception will be thrown.
+        ///
+        /// If a query has terminated, then subsequent calls to <see cref="AwaitAnyTermination()"/>
+        /// will either return immediately (if the query was terminated by
+        /// <see cref="StreamingQuery.Stop"/>), or throw the exception immediately (if the query
+        /// was terminated with exception). Use <see cref="ResetTerminated"/> to clear past
+        /// terminations and wait for new terminations.
+        ///
+        /// In the case where multiple queries have terminated since <see cref="ResetTerminated"/>
+        /// was called, if any query has terminated with exception, then
+        /// <see cref="AwaitAnyTermination()"/> will throw any of the exception. For correctly
+        /// documenting exceptions across multiple queries, users need to stop all of them after
+        /// any of them terminates with exception, and then check the
+        /// <see cref="StreamingQuery.Exception"/> for each query.
+        ///
+        /// Throws StreamingQueryException on the JVM if any query has terminated with an
+        /// exception.
+        /// </summary>
+        public void AwaitAnyTermination() => _jvmObject.Invoke("awaitAnyTermination");
+
+        /// <summary>
+        /// Wait until any of the queries on the associated SQLContext has terminated since the
+        /// creation of the context, or since <see cref="ResetTerminated"/> was called. If any
+        /// query was terminated with an exception, then the exception will be thrown.
+        ///
+        /// If a query has terminated, then subsequent calls to <see cref="AwaitAnyTermination()"/>
+        /// will either return immediately (if the query was terminated by
+        /// <see cref="StreamingQuery.Stop"/>), or throw the exception immediately (if the query
+        /// was terminated with exception). Use <see cref="ResetTerminated"/> to clear past
+        /// terminations and wait for new terminations.
+        ///
+        /// In the case where multiple queries have terminated since <see cref="ResetTerminated"/>
+        /// was called, if any query has terminated with exception, then
+        /// <see cref="AwaitAnyTermination()"/> will throw any of the exception. For correctly
+        /// documenting exceptions across multiple queries, users need to stop all of them after
+        /// any of them terminates with exception, and then check the
+        /// <see cref="StreamingQuery.Exception"/> for each query.
+        ///
+        /// Throws StreamingQueryException on the JVM if any query has terminated with an
+        /// exception.
+        /// </summary>
+        /// <param name="timeoutMs">
+        /// Milliseconds to wait for query to terminate. Returns whether the query has terminated
+        /// or not.
+        /// </param>
+        public void AwaitAnyTermination(long timeoutMs) =>
+            _jvmObject.Invoke("awaitAnyTermination", timeoutMs);
+
+        /// <summary>
+        /// Forget about past terminated queries so that <see cref="AwaitAnyTermination()"/> can be
+        /// used again to wait for new terminations
+        /// </summary>
+        public void ResetTerminated() => _jvmObject.Invoke("resetTerminated");
+    }
+}

src/scala/microsoft-spark-2.3.x/src/main/scala/org/apache/spark/sql/test/TestUtils.scala

@@ -0,0 +1,30 @@
+/*	
+ * Licensed to the .NET Foundation under one or more agreements.	
+ * The .NET Foundation licenses this file to you under the MIT license.	
+ * See the LICENSE file in the project root for more information.	
+ */
+
+package org.apache.spark.sql.test
+
+import org.apache.spark.sql.SQLContext
+import org.apache.spark.sql.execution.streaming.MemoryStream
+
+object TestUtils {
+
+  /**
+   * Helper method to create typed MemoryStreams intended for use in unit tests.	
+   * @param sqlContext The SQLContext.
+   * @param streamType The type of memory stream to create. This string is the `Name`
+   *                   property of the dotnet type.
+   * @return A typed MemoryStream.
+   */
+  def createMemoryStream(implicit sqlContext: SQLContext, streamType: String): MemoryStream[_] = {
+    import sqlContext.implicits._
+
+    streamType match {
+      case "Int32" => MemoryStream[Int]
+      case "String" => MemoryStream[String]
+      case _ => throw new Exception(s"$streamType not supported")
+    }
+  }
+}