PgQuerySplitResult pg_query_split_with_scanner()

in ext/pg_query/pg_query_split.c [10:162]


PgQuerySplitResult pg_query_split_with_scanner(const char* input)
{
  MemoryContext ctx = NULL;
  PgQuerySplitResult result = {0};
  core_yyscan_t yyscanner;
  core_yy_extra_type yyextra;
  core_YYSTYPE yylval;
  YYLTYPE    yylloc;
  size_t curstmt = 0;
  bool keyword_before_terminator = false;
  int stmtstart = 0;
  bool is_keyword = false;
  size_t open_parens = 0;

  ctx = pg_query_enter_memory_context();

  MemoryContext parse_context = CurrentMemoryContext;

  char stderr_buffer[STDERR_BUFFER_LEN + 1] = {0};
#ifndef DEBUG
  int stderr_global;
  int stderr_pipe[2];
#endif

#ifndef DEBUG
  // Setup pipe for stderr redirection
  if (pipe(stderr_pipe) != 0) {
    PgQueryError* error = malloc(sizeof(PgQueryError));

    error->message = strdup("Failed to open pipe, too many open file descriptors")

    result.error = error;

    return result;
  }

  fcntl(stderr_pipe[0], F_SETFL, fcntl(stderr_pipe[0], F_GETFL) | O_NONBLOCK);

  // Redirect stderr to the pipe
  stderr_global = dup(STDERR_FILENO);
  dup2(stderr_pipe[1], STDERR_FILENO);
  close(stderr_pipe[1]);
#endif

  PG_TRY();
  {
    // Really this is stupid, we only run twice so we can pre-allocate the output array correctly
    yyscanner = scanner_init(input, &yyextra, &ScanKeywords, ScanKeywordTokens);
    while (true)
    {
      int tok = core_yylex(&yylval, &yylloc, yyscanner);
      switch (tok) {
      #define PG_KEYWORD(a,b,c) case b: is_keyword = true; break;
      #include "parser/kwlist.h"
      default: is_keyword = false;
      }
      if (is_keyword)
        keyword_before_terminator = true;
      else if (tok == '(')
        open_parens++;
      else if (tok == ')')
        open_parens--;
      else if (keyword_before_terminator && open_parens == 0 && (tok == ';' || tok == 0))
      {
        result.n_stmts++;
        keyword_before_terminator = false;
      }
      if (tok == 0) break;
    }
    scanner_finish(yyscanner);

    result.stmts = malloc(sizeof(PgQuerySplitStmt *) * result.n_stmts);

    // Now actually set the output values
    keyword_before_terminator = false;
    open_parens = 0;
    yyscanner = scanner_init(input, &yyextra, &ScanKeywords, ScanKeywordTokens);
    while (true)
    {
      int tok = core_yylex(&yylval, &yylloc, yyscanner);
      switch (tok) {
      #define PG_KEYWORD(a,b,c) case b: is_keyword = true; break;
      #include "parser/kwlist.h"
      default: is_keyword = false;
      }
      if (is_keyword)
        keyword_before_terminator = true;
      else if (tok == '(')
        open_parens++;
      else if (tok == ')')
        open_parens--;
      else if (keyword_before_terminator && open_parens == 0 && (tok == ';' || tok == 0))
      {
        // Add statement up to the current position
        result.stmts[curstmt] = malloc(sizeof(PgQuerySplitStmt));
        result.stmts[curstmt]->stmt_location = stmtstart;
        result.stmts[curstmt]->stmt_len = yylloc - stmtstart;

        stmtstart = yylloc + 1;
        keyword_before_terminator = false;

        curstmt++;
      }
      else if (open_parens == 0 && tok == ';') // Advance statement start in case we skip an empty statement
      {
        stmtstart = yylloc + 1;
      }

      if (tok == 0) break;
    }

    scanner_finish(yyscanner);

#ifndef DEBUG
    // Save stderr for result
    read(stderr_pipe[0], stderr_buffer, STDERR_BUFFER_LEN);
#endif

    result.stderr_buffer = strdup(stderr_buffer);
  }
  PG_CATCH();
  {
    ErrorData* error_data;
    PgQueryError* error;

    MemoryContextSwitchTo(parse_context);
    error_data = CopyErrorData();

    // Note: This is intentionally malloc so exiting the memory context doesn't free this
    error = malloc(sizeof(PgQueryError));
    error->message   = strdup(error_data->message);
    error->filename  = strdup(error_data->filename);
    error->funcname  = strdup(error_data->funcname);
    error->context   = NULL;
    error->lineno    = error_data->lineno;
    error->cursorpos = error_data->cursorpos;

    result.error = error;
    FlushErrorState();
  }
  PG_END_TRY();

#ifndef DEBUG
  // Restore stderr, close pipe
  dup2(stderr_global, STDERR_FILENO);
  close(stderr_pipe[0]);
  close(stderr_global);
#endif

  pg_query_exit_memory_context(ctx);

  return result;
}