internal ClickHouseDataReader(ClickHouseTable table, ClickHouseTcpClient.Session session, CancellationTokenSource?sessionTokenSource)
{
_currentTable = table ?? throw new ArgumentNullException(nameof(table));
_session = session ?? throw new ArgumentNullException(nameof(session));
_sessionTokenSource = sessionTokenSource;
_reinterpretedColumnsCache = new IClickHouseTableColumn[_currentTable.Columns.Count];
_recordsAffected = _currentTable.Header.RowCount;
_blockHeader = _currentTable.Header;
State = ClickHouseDataReaderState.Data;
}
public async Task SkipNextResult(ClickHouseDataReaderState readBlock, int expectedCount)
{
await using var cn = await OpenConnectionAsync();
await using var cmd = cn.CreateCommand();
cmd.Extremes = true;
cmd.CommandText = "select x, sum(y) as v from (SELECT number%3 + 1 as x, number as y FROM numbers(10)) group by x with totals;";
await using var reader = await cmd.ExecuteReaderAsync();
Assert.Equal(ClickHouseDataReaderState.Data, reader.State);
do
{
switch (reader.State)
{
case ClickHouseDataReaderState.Data:
case ClickHouseDataReaderState.Totals:
case ClickHouseDataReaderState.Extremes:
if (reader.State == readBlock)
{
break;
}
continue;
default:
Assert.True(false, $"Unexpected state: {reader.State}.");
break;
}
int count = 0;
while (await reader.ReadAsync())
{
++count;
}
Assert.False(await reader.ReadAsync());
Assert.Equal(expectedCount, count);
Assert.True(reader.State == ClickHouseDataReaderState.NextResultPending || reader.State == ClickHouseDataReaderState.Closed);
} while (await reader.NextResultAsync());
Assert.Equal(ClickHouseDataReaderState.Closed, reader.State);
}
private async ValueTask Close(bool disposing, bool async)
{
if (_session.IsDisposed || _session.IsFailed)
{
return;
}
if (!(State == ClickHouseDataReaderState.Closed || State == ClickHouseDataReaderState.Broken))
{
await _session.SendCancel(async);
while (true)
{
try
{
var message = _nextResultMessage ?? await _session.ReadMessage(async, CancellationToken.None);
_nextResultMessage = null;
switch (message.MessageCode)
{
case ServerMessageCode.Data:
case ServerMessageCode.Totals:
case ServerMessageCode.Extremes:
var dataMessage = (ServerDataMessage)message;
await _session.SkipTable(dataMessage, async, CancellationToken.None);
continue;
case ServerMessageCode.Error:
State = ClickHouseDataReaderState.Closed;
if (disposing)
{
break;
}
throw ((ServerErrorMessage)message).Exception;
case ServerMessageCode.Progress:
var progressMessage = (ServerProgressMessage)message;
_recordsAffected = progressMessage.Rows;
continue;
case ServerMessageCode.EndOfStream:
State = ClickHouseDataReaderState.Closed;
break;
case ServerMessageCode.ProfileInfo:
continue;
case ServerMessageCode.Pong:
case ServerMessageCode.Hello:
case ServerMessageCode.Log:
throw new ClickHouseException(ClickHouseErrorCodes.ProtocolUnexpectedResponse, $"Unexpected server message: \"{message.MessageCode}\".");
default:
throw new NotSupportedException($"Internal error. Message code \"{message.MessageCode}\" not supported.");
}
}
catch (ClickHouseHandledException ex)
{
if (!disposing)
{
throw;
}
State = ClickHouseDataReaderState.Broken;
await _session.SetFailed(ex.InnerException, false, async);
return;
}
catch (Exception ex)
{
State = ClickHouseDataReaderState.Broken;
var aggrEx = await _session.SetFailed(ex, false, async);
if (disposing)
{
return;
}
if (aggrEx != null)
{
throw aggrEx;
}
throw;
}
break;
}
}
_session.Dispose();
_sessionTokenSource?.Dispose();
}
private async ValueTask <bool> Read(bool nextResult, bool async, CancellationToken cancellationToken)
{
if (!nextResult)
{
if (_rowIndex == _currentTable.Header.RowCount)
{
return(false);
}
if (++_rowIndex < _currentTable.Header.RowCount)
{
return(true);
}
}
while (true)
{
ClickHouseTable nextTable;
try
{
var message = _nextResultMessage;
if (message == null)
{
message = await _session.ReadMessage(async, cancellationToken);
}
else
{
_nextResultMessage = null;
}
switch (message.MessageCode)
{
case ServerMessageCode.Data:
switch (State)
{
case ClickHouseDataReaderState.NextResultPending:
State = ClickHouseDataReaderState.Data;
goto case ClickHouseDataReaderState.Data;
case ClickHouseDataReaderState.Data:
{
var dataMessage = (ServerDataMessage)message;
nextTable = await _session.ReadTable(dataMessage, _columnSettings, async, cancellationToken);
break;
}
case ClickHouseDataReaderState.Totals:
case ClickHouseDataReaderState.Extremes:
{
var dataMessage = (ServerDataMessage)message;
var table = await _session.SkipTable(dataMessage, async, cancellationToken);
if (table.RowCount != 0 || table.Columns.Count != 0)
{
throw new ClickHouseException(ClickHouseErrorCodes.ProtocolUnexpectedResponse, "Unexpected data block after totals or extremes.");
}
continue;
}
default:
goto UNEXPECTED_RESPONSE;
}
break;
case ServerMessageCode.Error:
State = ClickHouseDataReaderState.Closed;
_session.Dispose();
throw ((ServerErrorMessage)message).Exception;
case ServerMessageCode.Progress:
var progressMessage = (ServerProgressMessage)message;
_recordsAffected = progressMessage.Rows;
continue;
case ServerMessageCode.EndOfStream:
State = ClickHouseDataReaderState.Closed;
_session.Dispose();
return(false);
case ServerMessageCode.ProfileInfo:
continue;
case ServerMessageCode.Totals:
switch (State)
{
case ClickHouseDataReaderState.NextResultPending:
State = ClickHouseDataReaderState.Totals;
goto case ClickHouseDataReaderState.Totals;
case ClickHouseDataReaderState.Totals:
var totalsMessage = (ServerDataMessage)message;
nextTable = await _session.ReadTable(totalsMessage, _columnSettings, async, cancellationToken);
break;
case ClickHouseDataReaderState.Data:
case ClickHouseDataReaderState.Extremes:
_nextResultMessage = message;
State = ClickHouseDataReaderState.NextResultPending;
return(false);
default:
goto UNEXPECTED_RESPONSE;
}
break;
case ServerMessageCode.Extremes:
switch (State)
{
case ClickHouseDataReaderState.NextResultPending:
State = ClickHouseDataReaderState.Extremes;
goto case ClickHouseDataReaderState.Extremes;
case ClickHouseDataReaderState.Extremes:
var extremesMessage = (ServerDataMessage)message;
nextTable = await _session.ReadTable(extremesMessage, _columnSettings, async, cancellationToken);
break;
case ClickHouseDataReaderState.Data:
case ClickHouseDataReaderState.Totals:
_nextResultMessage = message;
State = ClickHouseDataReaderState.NextResultPending;
return(false);
default:
goto UNEXPECTED_RESPONSE;
}
break;
case ServerMessageCode.Pong:
case ServerMessageCode.Hello:
case ServerMessageCode.Log:
UNEXPECTED_RESPONSE:
throw new ClickHouseException(ClickHouseErrorCodes.ProtocolUnexpectedResponse, $"Unexpected server message: \"{message.MessageCode}\".");
default:
throw new NotSupportedException($"Internal error. Message code \"{message.MessageCode}\" not supported.");
}
}
catch (ClickHouseHandledException)
{
throw;
}
catch (ClickHouseServerException)
{
throw;
}
catch (Exception ex)
{
State = ClickHouseDataReaderState.Broken;
var aggrEx = await _session.SetFailed(ex, true, async);
if (aggrEx != null)
{
throw aggrEx;
}
throw;
}
if (nextTable.Header.RowCount == 0)
{
continue;
}
_currentTable = nextTable;
_reinterpretedColumnsCache = new IClickHouseTableColumn[_currentTable.Columns.Count];
_recordsAffected += nextTable.Header.RowCount;
_rowIndex = nextResult ? -1 : 0;
return(true);
}
}
