using System;

using System.Diagnostics;

using System.IO;

using System.Threading;

using System.Threading.Tasks;

using Microsoft.Extensions.Logging;

using Npgsql.BackendMessages;

using Npgsql.Internal;

using InfiniteTimeout = System.Threading.Timeout;

using static Npgsql.Util.Statics;

#pragma warning disable 1591

namespace Npgsql;

public sealed class NpgsqlRawCopyStream : Stream, ICancelable

}

#endregion

#region Write

public override void Write(byte[] buffer, int offset, int count)

{

ValidateArguments(buffer, offset, count);

Write(new ReadOnlySpan<byte>(buffer, offset, count));

}

public override Task WriteAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)

{

ValidateArguments(buffer, offset, count);

return WriteAsync(new Memory<byte>(buffer, offset, count), cancellationToken).AsTask();

}

public override void Write(ReadOnlySpan<byte> buffer)

{

CheckDisposed();

if (!CanWrite)

throw new InvalidOperationException("Stream not open for writing");

if (buffer.Length == 0) { return; }

if (buffer.Length <= _writeBuf.WriteSpaceLeft)

{

_writeBuf.WriteBytes(buffer);

return;

}

// Value is too big, flush.

Flush();

if (buffer.Length <= _writeBuf.WriteSpaceLeft)

{

_writeBuf.WriteBytes(buffer);

return;

}

// Value is too big even after a flush - bypass the buffer and write directly.

_writeBuf.DirectWrite(buffer);

}

public override ValueTask WriteAsync(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken = default)

{

CheckDisposed();

if (!CanWrite)

throw new InvalidOperationException("Stream not open for writing");

cancellationToken.ThrowIfCancellationRequested();

return WriteAsyncInternal(buffer, cancellationToken);

async ValueTask WriteAsyncInternal(ReadOnlyMemory<byte> buffer, CancellationToken cancellationToken)

{

if (buffer.Length == 0)

return;

if (buffer.Length <= _writeBuf.WriteSpaceLeft)

{

_writeBuf.WriteBytes(buffer.Span);

return;

}

// Value is too big, flush.

await FlushAsync(true, cancellationToken).ConfigureAwait(false);

if (buffer.Length <= _writeBuf.WriteSpaceLeft)

{

_writeBuf.WriteBytes(buffer.Span);

return;

}

// Value is too big even after a flush - bypass the buffer and write directly.

await _writeBuf.DirectWrite(buffer, true, cancellationToken).ConfigureAwait(false);

}

}

public override void Flush() => FlushAsync(async: false).GetAwaiter().GetResult();

public override Task FlushAsync(CancellationToken cancellationToken)

{

if (cancellationToken.IsCancellationRequested)

return Task.FromCanceled(cancellationToken);

return FlushAsync(async: true, cancellationToken);

}

Task FlushAsync(bool async, CancellationToken cancellationToken = default)

{

CheckDisposed();

return _writeBuf.Flush(async, cancellationToken);

}

#endregion

#region Read

public override int Read(byte[] buffer, int offset, int count)

{

ValidateArguments(buffer, offset, count);

return Read(new Span<byte>(buffer, offset, count));

}

public override Task<int> ReadAsync(byte[] buffer, int offset, int count, CancellationToken cancellationToken)

{

ValidateArguments(buffer, offset, count);

return ReadAsync(new Memory<byte>(buffer, offset, count), cancellationToken).AsTask();

}

public override int Read(Span<byte> span)

{

CheckDisposed();

if (!CanRead)

throw new InvalidOperationException("Stream not open for reading");

var count = ReadCore(span.Length, false).GetAwaiter().GetResult();

if (count > 0)

_readBuf.ReadBytes(span.Slice(0, count));

return count;

}

public override ValueTask<int> ReadAsync(Memory<byte> buffer, CancellationToken cancellationToken)

{

CheckDisposed();

if (!CanRead)

throw new InvalidOperationException("Stream not open for reading");

cancellationToken.ThrowIfCancellationRequested();

return ReadAsyncInternal();

async ValueTask<int> ReadAsyncInternal()

{

var count = await ReadCore(buffer.Length, true, cancellationToken).ConfigureAwait(false);

if (count > 0)

_readBuf.ReadBytes(buffer.Slice(0, count).Span);

return count;

}

}

async ValueTask<int> ReadCore(int count, bool async, CancellationToken cancellationToken = default)

{

if (_state == CopyStreamState.Consumed)

return 0;

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

if (_leftToReadInDataMsg == 0)

{

IBackendMessage msg;

try

{

// We've consumed the current DataMessage (or haven't yet received the first),

// read the next message

msg = await _connector.ReadMessage(async).ConfigureAwait(false);

}

catch (Exception e)

{

if (_state != CopyStreamState.Disposed)

{

TraceSetException(e);

Cleanup();

}

throw;

}

switch (msg.Code)

{

case BackendMessageCode.CopyData:

_leftToReadInDataMsg = ((CopyDataMessage)msg).Length;

break;

case BackendMessageCode.CopyDone:

Expect<CommandCompleteMessage>(await _connector.ReadMessage(async).ConfigureAwait(false), _connector);

Expect<ReadyForQueryMessage>(await _connector.ReadMessage(async).ConfigureAwait(false), _connector);

_state = CopyStreamState.Consumed;

return 0;

default:

throw _connector.UnexpectedMessageReceived(msg.Code);

}

}

Debug.Assert(_leftToReadInDataMsg > 0);

// If our buffer is empty, read in more. Otherwise return whatever is there, even if the

// user asked for more (normal socket behavior)

if (_readBuf.ReadBytesLeft == 0)

await _readBuf.ReadMore(async).ConfigureAwait(false);

Debug.Assert(_readBuf.ReadBytesLeft > 0);

var maxCount = Math.Min(_readBuf.ReadBytesLeft, _leftToReadInDataMsg);

if (count > maxCount)

count = maxCount;

_leftToReadInDataMsg -= count;

return count;

}

