Skip to content

Commit

Permalink
expose raw CopyData command (#1077)
Browse files Browse the repository at this point in the history
  • Loading branch information
otan committed Apr 12, 2022
1 parent 1ef134d commit b3b8332
Showing 1 changed file with 33 additions and 0 deletions.
33 changes: 33 additions & 0 deletions copy.go
@@ -1,6 +1,7 @@
package pq

import (
"context"
"database/sql/driver"
"encoding/binary"
"errors"
Expand Down Expand Up @@ -273,6 +274,38 @@ func (ci *copyin) Exec(v []driver.Value) (r driver.Result, err error) {
return driver.RowsAffected(0), nil
}

// CopyData executes a raw CopyData command using the PostgreSQL Frontend/Backend
// protocol. Use Exec(nil) to finish the command.
func (ci *copyin) CopyData(ctx context.Context, line string) (r driver.Result, err error) {
if ci.closed {
return nil, errCopyInClosed
}

if finish := ci.cn.watchCancel(ctx); finish != nil {
defer finish()
}

if err := ci.getBad(); err != nil {
return nil, err
}
defer ci.cn.errRecover(&err)

if err := ci.err(); err != nil {
return nil, err
}

ci.buffer = append(ci.buffer, []byte(line)...)
ci.buffer = append(ci.buffer, '\n')

if len(ci.buffer) > ciBufferFlushSize {
ci.flush(ci.buffer)
// reset buffer, keep bytes for message identifier and length
ci.buffer = ci.buffer[:5]
}

return driver.RowsAffected(0), nil
}

func (ci *copyin) Close() (err error) {
if ci.closed { // Don't do anything, we're already closed
return nil
Expand Down

0 comments on commit b3b8332

Please sign in to comment.