diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c index 3e199bd..c5be17e 100644 --- a/src/backend/commands/copy.c +++ b/src/backend/commands/copy.c @@ -367,6 +367,7 @@ static bool CopyReadLine(CopyState cstate); static bool CopyReadLineText(CopyState cstate); static int CopyReadAttributesText(CopyState cstate); static int CopyReadAttributesCSV(CopyState cstate); +static int CopyReadFromRawBuf(CopyState cstate, void *saveTo, int nbytes); static Datum CopyReadBinaryAttribute(CopyState cstate, FmgrInfo *flinfo, Oid typioparam, int32 typmod, bool *isnull); @@ -739,7 +740,7 @@ CopyGetInt32(CopyState cstate, int32 *val) { uint32 buf; - if (CopyGetData(cstate, &buf, sizeof(buf), sizeof(buf)) != sizeof(buf)) + if (CopyReadFromRawBuf(cstate, &buf, sizeof(buf)) != sizeof(buf)) { *val = 0; /* suppress compiler warning */ return false; @@ -768,7 +769,7 @@ CopyGetInt16(CopyState cstate, int16 *val) { uint16 buf; - if (CopyGetData(cstate, &buf, sizeof(buf), sizeof(buf)) != sizeof(buf)) + if (CopyReadFromRawBuf(cstate, &buf, sizeof(buf)) != sizeof(buf)) { *val = 0; /* suppress compiler warning */ return false; @@ -813,6 +814,62 @@ CopyLoadRawBuf(CopyState cstate) return (inbytes > 0); } +/* + * CopyReadFromRawBuf + * Reads 'nbytes' bytes from cstate->copy_file via cstate->raw_buf and + * writes then to 'saveTo' + * + * Useful when reading binary data from the file. + */ +#define DRAIN_COPY_RAW_BUF(cstate, dest, nbytes)\ + do {\ + memcpy((dest), (cstate)->raw_buf + (cstate)->raw_buf_index, (nbytes));\ + (cstate)->raw_buf_index += (nbytes);\ + } while(0) + +#define BUF_BYTES (cstate->raw_buf_len - cstate->raw_buf_index) + +static int +CopyReadFromRawBuf(CopyState cstate, void *saveTo, int nbytes) +{ + char *dest = saveTo; + int copied_bytes = 0; + + if (BUF_BYTES >= nbytes) + { + /* Enough bytes are present in the buffer. */ + DRAIN_COPY_RAW_BUF(cstate, dest, nbytes); + copied_bytes = nbytes; + } + else + { + /* + * Not enough bytes in the buffer, so must read from the file. Need + * the loop considering that 'nbytes' may be larger than the maximum + * bytes that the buffer can hold. + */ + do + { + int copy_bytes; + + /* + * This tries to read up to RAW_BUF_SIZE bytes into raw_buf, + * returning if no more were read. + */ + if (BUF_BYTES == 0 && !CopyLoadRawBuf(cstate)) + return copied_bytes; + + copy_bytes = Min(nbytes - copied_bytes, BUF_BYTES); + DRAIN_COPY_RAW_BUF(cstate, dest, copy_bytes); + dest += copy_bytes; + copied_bytes += copy_bytes; + } while (copied_bytes < nbytes); + } + + return copied_bytes; +} +#undef DRAIN_COPY_RAW_BUF +#undef BUF_BYTES /* * DoCopy executes the SQL COPY statement @@ -4731,8 +4788,8 @@ CopyReadBinaryAttribute(CopyState cstate, FmgrInfo *flinfo, resetStringInfo(&cstate->attribute_buf); enlargeStringInfo(&cstate->attribute_buf, fld_size); - if (CopyGetData(cstate, cstate->attribute_buf.data, - fld_size, fld_size) != fld_size) + if (CopyReadFromRawBuf(cstate, cstate->attribute_buf.data, + fld_size) != fld_size) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("unexpected EOF in COPY data")));