diff --git a/csharp/src/DatabricksConnection.cs b/csharp/src/DatabricksConnection.cs index e0d72cca5..b572cac89 100644 --- a/csharp/src/DatabricksConnection.cs +++ b/csharp/src/DatabricksConnection.cs @@ -132,6 +132,23 @@ internal class DatabricksConnection : SparkHttpConnection private bool _sessionDeleteTelemetryEmitted; internal TelemetrySessionContext? TelemetrySession => _telemetry.Session; + // Set once the server-side session is known to be closed/expired (e.g. inactivity + // timeout). The session handle held by this connection is then permanently stale, so + // every subsequent operation should fail fast rather than reusing the dead handle. + private volatile bool _sessionInvalid; + + /// + /// True when the server-side session has been closed or has expired. Once set, the + /// connection can no longer execute statements and must be disposed and re-created. + /// + internal bool IsSessionInvalid => _sessionInvalid; + + /// + /// Marks the connection's server-side session as closed/expired. Idempotent and + /// thread-safe. Subsequent operations fail fast with . + /// + internal void MarkSessionInvalid() => _sessionInvalid = true; + /// /// RecyclableMemoryStreamManager for LZ4 decompression. /// If provided by Database, this is shared across all connections for optimal pooling. diff --git a/csharp/src/DatabricksSessionExpiredException.cs b/csharp/src/DatabricksSessionExpiredException.cs new file mode 100644 index 000000000..5f527df58 --- /dev/null +++ b/csharp/src/DatabricksSessionExpiredException.cs @@ -0,0 +1,93 @@ +/* +* Copyright (c) 2025 ADBC Drivers Contributors +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +using System; +using Apache.Arrow.Adbc; + +namespace AdbcDrivers.Databricks +{ + /// + /// Raised when the server-side session backing a connection has been closed or has + /// expired (e.g. due to an inactivity timeout) and the connection can no longer be used. + /// + /// The server reports this as HTTP 400 with a Thrift error message containing + /// "Invalid SessionHandle". Without this typed exception the condition surfaces as a + /// generic, misleading "An unexpected error occurred while fetching results / Couldn't + /// connect to server" error, which makes it impossible for callers to distinguish a + /// recoverable "reconnect" situation from a genuine network/transport failure. + /// + /// Callers that catch this exception should dispose the connection and open a new one. + /// + public class DatabricksSessionExpiredException : DatabricksException + { + /// + /// Substring present in the server's error message for a closed/expired session. + /// Both the inactivity-timeout variant ("Invalid SessionHandle: Session [..] is closed") + /// and the explicitly-closed variant ("Invalid SessionHandle: SessionHandle [..]") + /// contain this phrase. + /// + internal const string ServerErrorSignature = "Invalid SessionHandle"; + + /// + /// Message used when failing fast on a connection already known to have an invalid session. + /// + internal const string FastFailMessage = + "The Databricks session has expired or was closed by the server and is no longer usable. " + + "Open a new connection to continue."; + + public DatabricksSessionExpiredException(string message) + : base(message, AdbcStatusCode.InvalidState) + { + } + + public DatabricksSessionExpiredException(string message, Exception innerException) + : base(message, AdbcStatusCode.InvalidState, innerException) + { + } + + /// + /// Determines whether the given exception (or any exception in its inner / aggregate + /// chain) represents a closed or expired server-side session. + /// + internal static bool IsSessionExpired(Exception? exception) + { + switch (exception) + { + case null: + return false; + case DatabricksSessionExpiredException: + return true; + case AggregateException aggregate: + foreach (Exception inner in aggregate.InnerExceptions) + { + if (IsSessionExpired(inner)) + { + return true; + } + } + return false; + } + + if (exception.Message != null && + exception.Message.IndexOf(ServerErrorSignature, StringComparison.OrdinalIgnoreCase) >= 0) + { + return true; + } + + return IsSessionExpired(exception.InnerException); + } + } +} diff --git a/csharp/src/DatabricksStatement.cs b/csharp/src/DatabricksStatement.cs index d0b56ec0e..183750544 100644 --- a/csharp/src/DatabricksStatement.cs +++ b/csharp/src/DatabricksStatement.cs @@ -211,60 +211,118 @@ private void RecordError(StatementTelemetryContext ctx, Exception ex) CaptureRetryCount(ctx); } - public override QueryResult ExecuteQuery() + /// + /// Tier 2 fast-fail: if the connection's server-side session is already known to be + /// closed/expired, throw immediately instead of issuing an RPC that would reuse the + /// stale handle and return the same opaque error. + /// + private void ThrowIfSessionInvalid() + { + if (((DatabricksConnection)Connection).IsSessionInvalid) + { + throw new DatabricksSessionExpiredException(DatabricksSessionExpiredException.FastFailMessage); + } + } + + /// + /// Tier 1: when an execution failure indicates a closed/expired server-side session, + /// mark the connection invalid and rethrow as a clear . + /// Otherwise returns so the caller can rethrow the original exception (preserving its stack). + /// + private void ThrowIfSessionExpired(Exception ex) { - var ctx = IsMetadataCommand - ? CreateMetadataTelemetryContext() - : CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Query); - if (ctx == null) return MaybeWrapComplexTypes(base.ExecuteQuery()); + if (!DatabricksSessionExpiredException.IsSessionExpired(ex)) + { + return; + } + + ((DatabricksConnection)Connection).MarkSessionInvalid(); - // Expose ctx to NewReader so the operation status poller can update PollCount/PollLatencyMs (PECO-2992). - PendingTelemetryContext = ctx; + // Already the clean typed exception (e.g. surfaced by a nested statement) — preserve it. + if (ex is DatabricksSessionExpiredException sessionEx) + { + throw sessionEx; + } + + throw new DatabricksSessionExpiredException( + "The Databricks session has expired or was closed by the server " + + "(e.g. due to an inactivity timeout). The connection can no longer be used; " + + "open a new connection to continue. Server error: " + ex.Message, + ex); + } + + public override QueryResult ExecuteQuery() + { + ThrowIfSessionInvalid(); try { - QueryResult result = base.ExecuteQuery(); - // Store the UNWRAPPED result for telemetry: EmitTelemetry inspects - // _lastQueryResult.Stream via `is CloudFetchReader/DatabricksCompositeReader` - // to read chunk metrics and IsCompressed/ResultFormat. ComplexTypeSerializingStream - // would mask those types, so keep the real reader here and wrap only on return. - _lastQueryResult = result; - RecordSuccess(ctx); - return MaybeWrapComplexTypes(result); + var ctx = IsMetadataCommand + ? CreateMetadataTelemetryContext() + : CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Query); + if (ctx == null) return MaybeWrapComplexTypes(base.ExecuteQuery()); + + // Expose ctx to NewReader so the operation status poller can update PollCount/PollLatencyMs (PECO-2992). + PendingTelemetryContext = ctx; + try + { + QueryResult result = base.ExecuteQuery(); + // Store the UNWRAPPED result for telemetry: EmitTelemetry inspects + // _lastQueryResult.Stream via `is CloudFetchReader/DatabricksCompositeReader` + // to read chunk metrics and IsCompressed/ResultFormat. ComplexTypeSerializingStream + // would mask those types, so keep the real reader here and wrap only on return. + _lastQueryResult = result; + RecordSuccess(ctx); + return MaybeWrapComplexTypes(result); + } + catch (Exception ex) + { + RecordError(ctx, ex); + // Emit telemetry immediately on error (won't reach Dispose) + EmitTelemetry(ctx); + PendingTelemetryContext = null; // Clear to avoid double emission + throw; + } } catch (Exception ex) { - RecordError(ctx, ex); - // Emit telemetry immediately on error (won't reach Dispose) - EmitTelemetry(ctx); - PendingTelemetryContext = null; // Clear to avoid double emission + ThrowIfSessionExpired(ex); // Tier 1: reclassify a closed/expired session; else rethrow as-is. throw; } } public override async ValueTask ExecuteQueryAsync() { - var ctx = IsMetadataCommand - ? CreateMetadataTelemetryContext() - : CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Query); - if (ctx == null) return MaybeWrapComplexTypes(await base.ExecuteQueryAsync()); - - // Expose ctx to NewReader so the operation status poller can update PollCount/PollLatencyMs (PECO-2992). - PendingTelemetryContext = ctx; + ThrowIfSessionInvalid(); try { - QueryResult result = await base.ExecuteQueryAsync(); - // Store the UNWRAPPED result for telemetry (see ExecuteQuery for rationale): - // the wrapper would mask CloudFetchReader/DatabricksCompositeReader from EmitTelemetry. - _lastQueryResult = result; - RecordSuccess(ctx); - return MaybeWrapComplexTypes(result); + var ctx = IsMetadataCommand + ? CreateMetadataTelemetryContext() + : CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Query); + if (ctx == null) return MaybeWrapComplexTypes(await base.ExecuteQueryAsync()); + + // Expose ctx to NewReader so the operation status poller can update PollCount/PollLatencyMs (PECO-2992). + PendingTelemetryContext = ctx; + try + { + QueryResult result = await base.ExecuteQueryAsync(); + // Store the UNWRAPPED result for telemetry (see ExecuteQuery for rationale): + // the wrapper would mask CloudFetchReader/DatabricksCompositeReader from EmitTelemetry. + _lastQueryResult = result; + RecordSuccess(ctx); + return MaybeWrapComplexTypes(result); + } + catch (Exception ex) + { + RecordError(ctx, ex); + // Emit telemetry immediately on error (won't reach Dispose) + EmitTelemetry(ctx); + PendingTelemetryContext = null; // Clear to avoid double emission + throw; + } } catch (Exception ex) { - RecordError(ctx, ex); - // Emit telemetry immediately on error (won't reach Dispose) - EmitTelemetry(ctx); - PendingTelemetryContext = null; // Clear to avoid double emission + ThrowIfSessionExpired(ex); // Tier 1: reclassify a closed/expired session; else rethrow as-is. throw; } } @@ -286,44 +344,62 @@ private QueryResult MaybeWrapComplexTypes(QueryResult result) public override UpdateResult ExecuteUpdate() { - var ctx = CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Update); - if (ctx == null) return base.ExecuteUpdate(); - - PendingTelemetryContext = ctx; + ThrowIfSessionInvalid(); try { - UpdateResult result = base.ExecuteUpdate(); - RecordSuccess(ctx); - return result; + var ctx = CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Update); + if (ctx == null) return base.ExecuteUpdate(); + + PendingTelemetryContext = ctx; + try + { + UpdateResult result = base.ExecuteUpdate(); + RecordSuccess(ctx); + return result; + } + catch (Exception ex) + { + RecordError(ctx, ex); + // Emit telemetry immediately on error (won't reach Dispose) + EmitTelemetry(ctx); + PendingTelemetryContext = null; // Clear to avoid double emission + throw; + } } catch (Exception ex) { - RecordError(ctx, ex); - // Emit telemetry immediately on error (won't reach Dispose) - EmitTelemetry(ctx); - PendingTelemetryContext = null; // Clear to avoid double emission + ThrowIfSessionExpired(ex); // Tier 1: reclassify a closed/expired session; else rethrow as-is. throw; } } public override async Task ExecuteUpdateAsync() { - var ctx = CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Update); - if (ctx == null) return await base.ExecuteUpdateAsync(); - - PendingTelemetryContext = ctx; + ThrowIfSessionInvalid(); try { - UpdateResult result = await base.ExecuteUpdateAsync(); - RecordSuccess(ctx); - return result; + var ctx = CreateTelemetryContext(Telemetry.Proto.Statement.Types.Type.Update); + if (ctx == null) return await base.ExecuteUpdateAsync(); + + PendingTelemetryContext = ctx; + try + { + UpdateResult result = await base.ExecuteUpdateAsync(); + RecordSuccess(ctx); + return result; + } + catch (Exception ex) + { + RecordError(ctx, ex); + // Emit telemetry immediately on error (won't reach Dispose) + EmitTelemetry(ctx); + PendingTelemetryContext = null; // Clear to avoid double emission + throw; + } } catch (Exception ex) { - RecordError(ctx, ex); - // Emit telemetry immediately on error (won't reach Dispose) - EmitTelemetry(ctx); - PendingTelemetryContext = null; // Clear to avoid double emission + ThrowIfSessionExpired(ex); // Tier 1: reclassify a closed/expired session; else rethrow as-is. throw; } } diff --git a/csharp/test/E2E/InvalidSessionHandleE2ETest.cs b/csharp/test/E2E/InvalidSessionHandleE2ETest.cs new file mode 100644 index 000000000..0f9f3d6b4 --- /dev/null +++ b/csharp/test/E2E/InvalidSessionHandleE2ETest.cs @@ -0,0 +1,131 @@ +/* +* Copyright (c) 2025 ADBC Drivers Contributors +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +using System; +using System.Collections.Generic; +using System.Threading.Tasks; +using Apache.Arrow.Adbc; +using Apache.Arrow.Adbc.Tests; +using Apache.Hive.Service.Rpc.Thrift; +using AdbcDrivers.HiveServer2.Hive2; +using Xunit; +using Xunit.Abstractions; + +namespace AdbcDrivers.Databricks.Tests +{ + /// + /// Reproduces the customer-reported "Invalid SessionHandle: Session [...] is closed" + /// (HTTP 400 BAD_REQUEST) error WITHOUT waiting for the server-side session timeout + /// (8-12 hours). + /// + /// Context: a customer hit + /// HTTP Response code: 400, BAD_REQUEST: Invalid SessionHandle: Session [..] is closed + /// because the server closed the session due to inactivity timeout while the driver + /// still held the (now stale) session handle. Lowering the server timeout to make this + /// testable was rejected as too risky (impacts other workspaces). + /// + /// This test induces the identical server-side condition deterministically: + /// 1. Open a connection (server opens a session, driver stores the TSessionHandle). + /// 2. Send a CloseSession RPC for that handle directly, WITHOUT disposing the + /// connection object — so the driver keeps the now-stale handle, exactly as it + /// would after a server-side timeout. The HTTP transport stays alive. + /// 3. Execute a query, which sends ExecuteStatement with the stale handle. + /// 4. Server returns INVALID_HANDLE_STATUS / "Invalid SessionHandle ... is closed". + /// + /// This lets us exercise any Thrift API call against a closed session on demand, + /// covering the scenarios the customer needs without a config change or a multi-hour wait. + /// + public class InvalidSessionHandleE2ETest : TestBase + { + public InvalidSessionHandleE2ETest(ITestOutputHelper? outputHelper) + : base(outputHelper, new DatabricksTestEnvironment.Factory()) + { + Skip.IfNot(Utils.CanExecuteTestConfig(TestConfigVariable)); + } + + [SkippableFact] + public async Task ExecuteAgainstClosedSessionReturnsInvalidSessionHandle() + { + var parameters = new Dictionary + { + [DatabricksParameters.Protocol] = "thrift", + }; + + var connection = NewConnection(TestConfiguration, parameters); + try + { + // Reach the underlying driver connection to access the live session handle + // and the Thrift client. (Test assembly has InternalsVisibleTo access.) + var databricksConnection = (DatabricksConnection)connection; + + // 1. Establish the session by running a trivial query. + using (var warmup = connection.CreateStatement()) + { + warmup.SqlQuery = "SELECT 1"; + await warmup.ExecuteQueryAsync(); + } + + TSessionHandle? sessionHandle = databricksConnection.SessionHandle; + Assert.NotNull(sessionHandle); + OutputHelper?.WriteLine( + $"Session opened: {new Guid(sessionHandle!.SessionId.Guid)}"); + + // 2. Close the session server-side WITHOUT disposing the connection. + // This mimics the server's inactivity-timeout cleanup: the session is + // gone server-side, but the driver still holds the handle and the + // HTTP transport is still open. + var closeResp = await databricksConnection.Client.CloseSession( + new TCloseSessionReq(sessionHandle)); + OutputHelper?.WriteLine( + $"CloseSession status: {closeResp.Status.StatusCode}"); + + // 3. + 4. Execute against the now-closed session. Tier 1: the driver must + // surface a clear, typed DatabricksSessionExpiredException rather than the + // old opaque "An unexpected error occurred while fetching results / + // Couldn't connect to server" wrapper. + using var statement = connection.CreateStatement(); + statement.SqlQuery = "SELECT 1"; + + var ex = await Assert.ThrowsAnyAsync( + () => statement.ExecuteQueryAsync().AsTask()); + + OutputHelper?.WriteLine($"Tier 1 — got typed exception: {ex.GetType().Name}: {ex.Message}"); + + // Message clearly states the session expired/closed (not a connectivity error)... + Assert.Contains("session has expired or was closed", ex.Message, StringComparison.OrdinalIgnoreCase); + // ...and the underlying server signature is preserved for diagnostics. + Assert.Contains("Invalid SessionHandle", ex.Message, StringComparison.OrdinalIgnoreCase); + + // Tier 2: the connection is now marked invalid, so a subsequent execute must + // fail fast with the same typed exception instead of reusing the stale handle. + using var secondStatement = connection.CreateStatement(); + secondStatement.SqlQuery = "SELECT 1"; + + var fastFailEx = await Assert.ThrowsAnyAsync( + () => secondStatement.ExecuteQueryAsync().AsTask()); + + OutputHelper?.WriteLine($"Tier 2 — fast-fail exception: {fastFailEx.GetType().Name}: {fastFailEx.Message}"); + Assert.Contains("no longer usable", fastFailEx.Message, StringComparison.OrdinalIgnoreCase); + } + finally + { + // Connection is already server-side-closed; Dispose may emit a benign + // CloseSession failure. Swallow it so the test result reflects the assertions. + try { connection.Dispose(); } catch { /* session already closed */ } + } + } + } +} diff --git a/csharp/test/Unit/DatabricksSessionExpiredExceptionTests.cs b/csharp/test/Unit/DatabricksSessionExpiredExceptionTests.cs new file mode 100644 index 000000000..620e304b9 --- /dev/null +++ b/csharp/test/Unit/DatabricksSessionExpiredExceptionTests.cs @@ -0,0 +1,103 @@ +/* +* Copyright (c) 2025 ADBC Drivers Contributors +* +* Licensed under the Apache License, Version 2.0 (the "License"); +* you may not use this file except in compliance with the License. +* You may obtain a copy of the License at +* +* http://www.apache.org/licenses/LICENSE-2.0 +* +* Unless required by applicable law or agreed to in writing, software +* distributed under the License is distributed on an "AS IS" BASIS, +* WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. +* See the License for the specific language governing permissions and +* limitations under the License. +*/ + +using System; +using System.Net.Http; +using AdbcDrivers.Databricks; +using Xunit; + +namespace AdbcDrivers.Databricks.Tests.Unit +{ + /// + /// Unit tests for 's detection logic. + /// These run without a workspace, so they cover the classifier in CI. + /// + public class DatabricksSessionExpiredExceptionTests + { + // The shape the driver actually produces when a stale session handle hits the server: + // HiveServer2Reader wraps a TTransportException ("Couldn't connect to server") that wraps + // the HttpRequestException carrying the Thrift "Invalid SessionHandle" header text. + private static Exception BuildRealWorldChain() + { + var httpEx = new HttpRequestException( + "Thrift server error: INVALID_STATE: Invalid SessionHandle: SessionHandle " + + "[01f1604b-abec-124f-a31b-292033255223]. (HTTP 400 Bad Request)"); + var transportEx = new Exception("Couldn't connect to server: " + httpEx.Message, httpEx); + return new Exception("An unexpected error occurred while fetching results. '" + transportEx.Message + "'", transportEx); + } + + [Fact] + public void IsSessionExpired_DirectInvalidSessionHandleMessage_ReturnsTrue() + { + var ex = new Exception("INVALID_STATE: Invalid SessionHandle: SessionHandle [abc]."); + Assert.True(DatabricksSessionExpiredException.IsSessionExpired(ex)); + } + + [Fact] + public void IsSessionExpired_TimeoutClosedVariant_ReturnsTrue() + { + // The inactivity-timeout variant the customer reported. + var ex = new Exception("BAD_REQUEST: Invalid SessionHandle: Session [abc] is closed"); + Assert.True(DatabricksSessionExpiredException.IsSessionExpired(ex)); + } + + [Fact] + public void IsSessionExpired_NestedTransportWrappedChain_ReturnsTrue() + { + Assert.True(DatabricksSessionExpiredException.IsSessionExpired(BuildRealWorldChain())); + } + + [Fact] + public void IsSessionExpired_InsideAggregateException_ReturnsTrue() + { + var agg = new AggregateException("One or more errors occurred.", BuildRealWorldChain()); + Assert.True(DatabricksSessionExpiredException.IsSessionExpired(agg)); + } + + [Fact] + public void IsSessionExpired_AlreadyTypedException_ReturnsTrue() + { + var ex = new DatabricksSessionExpiredException(DatabricksSessionExpiredException.FastFailMessage); + Assert.True(DatabricksSessionExpiredException.IsSessionExpired(ex)); + } + + [Fact] + public void IsSessionExpired_Null_ReturnsFalse() + { + Assert.False(DatabricksSessionExpiredException.IsSessionExpired(null)); + } + + [Fact] + public void IsSessionExpired_UnrelatedError_ReturnsFalse() + { + // A genuine connectivity failure must NOT be misclassified as a session expiry. + var ex = new Exception( + "An unexpected error occurred while fetching results. " + + "'Couldn't connect to server: System.Net.Http.HttpRequestException: " + + "Connection refused (HTTP 503 Service Unavailable)'"); + Assert.False(DatabricksSessionExpiredException.IsSessionExpired(ex)); + } + + [Fact] + public void IsSessionExpired_OtherInvalidHandleButNotSession_ReturnsFalse() + { + // An invalid *operation* handle is a different condition and must not be treated + // as a session expiry (the connection's session is still valid). + var ex = new Exception("INVALID_STATE: Invalid OperationHandle: OperationHandle [abc]."); + Assert.False(DatabricksSessionExpiredException.IsSessionExpired(ex)); + } + } +}