Skip to content

Commit 241477e

Browse files
authored
feat(csharp/src/Drivers/Databricks): Remove redundant CloseOperation for GetColumnsAsync (#3132)
[A previous PR](#3093) fixed redundant CloseOperations that were throwing errors and causing failure for dbr 11.3 and below. GetColumnsAsync uses a seperate response handling path, which does not call GetQueryResults (where the previous fix was implemented), so this operation still fails. This PR adds a fix
1 parent 5a3bfcd commit 241477e

File tree

2 files changed

+70
-14
lines changed

2 files changed

+70
-14
lines changed

csharp/src/Drivers/Apache/Hive2/HiveServer2Statement.cs

Lines changed: 3 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -523,6 +523,9 @@ protected virtual async Task<QueryResult> GetColumnsAsync(CancellationToken canc
523523
cancellationToken);
524524
OperationHandle = resp.OperationHandle;
525525

526+
// Set _directResults so that dispose logic can check if operation was already closed
527+
_directResults = resp.DirectResults;
528+
526529
// Common variables declared upfront
527530
TGetResultSetMetadataResp metadata;
528531
Schema schema;

csharp/test/Drivers/Databricks/E2E/StatementTests.cs

Lines changed: 67 additions & 14 deletions
Original file line numberDiff line numberDiff line change
@@ -133,25 +133,78 @@ public async Task CanGetCrossReferenceFromChildTableDatabricks()
133133
}
134134

135135
/// <summary>
136-
/// Verifies that Dispose() can be called on metadata query statements without throwing
137-
/// "Invalid OperationHandle" errors. This tests the fix for the issue where the server
138-
/// auto-closes operations but the client still tries to close them during disposal.
136+
/// Comprehensive test that verifies disposal works for all statement types without throwing exceptions.
137+
/// This prevents regressions like the GetColumns disposal bug where _directResults wasn't set properly.
139138
/// </summary>
140-
[SkippableFact]
141-
public async Task CanDisposeMetadataQueriesWithoutError()
139+
[SkippableTheory]
140+
[InlineData("ExecuteStatement", "SELECT 1 as test_column")]
141+
[InlineData("GetCatalogs", "GetCatalogs")]
142+
[InlineData("GetSchemas", "GetSchemas")]
143+
[InlineData("GetTables", "GetTables")]
144+
[InlineData("GetColumns", "GetColumns")]
145+
[InlineData("GetPrimaryKeys", "GetPrimaryKeys")]
146+
[InlineData("GetCrossReference", "GetCrossReference")]
147+
[InlineData("GetColumnsExtended", "GetColumnsExtended")]
148+
public async Task AllStatementTypesDisposeWithoutErrors(string statementType, string sqlCommand)
142149
{
143-
// Test a simple metadata command that's most likely to trigger the issue
144150
var statement = Connection.CreateStatement();
145-
statement.SetOption(ApacheParameters.IsMetadataCommand, "true");
146-
statement.SqlQuery = "GetSchemas";
147151

148-
// Execute the metadata query
149-
QueryResult queryResult = await statement.ExecuteQueryAsync();
150-
Assert.NotNull(queryResult.Stream);
152+
try
153+
{
154+
if (statementType == "ExecuteStatement")
155+
{
156+
// Regular SQL statement
157+
statement.SqlQuery = sqlCommand;
158+
}
159+
else
160+
{
161+
// Metadata command
162+
statement.SetOption(ApacheParameters.IsMetadataCommand, "true");
163+
statement.SqlQuery = sqlCommand;
164+
165+
// Set required parameters for specific metadata commands
166+
if (sqlCommand is "GetColumns" or "GetPrimaryKeys" or "GetCrossReference" or "GetColumnsExtended")
167+
{
168+
statement.SetOption(ApacheParameters.CatalogName, TestConfiguration.Metadata.Catalog);
169+
statement.SetOption(ApacheParameters.SchemaName, TestConfiguration.Metadata.Schema);
170+
statement.SetOption(ApacheParameters.TableName, TestConfiguration.Metadata.Table);
171+
}
172+
173+
if (sqlCommand == "GetCrossReference")
174+
{
175+
// GetCrossReference needs foreign table parameters too
176+
statement.SetOption(ApacheParameters.ForeignCatalogName, TestConfiguration.Metadata.Catalog);
177+
statement.SetOption(ApacheParameters.ForeignSchemaName, TestConfiguration.Metadata.Schema);
178+
statement.SetOption(ApacheParameters.ForeignTableName, TestConfiguration.Metadata.Table);
179+
}
180+
}
181+
182+
// Execute the statement
183+
QueryResult queryResult = await statement.ExecuteQueryAsync();
184+
Assert.NotNull(queryResult.Stream);
151185

152-
// This should not throw "Invalid OperationHandle" errors
153-
// The fix ensures _directResults is set so dispose logic works correctly
154-
statement.Dispose();
186+
// Consume at least one batch to ensure the operation completes
187+
var batch = await queryResult.Stream.ReadNextRecordBatchAsync();
188+
// Note: batch might be null for empty results, that's OK
189+
190+
// The critical test: disposal should not throw any exceptions
191+
// This specifically tests the fix for the GetColumns bug where _directResults wasn't set
192+
var exception = Record.Exception(() => statement.Dispose());
193+
Assert.Null(exception);
194+
}
195+
catch (Exception ex)
196+
{
197+
// If execution fails, we still want to test disposal
198+
OutputHelper?.WriteLine($"Statement execution failed for {statementType}: {ex.Message}");
199+
200+
// Even if execution failed, disposal should not throw
201+
var disposalException = Record.Exception(() => statement.Dispose());
202+
Assert.Null(disposalException);
203+
204+
// Re-throw the original exception if we want to investigate execution failures
205+
// For now, we'll skip the test if execution fails since disposal is our main concern
206+
Skip.If(true, $"Skipping disposal test for {statementType} due to execution failure: {ex.Message}");
207+
}
155208
}
156209

157210
[SkippableFact]

0 commit comments

Comments
 (0)