sqlCopy.c

Go to the documentation of this file.

/*
 SPDX-FileCopyrightText: © 2011 Hewlett-Packard Development Company, L.P.
 
 SPDX-License-Identifier: LGPL-2.1-only
*/
 
#include "sqlCopy.h"
 
psqlCopy_t fo_sqlCopyCreate(PGconn* pGconn, char* TableName, int BufSize, int NumColumns, ...)
{
  psqlCopy_t pCopy;
  va_list ColumnNameArg;
  int ColIdx;
  int ColStrLength;
  char* ColStr;
 
  va_start(ColumnNameArg, NumColumns);
 
  /* Allocate the structure */
  pCopy = malloc(sizeof(sqlCopy_t));
  if (!pCopy)
  {
    va_end(ColumnNameArg);
    ERROR_RETURN("sqlCopy malloc")
  }
 
  /* Allocate storage for the data buffer */
  if (BufSize < 1) BufSize = 1;
  pCopy->DataBuf = calloc(BufSize, sizeof(char));
 
  /* Save TableName */
  pCopy->TableName =g_strdup(TableName);
 
  /* check for malloc failures */
  if ((!pCopy->DataBuf) || (!pCopy->TableName))
  {
    fo_sqlCopyDestroy(pCopy, 0);
    va_end(ColumnNameArg);
    ERROR_RETURN("sqlCopyCreate")
  }
 
  /* Save the DB connection */
  pCopy->pGconn = pGconn;
 
  /* Save the data buffer size  */
  pCopy->BufSize = BufSize;
 
  /* Data buffer is empty */
  pCopy->DataIdx = 0;
 
  /* Build the column name string pCopy->ColumnNames  */
  ColStrLength = 0;
  pCopy->ColumnNames[0] = 0;
  for (ColIdx = 0; ColIdx < NumColumns; ColIdx++)
  {
    ColStr = va_arg(ColumnNameArg, char *);
    ColStrLength += strlen(ColStr) + 1;  /* extra 1 for the comma */
    if (ColStrLength < sizeof(pCopy->ColumnNames))
    {
      if (ColIdx != 0) strncat(pCopy->ColumnNames, ",", 2);
      strncat(pCopy->ColumnNames, ColStr, ColStrLength);
    }
    else
    {
      fo_sqlCopyDestroy(pCopy, 0);
      va_end(ColumnNameArg);
      ERROR_RETURN("pCopy->ColumnNames size too small")
    }
  }
  va_end(ColumnNameArg);
  return (pCopy);
}  /* End fo_sqlCopyCreate()  */
 
 
#ifdef DEBUG
int tmp_printhex(char * str)
{
  while(*str) printf("%02x", *str++);
  return(0);
}
#endif
 
#define growby  128  
 
