Skip to content

Redo binary exporter column reading #5464

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
NinoFloris merged 3 commits into from
Dec 3, 2023
Merged

Redo binary exporter column reading #5464

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
3 commits
Select commit Hold shift + click to select a range
28927ed
Redo binary exporter column reading
NinoFloris Nov 27, 2023
511e038
Codify why we don't commit directly after a read call is finished
NinoFloris Dec 2, 2023
eb9c410
Fix naming
NinoFloris Dec 3, 2023
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
8 changes: 4 additions & 4 deletions src/Npgsql/Internal/PgReader.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -12,6 +12,9 @@ namespace Npgsql.Internal;

public class PgReader
{
// We don't want to add a ton of memory pressure for large strings.
internal const int MaxPreparedTextReaderSize = 1024 * 64;

readonly NpgsqlReadBuffer _buffer;

bool _resumable;
Expand Down Expand Up @@ -210,11 +213,8 @@ public ValueTask<TextReader> GetTextReaderAsync(Encoding encoding, CancellationT

async ValueTask<TextReader> GetTextReader(bool async, Encoding encoding, CancellationToken cancellationToken)
{
// We don't want to add a ton of memory pressure for large strings.
const int maxPreparedSize = 1024 * 64;

_requiresCleanup = true;
if (CurrentRemaining > _buffer.ReadBytesLeft || CurrentRemaining > maxPreparedSize)
if (CurrentRemaining > _buffer.ReadBytesLeft || CurrentRemaining > MaxPreparedTextReaderSize)
return new StreamReader(GetColumnStream(), encoding, detectEncodingFromByteOrderMarks: false);

if (_preparedTextReader is { IsDisposed: false })
Expand Down
153 changes: 77 additions & 76 deletions src/Npgsql/NpgsqlBinaryExporter.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -35,7 +35,7 @@ public sealed class NpgsqlBinaryExporter : ICancelable
/// <summary>
/// The number of columns, as returned from the backend in the CopyInResponse.
/// </summary>
internal int NumColumns { get; private set; }
int NumColumns { get; set; }

PgConverterInfo[] _columnInfoCache;

Expand Down Expand Up @@ -140,16 +140,18 @@ async Task ReadHeader(bool async)

async ValueTask<int> StartRow(bool async, CancellationToken cancellationToken = default)
{

CheckDisposed();
ThrowIfDisposed();
if (_isConsumed)
return -1;

using var registration = _connector.StartNestedCancellableOperation(cancellationToken);

// Consume and advance any active column.
if (_column >= 0)
await Commit(async, resumableOp: false).ConfigureAwait(false);
{
await Commit(async).ConfigureAwait(false);
_column++;
}

// The very first row (i.e. _column == -1) is included in the header's CopyData message.
// Otherwise we need to read in a new CopyData row (the docs specify that there's a CopyData
Expand Down Expand Up @@ -210,29 +212,6 @@ public ValueTask<T> ReadAsync<T>(CancellationToken cancellationToken = default)
ValueTask<T> Read<T>(bool async, CancellationToken cancellationToken = default)
=> Read<T>(async, null, cancellationToken);

PgConverterInfo CreateConverterInfo(Type type, NpgsqlDbType? npgsqlDbType = null)
{
var options = _connector.SerializerOptions;
PgTypeId? pgTypeId = null;
if (npgsqlDbType.HasValue)
{
pgTypeId = npgsqlDbType.Value.ToDataTypeName() is { } name
? options.GetCanonicalTypeId(name)
// Handle plugin types via lookup.
: GetRepresentationalOrDefault(npgsqlDbType.Value.ToUnqualifiedDataTypeNameOrThrow());
}
var info = options.GetTypeInfo(type, pgTypeId)
?? throw new NotSupportedException($"Reading is not supported for type '{type}'{(npgsqlDbType is null ? "" : $" and NpgsqlDbType '{npgsqlDbType}'")}");
// Binary export has no type info so we only do caller-directed interpretation of data.
return info.Bind(new Field("?", info.PgTypeId!.Value, -1), DataFormat.Binary);

PgTypeId GetRepresentationalOrDefault(string dataTypeName)
{
var type = options.DatabaseInfo.GetPostgresType(dataTypeName);
return options.ToCanonicalTypeId(type.GetRepresentationalType());
}
}

/// <summary>
/// Reads the current column, returns its value according to <paramref name="type"/> and
/// moves ahead to the next column.
Expand Down Expand Up @@ -269,39 +248,22 @@ public ValueTask<T> ReadAsync<T>(NpgsqlDbType type, CancellationToken cancellati

async ValueTask<T> Read<T>(bool async, NpgsqlDbType? type, CancellationToken cancellationToken)
{
CheckDisposed();
if (_column is BeforeRow)
ThrowHelper.ThrowInvalidOperationException("Not reading a row");
ThrowIfNotOnRow();

using var registration = _connector.StartNestedCancellableOperation(cancellationToken, attemptPgCancellation: false);

// Allow one more read if the field is a db null.
// We cannot allow endless rereads otherwise it becomes quite unclear when a column advance happens.
if (PgReader is { Initialized: true, Resumable: true, FieldSize: -1 })
{
await Commit(async, resumableOp: false).ConfigureAwait(false);
return DbNullOrThrow();
}
if (!IsInitializedAndAtStart)
await MoveNextColumn(async, resumableOp: false).ConfigureAwait(false);

// We must commit the current column before reading the next one unless it was an IsNull call.
PgConverterInfo info;
bool asObject;
if (!PgReader.Initialized || !PgReader.Resumable || PgReader.CurrentRemaining != PgReader.FieldSize)
if (PgReader.FieldSize is (-1 or 0) and var fieldSize)
{
await Commit(async, resumableOp: false).ConfigureAwait(false);
info = GetInfo(type, out asObject);

// We need to get info after potential I/O as we don't know beforehand at what column we're at.
var columnLen = await ReadColumnLenIfNeeded(async, resumableOp: false).ConfigureAwait(false);
if (_column == NumColumns)
ThrowHelper.ThrowInvalidOperationException("No more columns left in the current row");

if (columnLen is -1)
// Commit, otherwise we'll have no way of knowing this column is finished.
await Commit(async).ConfigureAwait(false);
if (fieldSize is -1)
return DbNullOrThrow();

}
else
info = GetInfo(type, out asObject);

var info = GetInfo(type, out var asObject);

T result;
if (async)
Expand All @@ -323,6 +285,14 @@ async ValueTask<T> Read<T>(bool async, NpgsqlDbType? type, CancellationToken can

return result;

static T DbNullOrThrow()
{
// When T is a Nullable<T>, we support returning null
if (default(T) is null && typeof(T).IsValueType)
return default!;
throw new InvalidCastException("Column is null");
}

PgConverterInfo GetInfo(NpgsqlDbType? type, out bool asObject)
{
ref var cachedInfo = ref _columnInfoCache[_column];
Expand All @@ -331,12 +301,27 @@ PgConverterInfo GetInfo(NpgsqlDbType? type, out bool asObject)
return converterInfo;
}

T DbNullOrThrow()
PgConverterInfo CreateConverterInfo(Type type, NpgsqlDbType? npgsqlDbType = null)
{
// When T is a Nullable<T>, we support returning null
if (default(T) is null && typeof(T).IsValueType)
return default!;
throw new InvalidCastException("Column is null");
var options = _connector.SerializerOptions;
PgTypeId? pgTypeId = null;
if (npgsqlDbType.HasValue)
{
pgTypeId = npgsqlDbType.Value.ToDataTypeName() is { } name
? options.GetCanonicalTypeId(name)
// Handle plugin types via lookup.
: GetRepresentationalOrDefault(npgsqlDbType.Value.ToUnqualifiedDataTypeNameOrThrow());
}
var info = options.GetTypeInfo(type, pgTypeId)
?? throw new NotSupportedException($"Reading is not supported for type '{type}'{(npgsqlDbType is null ? "" : $" and NpgsqlDbType '{npgsqlDbType}'")}");
// Binary export has no type info so we only do caller-directed interpretation of data.
return info.Bind(new Field("?", info.PgTypeId!.Value, -1), DataFormat.Binary);

PgTypeId GetRepresentationalOrDefault(string dataTypeName)
{
var type = options.DatabaseInfo.GetPostgresType(dataTypeName);
return options.ToCanonicalTypeId(type.GetRepresentationalType());
}
}
}

Expand All @@ -347,8 +332,11 @@ public bool IsNull
{
get
{
Commit(async: false, resumableOp: true);
return ReadColumnLenIfNeeded(async: false, resumableOp: true).GetAwaiter().GetResult() is -1;
ThrowIfNotOnRow();
if (!IsInitializedAndAtStart)
return MoveNextColumn(async: false, resumableOp: true).GetAwaiter().GetResult() is -1;

return PgReader.FieldSize is - 1;
}
}

Expand All @@ -365,46 +353,59 @@ public Task SkipAsync(CancellationToken cancellationToken = default)

async Task Skip(bool async, CancellationToken cancellationToken = default)
{
CheckDisposed();
ThrowIfNotOnRow();

using var registration = _connector.StartNestedCancellableOperation(cancellationToken);

// We allow IsNull to have been called before skip.
if (PgReader.Initialized && PgReader is not { Resumable: true, FieldSize: -1 })
await Commit(async, resumableOp: false).ConfigureAwait(false);
await ReadColumnLenIfNeeded(async, resumableOp: false).ConfigureAwait(false);
if (!IsInitializedAndAtStart)
await MoveNextColumn(async, resumableOp: false).ConfigureAwait(false);

await PgReader.Consume(async, cancellationToken: cancellationToken).ConfigureAwait(false);

// Commit, otherwise we'll have no way of knowing this column is finished.
if (PgReader.FieldSize is -1 or 0)
await Commit(async).ConfigureAwait(false);
}

#endregion

#region Utilities

ValueTask Commit(bool async, bool resumableOp)
{
var resuming = PgReader is { Initialized: true, Resumable: true } && resumableOp;
if (!resuming)
_column++;
bool IsInitializedAndAtStart => PgReader.Initialized && (PgReader.FieldSize is -1 || PgReader.FieldOffset is 0);

ValueTask Commit(bool async)
{
if (async)
return PgReader.CommitAsync(resuming);
return PgReader.CommitAsync(resuming: false);

PgReader.Commit(resuming);
PgReader.Commit(resuming: false);
return new();
}

async ValueTask<int> ReadColumnLenIfNeeded(bool async, bool resumableOp)
async ValueTask<int> MoveNextColumn(bool async, bool resumableOp)
{
if (PgReader is { Initialized: true, Resumable: true, FieldSize: -1 })
return -1;
if (async)
await PgReader.CommitAsync(resuming: false).ConfigureAwait(false);
else
PgReader.Commit(resuming: false);

if (_column + 1 == NumColumns)
ThrowHelper.ThrowInvalidOperationException("No more columns left in the current row");
_column++;
await _buf.Ensure(4, async).ConfigureAwait(false);
var columnLen = _buf.ReadInt32();
PgReader.Init(columnLen, DataFormat.Binary, resumableOp);
return PgReader.FieldSize;
}

void CheckDisposed()
void ThrowIfNotOnRow()
{
ThrowIfDisposed();
if (_column is BeforeRow)
ThrowHelper.ThrowInvalidOperationException("Not reading a row");
}

void ThrowIfDisposed()
{
if (_isDisposed)
ThrowHelper.ThrowObjectDisposedException(nameof(NpgsqlBinaryExporter), "The COPY operation has already ended.");
Expand Down
67 changes: 67 additions & 0 deletions test/Npgsql.Tests/CopyTests.cs
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -510,6 +510,73 @@ public async Task Wrong_table_definition_binary_export()
Assert.That(await conn.ExecuteScalarAsync("SELECT 1"), Is.EqualTo(1));
}

[Test, IssueLink("https://github.com/npgsql/npgsql/issues/5457")]
public async Task MixedOperations()
{
if (IsMultiplexing)
Assert.Ignore("Multiplexing: fails");
using var conn = await OpenConnectionAsync();

var reader = conn.BeginBinaryExport("""
COPY (values ('foo', 1), ('bar', null), (null, 2)) TO STDOUT BINARY
""");
while(reader.StartRow() != -1)
{
string? col1 = null;
if (reader.IsNull)
reader.Skip();
else
col1 = reader.Read<string>();
int? col2 = null;
if (reader.IsNull)
reader.Skip();
else
col2 = reader.Read<int>();
}
}

[Test]
public async Task ReadMoreColumnsThanExist()
{
if (IsMultiplexing)
Assert.Ignore("Multiplexing: fails");
using var conn = await OpenConnectionAsync();

var reader = conn.BeginBinaryExport("""
COPY (values ('foo', 1), ('bar', null), (null, 2)) TO STDOUT BINARY
""");
while(reader.StartRow() != -1)
{
string? col1 = null;
if (reader.IsNull)
reader.Skip();
else
col1 = reader.Read<string>();
int? col2 = null;
if (reader.IsNull)
reader.Skip();
else
col2 = reader.Read<int>();

Assert.Throws<InvalidOperationException>(() => _ = reader.IsNull);
}
}

[Test]
public async Task StreamingRead()
{
if (IsMultiplexing)
Assert.Ignore("Multiplexing: fails");
using var conn = await OpenConnectionAsync();

var str = new string('a', PgReader.MaxPreparedTextReaderSize + 1);
var reader = conn.BeginBinaryExport($"""COPY (values ('{str}')) TO STDOUT BINARY""");
while (reader.StartRow() != -1)
{
using var _ = reader.Read<TextReader>(NpgsqlDbType.Text);
}
}

[Test, IssueLink("https://github.com/npgsql/npgsql/issues/2330")]
public async Task Wrong_format_binary_export()
{
Expand Down