Command timeout change in adbc client & improve exception handling

apache · Nov 26, 2024 · 834221a · 834221a
1 parent 153c0d8
commit 834221a
Show file tree

Hide file tree

Showing 3 changed files with 19 additions and 6 deletions.
diff --git a/csharp/src/Client/AdbcCommand.cs b/csharp/src/Client/AdbcCommand.cs
@@ -117,7 +117,11 @@ public override CommandType CommandType
         public override int CommandTimeout
         {
             get => _timeout;
-            set => _timeout = value;
+            set
+            {
+                _timeout = value;
+                _adbcStatement.SetOption("adbc.apache.statement.query_timeout_s", value.ToString());
+            }
         }
 
         protected override DbParameterCollection DbParameterCollection

diff --git a/csharp/src/Drivers/Apache/Hive2/HiveServer2Connection.cs b/csharp/src/Drivers/Apache/Hive2/HiveServer2Connection.cs
@@ -70,9 +70,9 @@ internal TCLIService.Client Client
 
         internal async Task OpenAsync()
         {
+            CancellationToken cancellationToken = ApacheUtility.GetCancellationToken(ConnectTimeoutMilliseconds, ApacheUtility.TimeUnit.Milliseconds);
             try
             {
-                CancellationToken cancellationToken = ApacheUtility.GetCancellationToken(ConnectTimeoutMilliseconds, ApacheUtility.TimeUnit.Milliseconds);
                 TTransport transport = await CreateTransportAsync(cancellationToken);
                 TProtocol protocol = await CreateProtocolAsync(transport, cancellationToken);
                 _transport = protocol.Transport;
@@ -95,9 +95,11 @@ internal async Task OpenAsync()
 
                 SessionHandle = session.SessionHandle;
             }
-            catch (OperationCanceledException ex)
+            catch (Exception ex) when (
+                (ex is TTransportException || ex is OperationCanceledException) &&
+                cancellationToken.IsCancellationRequested)
             {
-                throw new TimeoutException("The operation timed out while attempting to open a session.", ex);
+                throw new HiveServer2Exception("The operation timed out while attempting to open a session. Please try increasing connect timeout.");
             }
             catch (Exception ex)
             {

diff --git a/csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs b/csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs
@@ -20,6 +20,7 @@
 using System.Threading.Tasks;
 using Apache.Arrow.Ipc;
 using Apache.Hive.Service.Rpc.Thrift;
+using Thrift.Transport;
 
 namespace Apache.Arrow.Adbc.Drivers.Apache.Hive2
 {
@@ -61,9 +62,15 @@ private async Task<QueryResult> ExecuteQueryAsyncInternal(CancellationToken canc
 
                 return new QueryResult(-1, Connection.NewReader(this, schema));
             }
-            catch (OperationCanceledException ex)
+            catch (Exception ex) when (
+                (ex is TTransportException || ex is OperationCanceledException) &&
+                cancellationToken.IsCancellationRequested)
             {
-                throw new TimeoutException("The query execution timed out.", ex);
+                throw new TimeoutException("The query execution timed out. Consider increasing the query timeout value.");
+            }
+            catch (Exception ex)
+            {
+                throw new HiveServer2Exception("The query execution timed out.", ex);
             }
         }