<commit_before>import { TransformOptions } from 'stream'
import {
  Mode,
  bindComplete,
  parseComplete,
  closeComplete,
  noData,
  portalSuspended,
  copyDone,
  replicationStart,
  emptyQuery,
  ReadyForQueryMessage,
  CommandCompleteMessage,
  CopyDataMessage,
  CopyResponse,
  NotificationResponseMessage,
  RowDescriptionMessage,
  Field,
  DataRowMessage,
  ParameterStatusMessage,
  BackendKeyDataMessage,
  DatabaseError,
  BackendMessage,
  MessageName,
  AuthenticationMD5Password,
  NoticeMessage,
} from './messages'
import { BufferReader } from './buffer-reader'
import assert from 'assert'

// every message is prefixed with a single bye
const CODE_LENGTH = 1
// every message has an int32 length which includes itself but does
// NOT include the code in the length
const LEN_LENGTH = 4

const HEADER_LENGTH = CODE_LENGTH + LEN_LENGTH

export type Packet = {
  code: number
  packet: Buffer
}

const emptyBuffer = Buffer.allocUnsafe(0)

type StreamOptions = TransformOptions & {
  mode: Mode
}

const enum MessageCodes {
  DataRow = 0x44, // D
  ParseComplete = 0x31, // 1
  BindComplete = 0x32, // 2
  CloseComplete = 0x33, // 3
  CommandComplete = 0x43, // C
  ReadyForQuery = 0x5a, // Z
  NoData = 0x6e, // n
  NotificationResponse = 0x41, // A
  AuthenticationResponse = 0x52, // R
  ParameterStatus = 0x53, // S
  BackendKeyData = 0x4b, // K
  ErrorMessage = 0x45, // E
  NoticeMessage = 0x4e, // N
  RowDescriptionMessage = 0x54, // T
  PortalSuspended = 0x73, // s
  ReplicationStart = 0x57, // W
  EmptyQuery = 0x49, // I
  CopyIn = 0x47, // G
  CopyOut = 0x48, // H
  CopyDone = 0x63, // c
  CopyData = 0x64, // d
}

export type MessageCallback = (msg: BackendMessage) => void

export class Parser {
  private remainingBuffer: Buffer = emptyBuffer
  private reader = new BufferReader()
  private mode: Mode

  constructor(opts?: StreamOptions) {
    if (opts?.mode === 'binary') {
      throw new Error('Binary mode not supported yet')
    }
    this.mode = opts?.mode || 'text'
  }

  public parse(buffer: Buffer, callback: MessageCallback) {
    let combinedBuffer = buffer
    if (this.remainingBuffer.byteLength) {
      combinedBuffer = Buffer.allocUnsafe(this.remainingBuffer.byteLength + buffer.byteLength)
      this.remainingBuffer.copy(combinedBuffer)
      buffer.copy(combinedBuffer, this.remainingBuffer.byteLength)
    }
    let offset = 0
    while (offset + HEADER_LENGTH <= combinedBuffer.byteLength) {
      // code is 1 byte long - it identifies the message type
      const code = combinedBuffer[offset]

      // length is 1 Uint32BE - it is the length of the message EXCLUDING the code
      const length = combinedBuffer.readUInt32BE(offset + CODE_LENGTH)

      const fullMessageLength = CODE_LENGTH + length

      if (fullMessageLength + offset <= combinedBuffer.byteLength) {
        const message = this.handlePacket(offset + HEADER_LENGTH, code, length, combinedBuffer)
        callback(message)
        offset += fullMessageLength
      } else {
        break
      }
    }

    if (offset === combinedBuffer.byteLength) {
      this.remainingBuffer = emptyBuffer
    } else {
      this.remainingBuffer = combinedBuffer.slice(offset)
    }
  }

