sqlCopy.c

Go to the documentation of this file.

/**************************************************************
Copyright (C) 2011 Hewlett-Packard Development Company, L.P.

This library is free software; you can redistribute it and/or
modify it under the terms of the GNU Lesser General Public
License version 2.1 as published by the Free Software Foundation.

This library is distributed in the hope that it will be useful,
but WITHOUT ANY WARRANTY; without even the implied warranty of
MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the GNU
Lesser General Public License for more details.

You should have received a copy of the GNU Lesser General Public License
along with this library; if not, write to the Free Software Foundation, Inc.0
51 Franklin Street, Fifth Floor, Boston, MA  02110-1301  USA
**************************************************************/

#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, ",", 1);
      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  ================");
}