#endregion

#region Cancel

/// <summary>

/// Cancels and terminates an ongoing operation. Any data already written will be discarded.

/// </summary>

public void Cancel() => Cancel(async: false).GetAwaiter().GetResult();

/// <summary>

/// Cancels and terminates an ongoing operation. Any data already written will be discarded.

/// </summary>

public Task CancelAsync() => Cancel(async: true);

async Task Cancel(bool async)

{

CheckDisposed();

if (CanWrite)

{

_writeBuf.EndCopyMode();

_writeBuf.Clear();

await _connector.WriteCopyFail(async).ConfigureAwait(false);

await _connector.Flush(async).ConfigureAwait(false);

try

{

var msg = await _connector.ReadMessage(async).ConfigureAwait(false);

// The CopyFail should immediately trigger an exception from the read above.

throw _connector.Break(

new NpgsqlException("Expected ErrorResponse when cancelling COPY but got: " + msg.Code));

}

catch (PostgresException e)

{

// TODO: NpgsqlBinaryImporter doesn't cleanup on cancellation

// And instead relies on users disposing the object

// We probably should do the same here

Cleanup();

if (e.SqlState != PostgresErrorCodes.QueryCanceled)

{

TraceSetException(e);

throw;

}

TraceStop();

}

}

else

{

_connector.PerformPostgresCancellation();

}

}

#endregion

#region Dispose

protected override void Dispose(bool disposing) => DisposeAsync(disposing, false).GetAwaiter().GetResult();

public override ValueTask DisposeAsync()

=> DisposeAsync(disposing: true, async: true);

async ValueTask DisposeAsync(bool disposing, bool async)

{

if (_state == CopyStreamState.Disposed || !disposing)

return;

try

{

_connector.CurrentCopyOperation = null;

if (CanWrite)

{

try

{

await FlushAsync(async).ConfigureAwait(false);

_writeBuf.EndCopyMode();

await _connector.WriteCopyDone(async).ConfigureAwait(false);

await _connector.Flush(async).ConfigureAwait(false);

Expect<CommandCompleteMessage>(await _connector.ReadMessage(async).ConfigureAwait(false), _connector);

Expect<ReadyForQueryMessage>(await _connector.ReadMessage(async).ConfigureAwait(false), _connector);

TraceStop();

}

catch (Exception e)

{

TraceSetException(e);

throw;

}

}

else

{

try

{

if (_state != CopyStreamState.Consumed && _state != CopyStreamState.Uninitialized)

{

if (_leftToReadInDataMsg > 0)

{

await _readBuf.Skip(async, _leftToReadInDataMsg).ConfigureAwait(false);

}

_connector.SkipUntil(BackendMessageCode.ReadyForQuery);

}

TraceStop();

}

catch (OperationCanceledException e) when (e.InnerException is PostgresException { SqlState: PostgresErrorCodes.QueryCanceled })

{

LogMessages.CopyOperationCancelled(_copyLogger, _connector.Id);

TraceStop();

}

catch (Exception e)

{

LogMessages.ExceptionWhenDisposingCopyOperation(_copyLogger, _connector.Id, e);

TraceSetException(e);

}

}

}

finally

{

Cleanup();

}

}

#pragma warning disable CS8625

void Cleanup()

{

Debug.Assert(_state != CopyStreamState.Disposed);

LogMessages.CopyOperationCompleted(_copyLogger, _connector.Id);

_connector.EndUserAction();

_connector.CurrentCopyOperation = null;

_connector = null;

_readBuf = null;

_writeBuf = null;

_state = CopyStreamState.Disposed;

}

#pragma warning restore CS8625

void CheckDisposed()

{

if (_state == CopyStreamState.Disposed) {

throw new ObjectDisposedException(nameof(NpgsqlRawCopyStream), "The COPY operation has already ended.");

}

}

#endregion

#region Unsupported

public override bool CanSeek => false;

public override long Seek(long offset, SeekOrigin origin) => throw new NotSupportedException();

public override void SetLength(long value) => throw new NotSupportedException();

public override long Length => throw new NotSupportedException();

public override long Position

{

get => throw new NotSupportedException();

set => throw new NotSupportedException();

}

#endregion

#region Input validation

static void ValidateArguments(byte[] buffer, int offset, int count)

{

ArgumentNullException.ThrowIfNull(buffer);

ArgumentOutOfRangeException.ThrowIfNegative(offset);

ArgumentOutOfRangeException.ThrowIfNegative(count);

if (buffer.Length - offset < count)

ThrowHelper.ThrowArgumentException("Offset and length were out of bounds for the array or count is greater than the number of elements from index to the end of the source collection.");

}

#endregion

#region Tracing

private void TraceSetImport()

{

if (_activity is not null)

{

NpgsqlActivitySource.SetOperation(_activity, "COPY FROM");

}

}

private void TraceSetExport()

{

if (_activity is not null)

{

NpgsqlActivitySource.SetOperation(_activity, "COPY TO");

}

}

private void TraceStop()

{

if (_activity is not null)

{

NpgsqlActivitySource.CopyStop(_activity);

_activity = null;

}

}

private void TraceSetException(Exception e)

{

if (_activity is not null)

{

NpgsqlActivitySource.SetException(_activity, e);

_activity = null;

}

}

#endregion

#region Enums

enum CopyStreamState

}

#endregion Enums

public sealed class NpgsqlCopyTextWriter : StreamWriter, ICancelable

public sealed class NpgsqlCopyTextReader : StreamReader, ICancelable