[Issue4484] Support specifying command timeout while using the database loader

src/Microsoft.ML.Data/DataLoadSave/Database/DatabaseLoaderCursor.cs

@@ -75,6 +75,7 @@ public DbCommand Command
                     {
                         _command = Connection.CreateCommand();
                         _command.CommandText = _source.CommandText;
+                        _command.CommandTimeout = _source.CommandTimeoutInSeconds;
                     }
                     return _command;
                 }

src/Microsoft.ML.Data/DataLoadSave/Database/DatabaseSource.cs

@@ -10,21 +10,37 @@ namespace Microsoft.ML.Data
     /// <summary>Exposes the data required for opening a database for reading.</summary>
     public sealed class DatabaseSource
     {
+        public static readonly int DefaultCommandTimeoutInSeconds = 30;
         /// <summary>Creates a new instance of the <see cref="DatabaseSource" /> class.</summary>
         /// <param name="providerFactory">The factory used to create the <see cref="DbConnection"/>..</param>
         /// <param name="connectionString">The string used to open the connection.</param>
         /// <param name="commandText">The text command to run against the data source.</param>
-        public DatabaseSource(DbProviderFactory providerFactory, string connectionString, string commandText)
+        public DatabaseSource(DbProviderFactory providerFactory, string connectionString, string commandText) :
+            this(providerFactory, connectionString, commandText, DefaultCommandTimeoutInSeconds)
+        {
+        }
+
+        /// <summary>Creates a new instance of the <see cref="DatabaseSource" /> class.</summary>
+        /// <param name="providerFactory">The factory used to create the <see cref="DbConnection"/>..</param>
+        /// <param name="connectionString">The string used to open the connection.</param>
+        /// <param name="commandText">The text command to run against the data source.</param>
+        /// <param name="commandTimeoutInSeconds">The timeout(in seconds) for database command.</param>
+        public DatabaseSource(DbProviderFactory providerFactory, string connectionString, string commandText, int commandTimeoutInSeconds)
         {
             Contracts.CheckValue(providerFactory, nameof(providerFactory));
             Contracts.CheckValue(connectionString, nameof(connectionString));
             Contracts.CheckValue(commandText, nameof(commandText));
+            Contracts.CheckUserArg(commandTimeoutInSeconds >= 0, nameof(commandTimeoutInSeconds));
 
             ProviderFactory = providerFactory;
             ConnectionString = connectionString;
             CommandText = commandText;
+            CommandTimeoutInSeconds = commandTimeoutInSeconds;
         }
 
+        /// <summary>Gets the timeout for database command.</summary>
+        public int CommandTimeoutInSeconds { get; }
+
         /// <summary>Gets the text command to run against the data source.</summary>
         public string CommandText { get; }
 

test/Microsoft.ML.Tests/DatabaseLoaderTests.cs

@@ -27,6 +27,22 @@ public DatabaseLoaderTests(ITestOutputHelper output)
 
         [LightGBMFact]
         public void IrisLightGbm()
+        {
+            DatabaseSource dbs = GetIrisDatabaseSource("SELECT * FROM {0}");
+            IrisLightGbmImpl(dbs);
+        }
+
+        [LightGBMFact]
+        public void IrisLightGbmWithTimeout()
+        {
+            if (!RuntimeInformation.IsOSPlatform(OSPlatform.Windows)) //sqlite does not have built-in command for sleep
+                return;
+            DatabaseSource dbs = GetIrisDatabaseSource("WAITFOR DELAY '00:00:01'; SELECT * FROM {0}", 1);
+            var ex = Assert.Throws<System.Reflection.TargetInvocationException>(() => IrisLightGbmImpl(dbs));
+            Assert.Contains("Timeout", ex.InnerException.Message);
+        }
+
+        private void IrisLightGbmImpl(DatabaseSource dbs)
         {
             var mlContext = new MLContext(seed: 1);
 
@@ -41,7 +57,7 @@ public void IrisLightGbm()
 
             var loader = mlContext.Data.CreateDatabaseLoader(loaderColumns);
 
-            var trainingData = loader.Load(GetIrisDatabaseSource("SELECT * FROM {0}"));
+            var trainingData = loader.Load(dbs);
 
             IEstimator<ITransformer> pipeline = mlContext.Transforms.Conversion.MapValueToKey("Label")
                 .Append(mlContext.Transforms.Concatenate("Features", "SepalLength", "SepalWidth", "PetalLength", "PetalWidth"))
@@ -211,18 +227,20 @@ public void IrisSdcaMaximumEntropy()
         /// SQLite database is used on Linux and MacOS builds.
         /// </summary>
         /// <returns>Return the appropiate Iris DatabaseSource according to build OS.</returns>
-        private DatabaseSource GetIrisDatabaseSource(string command)
+        private DatabaseSource GetIrisDatabaseSource(string command, int commandTimeoutInSeconds = 30)
         {
             if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
                 return new DatabaseSource(
                     SqlClientFactory.Instance,
                     GetMSSQLConnectionString(TestDatasets.irisDb.name),
-                    String.Format(command, $@"""{TestDatasets.irisDb.trainFilename}"""));
+                    String.Format(command, $@"""{TestDatasets.irisDb.trainFilename}"""),
+                    commandTimeoutInSeconds);
             else
                 return new DatabaseSource(
                     SQLiteFactory.Instance,
                     GetSQLiteConnectionString(TestDatasets.irisDbSQLite.name),
-                    String.Format(command, TestDatasets.irisDbSQLite.trainFilename));
+                    String.Format(command, TestDatasets.irisDbSQLite.trainFilename),
+                    commandTimeoutInSeconds);
         }
 
         private string GetMSSQLConnectionString(string databaseName)