int fo_sqlCopyAdd(psqlCopy_t pCopy, char* DataRow)
{
  int NewRowLen;
  int rncount = 0;
  char* dptr = DataRow;
  char* NewRow = 0, * nptr;
 
  /* As of Postgresql 8.4, COPY will not accept embedded literal carriage returns
   * or line feeds.  Use "\r" and "\n" instead.
   * Count how many we need to get rid of (and make space for).
   */
  while (*dptr)
  {
    if (((*dptr == '\n') || (*dptr == '\r')) && (*(dptr + 1))) rncount++;
    dptr++;
  }
 
  /* Substitute any literal '\n' or '\r' for string "\n", "\r"
   * (except for trailing \n which is required)
   */
  if (rncount)
  {
    NewRowLen = strlen(DataRow) + rncount;
    NewRow = malloc(NewRowLen + 1);  // plus 1 for potential required \n at end
    if (!NewRow)
    ERROR_RETURN("fo_sqlCopyAdd: out of memory");
    nptr = NewRow;
    dptr = DataRow;
    while (*dptr && *(dptr + 1))
    {
      if (*dptr == '\n')
      {
        *nptr++ = '\\';
        *nptr = 'n';
      }
      else if (*dptr == '\r')
      {
        *nptr++ = '\\';
        *nptr = 'r';
      }
      else
        *nptr = *dptr;
      ++dptr;
      ++nptr;
    }
    *nptr = 0;  // null terminator
    DataRow = NewRow;
  }
 
  /* Does new record fit in DataBuf?  */
  NewRowLen = strlen(DataRow);
  if ((pCopy->BufSize - pCopy->DataIdx) < (NewRowLen + 1))
  {
    /* if DataIdx is zero, then DataBuf isn't big enough to hold
     * this record.  In this case make DataBuf larger.
     */
    if (pCopy->DataIdx == 0)
    {
      pCopy->DataBuf = realloc(pCopy->DataBuf, NewRowLen + growby);
      if (!pCopy->DataBuf)
      ERROR_RETURN("fo_sqlCopyAdd: Realloc for DataBuf failed");
      pCopy->BufSize = NewRowLen + growby;
    }
    else
    {
      /* Execute a copy to make room in DataBuf */
      fo_sqlCopyExecute(pCopy);
    }
  }
 
  /* copy in DataRow */
  strcpy(pCopy->DataBuf + pCopy->DataIdx, DataRow);
  pCopy->DataIdx += NewRowLen;
 
  /* If the DataRow was missing a terminating newline, add one */
  if (DataRow[NewRowLen - 1] != '\n')
  {
    pCopy->DataBuf[pCopy->DataIdx++] = '\n';
    pCopy->DataBuf[pCopy->DataIdx] = 0;  // new null terminator
  }
 
  if (NewRow) free(NewRow);
  return (1);
}
 
int fo_sqlCopyExecute(psqlCopy_t pCopy)
{
  char copystmt[2048];
  PGresult* result;
 
  /* check pCopy */
  if (!pCopy)
  ERROR_RETURN("Null pCopy");
  if (pCopy->DataIdx == 0) return (1);  /* nothing to copy */
 
  /* Start the Copy command */
  sprintf(copystmt, "COPY %s(%s) from stdin",
    pCopy->TableName,
    pCopy->ColumnNames);
  result = PQexec(pCopy->pGconn, copystmt);
  if (PGRES_COPY_IN == PQresultStatus(result))
  {
    PQclear(result);
    if (PQputCopyData(pCopy->pGconn, pCopy->DataBuf, pCopy->DataIdx) != 1)
    ERROR_RETURN(PQresultErrorMessage(result))
  }
  else if (fo_checkPQresult(pCopy->pGconn, result, copystmt, __FILE__, __LINE__)) return 0;
 
 
  /* End copy  */
  if (PQputCopyEnd(pCopy->pGconn, NULL) == 1)
  {
    result = PQgetResult(pCopy->pGconn);
    if (fo_checkPQcommand(pCopy->pGconn, result, "copy end", __FILE__, __LINE__)) return 0;
  }
  PQclear(result);
 
  /* reset DataBuf */
  pCopy->DataIdx = 0;
 
  return (1);
}
 
 
void fo_sqlCopyDestroy(psqlCopy_t pCopy, int ExecuteFlag)
{
  if (!pCopy) return;
  if (ExecuteFlag) fo_sqlCopyExecute(pCopy);
  if (pCopy->TableName) free(pCopy->TableName);
  if (pCopy->DataBuf) free(pCopy->DataBuf);
  free(pCopy);
}
 
 
void fo_sqlCopyPrint(psqlCopy_t pCopy, int PrintBytes)
{
  int idx;
 
  printf("========== fo_sqlCopyPrint  Start  ================\n");
  printf("pCopy: %lx, TableName: %s, BufSize: %d, DataIdx: %d\n",
    (long) pCopy, pCopy->TableName, pCopy->BufSize, pCopy->DataIdx);
  printf("       ColumnNames: %s\n", pCopy->ColumnNames);
 
  if (PrintBytes == 0) PrintBytes = pCopy->DataIdx;
  for (idx = 0; idx < PrintBytes; idx++) putchar(pCopy->DataBuf[idx]);
 
  printf("========== fo_sqlCopyPrint  End  ================");
}