From 7c53f15dc8165d25c1c7600aa6e2f312d9158fab Mon Sep 17 00:00:00 2001 From: Stephen Zhao Date: Wed, 26 Mar 2025 12:04:01 -0400 Subject: [PATCH] Fix `CommandBehavior` flag read --- .../Trino.Data.ADO/Server/TrinoCommand.cs | 40 ++++++++++--------- 1 file changed, 22 insertions(+), 18 deletions(-) diff --git a/trino-csharp/Trino.Data.ADO/Server/TrinoCommand.cs b/trino-csharp/Trino.Data.ADO/Server/TrinoCommand.cs index 0047b96..0283be0 100644 --- a/trino-csharp/Trino.Data.ADO/Server/TrinoCommand.cs +++ b/trino-csharp/Trino.Data.ADO/Server/TrinoCommand.cs @@ -208,26 +208,30 @@ public async Task RunQuery(long bufferSizeBytes = Constants.Defa protected override async Task ExecuteDbDataReaderAsync(CommandBehavior behavior, CancellationToken cancellationToken) { RecordExecutor queryExecutor; - switch (behavior) + + if ((behavior == CommandBehavior.Default) + || ((behavior & CommandBehavior.SingleResult) == CommandBehavior.SingleResult)) { - case CommandBehavior.Default: // Single result means only run one query. Trino only supports one query. - case CommandBehavior.SingleResult: - queryExecutor = await RunQuery().ConfigureAwait(false); - break; - case CommandBehavior.SingleRow: - // Single row requires the reader to be created and the first row to be read. - queryExecutor = await RunQuery().ConfigureAwait(false); - return new TrinoDataReader(queryExecutor); - case CommandBehavior.SchemaOnly: - queryExecutor = await RunNonQuery().ConfigureAwait(false); - break; - case CommandBehavior.CloseConnection: - // Trino has no concept of a connection because every call is a new connection. - queryExecutor = await RunQuery().ConfigureAwait(false); - break; - default: - throw new NotSupportedException(); + queryExecutor = await RunQuery().ConfigureAwait(false); + } + else if ((behavior & CommandBehavior.SingleRow) == CommandBehavior.SingleRow) + { + // Single row requires the reader to be created and the first row to be read. + queryExecutor = await RunQuery().ConfigureAwait(false); + } + else if ((behavior & CommandBehavior.SchemaOnly) == CommandBehavior.SchemaOnly) + { + queryExecutor = await RunNonQuery().ConfigureAwait(false); + } + else if ((behavior & CommandBehavior.CloseConnection) == CommandBehavior.CloseConnection) + { + // Trino has no concept of a connection because every call is a new connection. + queryExecutor = await RunQuery().ConfigureAwait(false); + } + else + { + throw new NotSupportedException(); } // always wait for the schema when creating an IEnumerable