diff --git a/src/backend/commands/copy.c b/src/backend/commands/copy.c index 3e199bd..8a81bdd 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, char *dest, int nbytes); static Datum CopyReadBinaryAttribute(CopyState cstate, FmgrInfo *flinfo, Oid typioparam, int32 typmod, bool *isnull); @@ -388,9 +389,7 @@ static void CopySendEndOfRow(CopyState cstate); static int CopyGetData(CopyState cstate, void *databuf, int minread, int maxread); static void CopySendInt32(CopyState cstate, int32 val); -static bool CopyGetInt32(CopyState cstate, int32 *val); static void CopySendInt16(CopyState cstate, int16 val); -static bool CopyGetInt16(CopyState cstate, int16 *val); /* @@ -730,25 +729,6 @@ CopySendInt32(CopyState cstate, int32 val) } /* - * CopyGetInt32 reads an int32 that appears in network byte order - * - * Returns true if OK, false if EOF - */ -static bool -CopyGetInt32(CopyState cstate, int32 *val) -{ - uint32 buf; - - if (CopyGetData(cstate, &buf, sizeof(buf), sizeof(buf)) != sizeof(buf)) - { - *val = 0; /* suppress compiler warning */ - return false; - } - *val = (int32) pg_ntoh32(buf); - return true; -} - -/* * CopySendInt16 sends an int16 in network byte order */ static void @@ -760,23 +740,6 @@ CopySendInt16(CopyState cstate, int16 val) CopySendData(cstate, &buf, sizeof(buf)); } -/* - * CopyGetInt16 reads an int16 that appears in network byte order - */ -static bool -CopyGetInt16(CopyState cstate, int16 *val) -{ - uint16 buf; - - if (CopyGetData(cstate, &buf, sizeof(buf), sizeof(buf)) != sizeof(buf)) - { - *val = 0; /* suppress compiler warning */ - return false; - } - *val = (int16) pg_ntoh16(buf); - return true; -} - /* * CopyLoadRawBuf loads some more data into raw_buf @@ -813,6 +776,61 @@ 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, char *dest, int nbytes) +{ + 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 @@ -3514,16 +3532,18 @@ BeginCopyFrom(ParseState *pstate, int32 tmp; /* Signature */ - if (CopyGetData(cstate, readSig, 11, 11) != 11 || + if (CopyReadFromRawBuf(cstate, readSig, 11) != 11 || memcmp(readSig, BinarySignature, 11) != 0) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("COPY file signature not recognized"))); /* Flags field */ - if (!CopyGetInt32(cstate, &tmp)) + if (CopyReadFromRawBuf(cstate, (char *) &tmp, sizeof(tmp)) != + sizeof(tmp)) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("invalid COPY file header (missing flags)"))); + tmp = (int32) pg_ntoh32(tmp); if ((tmp & (1 << 16)) != 0) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), @@ -3534,15 +3554,15 @@ BeginCopyFrom(ParseState *pstate, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("unrecognized critical flags in COPY file header"))); /* Header extension length */ - if (!CopyGetInt32(cstate, &tmp) || - tmp < 0) + if (CopyReadFromRawBuf(cstate, (char *) &tmp, sizeof(tmp)) != + sizeof(tmp) || (tmp = (int32) pg_ntoh32(tmp)) < 0) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("invalid COPY file header (missing length)"))); /* Skip extension header, if present */ while (tmp-- > 0) { - if (CopyGetData(cstate, readSig, 1, 1) != 1) + if (CopyReadFromRawBuf(cstate, readSig, 1) != 1) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("invalid COPY file header (wrong length)"))); @@ -3735,12 +3755,14 @@ NextCopyFrom(CopyState cstate, ExprContext *econtext, cstate->cur_lineno++; - if (!CopyGetInt16(cstate, &fld_count)) + if (CopyReadFromRawBuf(cstate, (char *) &fld_count, + sizeof(fld_count)) != sizeof(fld_count)) { /* EOF detected (end of file, or protocol-level EOF) */ return false; } + fld_count = (int16) pg_ntoh16(fld_count); if (fld_count == -1) { /* @@ -3758,7 +3780,7 @@ NextCopyFrom(CopyState cstate, ExprContext *econtext, char dummy; if (cstate->copy_dest != COPY_OLD_FE && - CopyGetData(cstate, &dummy, 1, 1) > 0) + CopyReadFromRawBuf(cstate, &dummy, 1) > 0) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("received copy data after EOF marker"))); @@ -4713,10 +4735,12 @@ CopyReadBinaryAttribute(CopyState cstate, FmgrInfo *flinfo, int32 fld_size; Datum result; - if (!CopyGetInt32(cstate, &fld_size)) + if (CopyReadFromRawBuf(cstate, (char *) &fld_size, sizeof(fld_size)) != + sizeof(fld_size)) ereport(ERROR, (errcode(ERRCODE_BAD_COPY_FILE_FORMAT), errmsg("unexpected EOF in COPY data"))); + fld_size = (int32) pg_ntoh32(fld_size); if (fld_size == -1) { *isnull = true; @@ -4731,8 +4755,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")));