  private handlePacket(offset: number, code: number, length: number, bytes: Buffer): BackendMessage {
    switch (code) {
      case MessageCodes.BindComplete:
        return bindComplete
      case MessageCodes.ParseComplete:
        return parseComplete
      case MessageCodes.CloseComplete:
        return closeComplete
      case MessageCodes.NoData:
        return noData
      case MessageCodes.PortalSuspended:
        return portalSuspended
      case MessageCodes.CopyDone:
        return copyDone
      case MessageCodes.ReplicationStart:
        return replicationStart
      case MessageCodes.EmptyQuery:
        return emptyQuery
      case MessageCodes.DataRow:
        return this.parseDataRowMessage(offset, length, bytes)<commit_msg>fix: major performance issues with bytea performance #2240<commit_after>import { TransformOptions } from 'stream'
import {
  Mode,
  bindComplete,
  parseComplete,
  closeComplete,
  noData,
  portalSuspended,
  copyDone,
  replicationStart,
  emptyQuery,
  ReadyForQueryMessage,
  CommandCompleteMessage,
  CopyDataMessage,
  CopyResponse,
  NotificationResponseMessage,
  RowDescriptionMessage,
  Field,
  DataRowMessage,
  ParameterStatusMessage,
  BackendKeyDataMessage,
  DatabaseError,
  BackendMessage,
  MessageName,
  AuthenticationMD5Password,
  NoticeMessage,
} from './messages'
import { BufferReader } from './buffer-reader'
import assert from 'assert'

// every message is prefixed with a single bye
const CODE_LENGTH = 1
// every message has an int32 length which includes itself but does
// NOT include the code in the length
const LEN_LENGTH = 4

const HEADER_LENGTH = CODE_LENGTH + LEN_LENGTH

export type Packet = {
  code: number
  packet: Buffer
}

const emptyBuffer = Buffer.allocUnsafe(0)

type StreamOptions = TransformOptions & {
  mode: Mode
}

const enum MessageCodes {
  DataRow = 0x44, // D
  ParseComplete = 0x31, // 1
  BindComplete = 0x32, // 2
  CloseComplete = 0x33, // 3
  CommandComplete = 0x43, // C
  ReadyForQuery = 0x5a, // Z
  NoData = 0x6e, // n
  NotificationResponse = 0x41, // A
  AuthenticationResponse = 0x52, // R
  ParameterStatus = 0x53, // S
  BackendKeyData = 0x4b, // K
  ErrorMessage = 0x45, // E
  NoticeMessage = 0x4e, // N
  RowDescriptionMessage = 0x54, // T
  PortalSuspended = 0x73, // s
  ReplicationStart = 0x57, // W
  EmptyQuery = 0x49, // I
  CopyIn = 0x47, // G
  CopyOut = 0x48, // H
  CopyDone = 0x63, // c
  CopyData = 0x64, // d
}

export type MessageCallback = (msg: BackendMessage) => void

export class Parser {
  private remainingBuffer: Buffer = emptyBuffer
  private remainingBufferLength: number = 0
  private remainingBufferOffset: number = 0
  private reader = new BufferReader()
  private mode: Mode

  constructor(opts?: StreamOptions) {
    if (opts?.mode === 'binary') {
      throw new Error('Binary mode not supported yet')
    }
    this.mode = opts?.mode || 'text'
  }

  public parse(buffer: Buffer, callback: MessageCallback) {
    let combinedBuffer = buffer
    let combinedBufferOffset = 0
    let combinedBufferLength = buffer.byteLength
    const newRealLength = this.remainingBufferLength + combinedBufferLength
    if (this.remainingBufferLength) {
      const newLength = newRealLength + this.remainingBufferOffset
      if (newLength > this.remainingBuffer.byteLength) {
        let newBufferLength = this.remainingBufferLength * 2
        while (newRealLength >= newBufferLength) {
          newBufferLength *= 2
        }
        const newBuffer = Buffer.allocUnsafe(newBufferLength)
        this.remainingBuffer.copy(
          newBuffer,
          0,
          this.remainingBufferOffset,
          this.remainingBufferOffset + this.remainingBufferLength
        )
        this.remainingBuffer = newBuffer
        this.remainingBufferOffset = 0
      }
      buffer.copy(this.remainingBuffer, this.remainingBufferOffset + this.remainingBufferLength)
      combinedBuffer = this.remainingBuffer
      combinedBufferLength = newRealLength
      combinedBufferOffset = this.remainingBufferOffset
    }
    let offset = combinedBufferOffset
    while (offset + HEADER_LENGTH <= combinedBufferLength) {
      // code is 1 byte long - it identifies the message type
      const code = combinedBuffer[offset]

      // length is 1 Uint32BE - it is the length of the message EXCLUDING the code
      const length = combinedBuffer.readUInt32BE(offset + CODE_LENGTH)

      const fullMessageLength = CODE_LENGTH + length

      if (fullMessageLength + offset <= combinedBufferLength) {
        const message = this.handlePacket(offset + HEADER_LENGTH, code, length, combinedBuffer)
        callback(message)
        offset += fullMessageLength
      } else {
        break
      }
    }

    if (offset === combinedBufferLength) {
      this.remainingBuffer = emptyBuffer
      this.remainingBufferLength = 0
      this.remainingBufferOffset = 0
    } else {
      this.remainingBuffer = combinedBuffer
      this.remainingBufferLength = combinedBufferLength - offset
      this.remainingBufferOffset += offset
    }
  }

  private handlePacket(offset: number, code: number, length: number, bytes: Buffer): BackendMessage {
    switch (code) {
      case MessageCodes.BindComplete:
        return bindComplete
      case MessageCodes.ParseComplete:
        return parseComplete
      case MessageCodes.CloseComplete:
        return closeComplete
      case MessageCodes.NoData:
        return noData
      case MessageCodes.PortalSuspended:
        return portalSuspended
      case MessageCodes.CopyDone:
        return copyDone
      case MessageCodes.ReplicationStart:
        return replicationStart
      case MessageCodes.EmptyQuery:
        return emptyQuery
      case MessageCodes.DataRow:
        return this.parseDataRowMessage(offset, length, bytes)
{'fraction_non_alphanumeric': 0.07904085257548846, 'fraction_numerical': 0.017984014209591476, 'mean_word_length': 3.119341563786008, 'pattern_counts': {'":': 0, '<': 7, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '801177', 'n_tokens_mistral': 2691, 'n_tokens_neox': 2552, 'n_words': 711}