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

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

@@ -75,6 +75,10 @@ public DbCommand Command
                     {
                         _command = Connection.CreateCommand();
                         _command.CommandText = _source.CommandText;
+                        if (_source.CommandTimeOut.HasValue)
+                        {
+                            _command.CommandTimeout = _source.CommandTimeOut.Value;
+                        }
                     }
                     return _command;
                 }

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

@@ -25,6 +25,26 @@ public DatabaseSource(DbProviderFactory providerFactory, string connectionString
             CommandText = commandText;
         }
 
+        /// <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="commandTimeOut">The timeout(in seconds) for database command.</param>
+        public DatabaseSource(DbProviderFactory providerFactory, string connectionString, string commandText, int commandTimeOut)
+        {
+            Contracts.CheckValue(providerFactory, nameof(providerFactory));
+            Contracts.CheckValue(connectionString, nameof(connectionString));
+            Contracts.CheckValue(commandText, nameof(commandText));
+
+            ProviderFactory = providerFactory;
+            ConnectionString = connectionString;
+            CommandText = commandText;
+            CommandTimeOut = commandTimeOut;
+        }
+
+        /// <summary>Gets the timeout for database command.</summary>
+        public int? CommandTimeOut { 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,20 @@ public DatabaseLoaderTests(ITestOutputHelper output)
 
         [LightGBMFact]
         public void IrisLightGbm()
+        {
+            DatabaseSource dbs = GetIrisDatabaseSource("SELECT * FROM {0}");
+            IrisLightGbmImpl(dbs);
+        }
+
+        [LightGBMFact]
+        public void IrisLightGbmWithTimeOut()
+        {
+            DatabaseSource dbs = GetIrisDatabaseSourceWithTimeOut("WAITFOR DELAY '00:00:01'; SELECT * FROM {0}", 1);
+            var ex = Assert.Throws<System.Reflection.TargetInvocationException>(() => IrisLightGbmImpl(dbs));
+            Assert.Contains("Timeout expired", ex.InnerException.Message);
+        }
+
+        private void IrisLightGbmImpl(DatabaseSource dbs)
         {
             var mlContext = new MLContext(seed: 1);
 
@@ -41,7 +55,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"))
@@ -225,6 +239,22 @@ private DatabaseSource GetIrisDatabaseSource(string command)
                     String.Format(command, TestDatasets.irisDbSQLite.trainFilename));
         }
 
+        private DatabaseSource GetIrisDatabaseSourceWithTimeOut(string command, int commandTimeOut)
+        {
+            if (RuntimeInformation.IsOSPlatform(OSPlatform.Windows))
+                return new DatabaseSource(
+                    SqlClientFactory.Instance,
+                    GetMSSQLConnectionString(TestDatasets.irisDb.name),
+                    String.Format(command, $@"""{TestDatasets.irisDb.trainFilename}"""),
+                    commandTimeOut);
+            else
+                return new DatabaseSource(
+                    SQLiteFactory.Instance,
+                    GetSQLiteConnectionString(TestDatasets.irisDbSQLite.name),
+                    String.Format(command, TestDatasets.irisDbSQLite.trainFilename),
+                    commandTimeOut);
+        }
+
         private string GetMSSQLConnectionString(string databaseName)
         {
             var databaseFile = Path.GetFullPath(Path.Combine("TestDatabases", $"{databaseName}.mdf"));