<commit_before>#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>
#include <libpq-fe.h>
#include <log.h>
#include <mem.h>
#include <util.h>
#include <mstring.h>
#include <time.h>
#include <sys/time.h>
#include "database.h"

extern char *pg_conninfo;
extern char *pg_user;

typedef struct _psqlDB {
  PGconn *conn;
  int sender_cnt;
  time_t last_commit;
} PsqlDB;

typedef struct _pqlTable {
  MString *insert_stmt; /* Named statement for inserting into this table */
} PsqlTable;

static int sql_stmt(PsqlDB* self, const char* stmt);
char* psql_get_key_value (Database *database, const char *table,
                          const char *key_column, const char *value_column,
                          const char *key);
int psql_set_key_value (Database *database, const char *table,
                        const char *key_column, const char *value_column,
                        const char *key, const char *value);

static int
begin_transaction (PsqlDB *db)
{
  const char *sql = "BEGIN TRANSACTION;";
  return sql_stmt (db, sql);
}

static int
end_transaction (PsqlDB *db)
{
  const char *sql = "END TRANSACTION";
  return sql_stmt (db, sql);
}

static int
reopen_transaction (PsqlDB *db)
{
  if (end_transaction (db) == -1) return -1;
  if (begin_transaction (db) == -1) return -1;
  return 0;
}

static MString*
psql_make_sql_insert (DbTable* table)
{
  int n = 0;
  int max = table->schema->nfields;

  if (max <= 0) {
    logerror ("psql: Trying to insert 0 values into table '%s'\n", table->schema->name);
    goto fail_exit;
  }

  MString* mstr = mstring_create ();

  if (mstr == NULL) {
    logerror("psql: Failed to create managed string for preparing SQL INSERT statement\n");
    goto fail_exit;
  }

  /* Build SQL "INSERT INTO" statement */
  n += mstring_set (mstr, "INSERT INTO \"");
  n += mstring_cat (mstr, table->schema->name);
  n += mstring_cat (mstr, "\" VALUES ($1, $2, $3, $4"); /* metadata columns */
  while (max-- > 0)
    mstring_sprintf (mstr, ", $%d", 4 + table->schema->nfields - max);
  mstring_cat (mstr, ");");

  if (n != 0) goto fail_exit;
  return mstr;

 fail_exit:
  if (mstr) mstring_delete (mstr);
  return NULL;
}

static struct {
  const char *name;
  const char *sql;
} meta_tables [] = {
  { .name = "_experiment_metadata",
    .sql = "CREATE TABLE _experiment_metadata (key TEXT PRIMARY KEY, value TEXT);" },
  { .name = "_senders",
    .sql = "CREATE TABLE _senders (name TEXT PRIMARY KEY, id INTEGER UNIQUE);" },
};


static int
psql_table_create_meta (Database *db, const char *name)
{
  PsqlDB *self = (PsqlDB*)db->handle;
  size_t i = 0;
  for (i = 0; i < LENGTH (meta_tables); i++)
    if (strcmp (meta_tables[i].name, name) == 0)
      return sql_stmt (self, meta_tables[i].sql);
  return -1;
}

/**
 * @brief Release the psql database.
 *
 * This function closes the connection to the database and frees all
 * of the allocated memory associated with the database.
 *
 * @param db the database that contains the psql database to release.
 */
void
psql_release(Database* db)
{
  PsqlDB* self = (PsqlDB*)db->handle;
  end_transaction (self);
  PQfinish(self->conn);
  // TODO: Release table specific data

  xfree(self);
  db->handle = NULL;
}

static char*
psql_get_sender_id (Database *database, const char *name)
{
  return psql_get_key_value (database, "_senders", "name", "id", name);
}

static int
psql_set_sender_id (Database *database, const char *name, int id)
{
  MString *mstr = mstring_create();
  mstring_sprintf (mstr, "%d", id);
  int ret = psql_set_key_value (database, "_senders", "name", "id", name, mstring_buf (mstr));
  mstring_delete (mstr);
  return ret;
}

/**
 * @brief Add sender a new sender to the database, returning its
 * index.
 *
 * If a sender with the given id already exists, its pre-existing
 * index is returned.  Otherwise, a new sender is added to the table
 * with a new sender id, unique to this experiment.
 *
 * @param db the database that contains the experiment db
 * @param sender_id the sender ID
 * @return the index of the sender
 */
static int
psql_add_sender_id(Database *db, char *sender_id)
{
  PsqlDB *self = (PsqlDB*)db->handle;
  int index = -1;
  char *id_str = psql_get_sender_id (db, sender_id);

  if (id_str) {
    index = atoi (id_str);
    xfree (id_str);

  } else {
    PGresult *res = PQexec (self->conn, "SELECT MAX(id) FROM _senders;");

    if (PQresultStatus (res) != PGRES_TUPLES_OK) {
      logwarn("psql:%s: Failed to get maximum sender id from database: %s; starting at 0\n",
          db->name, PQerrorMessage (self->conn));
      PQclear (res);
      return -1;
    }
    int rows = PQntuples (res);
    if (rows == 0) {
      logerror("psql:%s: Failed to get maximum sender id from database: empty result\n",
            db->name);
      PQclear (res);
      return -1;
    }
    index = atoi (PQgetvalue (res, 0, 0)) + 1;
    PQclear (res);
    psql_set_sender_id (db, sender_id, index);

  }

  return index;
}

char*
psql_get_metadata (Database *db, const char *key)
{
  return psql_get_key_value (db, "_experiment_metadata", "key", "value", key);
}

int
psql_set_metadata (Database *db, const char *key, const char *value)
{
  return psql_set_key_value (db, "_experiment_metadata", "key", "value", key, value);
}

/**
 *  @brief Do a key-value style select on a database table.
 *
 *  This function does a key lookup on a database table that is set up
 *  in key-value style.  The table can have more than two columns, but
 *  this function SELECT's two of them and returns the value of the
 *  value column.  It checks to make sure that the key returned is the
 *  one requested, then returns its corresponding value.
 *
 *  This function makes a lot of assumptions about the database and
 *  the table:
 *<commit_msg>psql: start sender ID at 0 if not found<commit_after>#include <string.h>
#include <stdio.h>
#include <stdlib.h>
#include <inttypes.h>
#include <libpq-fe.h>
#include <log.h>
#include <mem.h>
#include <util.h>
#include <mstring.h>
#include <time.h>
#include <sys/time.h>
#include "database.h"

extern char *pg_conninfo;
extern char *pg_user;

typedef struct _psqlDB {
  PGconn *conn;
  int sender_cnt;
  time_t last_commit;
} PsqlDB;

typedef struct _pqlTable {
  MString *insert_stmt; /* Named statement for inserting into this table */
} PsqlTable;

static int sql_stmt(PsqlDB* self, const char* stmt);
char* psql_get_key_value (Database *database, const char *table,
                          const char *key_column, const char *value_column,
                          const char *key);
int psql_set_key_value (Database *database, const char *table,
                        const char *key_column, const char *value_column,
                        const char *key, const char *value);

static int
begin_transaction (PsqlDB *db)
{
  const char *sql = "BEGIN TRANSACTION;";
  return sql_stmt (db, sql);
}

static int
end_transaction (PsqlDB *db)
{
  const char *sql = "END TRANSACTION";
  return sql_stmt (db, sql);
}

static int
reopen_transaction (PsqlDB *db)
{
  if (end_transaction (db) == -1) return -1;
  if (begin_transaction (db) == -1) return -1;
  return 0;
}

static MString*
psql_make_sql_insert (DbTable* table)
{
  int n = 0;
  int max = table->schema->nfields;

  if (max <= 0) {
    logerror ("psql: Trying to insert 0 values into table '%s'\n", table->schema->name);
    goto fail_exit;
  }

  MString* mstr = mstring_create ();

  if (mstr == NULL) {
    logerror("psql: Failed to create managed string for preparing SQL INSERT statement\n");
    goto fail_exit;
  }

  /* Build SQL "INSERT INTO" statement */
  n += mstring_set (mstr, "INSERT INTO \"");
  n += mstring_cat (mstr, table->schema->name);
  n += mstring_cat (mstr, "\" VALUES ($1, $2, $3, $4"); /* metadata columns */
  while (max-- > 0)
    mstring_sprintf (mstr, ", $%d", 4 + table->schema->nfields - max);
  mstring_cat (mstr, ");");

  if (n != 0) goto fail_exit;
  return mstr;

 fail_exit:
  if (mstr) mstring_delete (mstr);
  return NULL;
}

static struct {
  const char *name;
  const char *sql;
} meta_tables [] = {
  { .name = "_experiment_metadata",
    .sql = "CREATE TABLE _experiment_metadata (key TEXT PRIMARY KEY, value TEXT);" },
  { .name = "_senders",
    .sql = "CREATE TABLE _senders (name TEXT PRIMARY KEY, id INTEGER UNIQUE);" },
};


static int
psql_table_create_meta (Database *db, const char *name)
{
  PsqlDB *self = (PsqlDB*)db->handle;
  size_t i = 0;
  for (i = 0; i < LENGTH (meta_tables); i++)
    if (strcmp (meta_tables[i].name, name) == 0)
      return sql_stmt (self, meta_tables[i].sql);
  return -1;
}

/**
 * @brief Release the psql database.
 *
 * This function closes the connection to the database and frees all
 * of the allocated memory associated with the database.
 *
 * @param db the database that contains the psql database to release.
 */
void
psql_release(Database* db)
{
  PsqlDB* self = (PsqlDB*)db->handle;
  end_transaction (self);
  PQfinish(self->conn);
  // TODO: Release table specific data

  xfree(self);
  db->handle = NULL;
}

static char*
psql_get_sender_id (Database *database, const char *name)
{
  return psql_get_key_value (database, "_senders", "name", "id", name);
}

static int
psql_set_sender_id (Database *database, const char *name, int id)
{
  MString *mstr = mstring_create();
  mstring_sprintf (mstr, "%d", id);
  int ret = psql_set_key_value (database, "_senders", "name", "id", name, mstring_buf (mstr));
  mstring_delete (mstr);
  return ret;
}

/**
 * @brief Add sender a new sender to the database, returning its
 * index.
 *
 * If a sender with the given id already exists, its pre-existing
 * index is returned.  Otherwise, a new sender is added to the table
 * with a new sender id, unique to this experiment.
 *
 * @param db the database that contains the experiment db
 * @param sender_id the sender ID
 * @return the index of the sender
 */
static int
psql_add_sender_id(Database *db, char *sender_id)
{
  PsqlDB *self = (PsqlDB*)db->handle;
  int index = -1;
  char *id_str = psql_get_sender_id (db, sender_id);

  if (id_str) {
    index = atoi (id_str);
    xfree (id_str);

  } else {
    PGresult *res = PQexec (self->conn, "SELECT MAX(id) FROM _senders;");

    if (PQresultStatus (res) != PGRES_TUPLES_OK) {
      logwarn("psql:%s: Failed to get maximum sender id from database: %s; starting at 0\n",
          db->name, PQerrorMessage (self->conn));
      PQclear (res);
      index = 0;
    } else {
      int rows = PQntuples (res);
      if (rows == 0) {
        logwarn("psql:%s: Failed to get maximum sender id from database: empty result; starting at 0\n",
            db->name);
        PQclear (res);
        index = 0;
      } else {
        index = atoi (PQgetvalue (res, 0, 0)) + 1;
        PQclear (res);
      }
    }

    psql_set_sender_id (db, sender_id, index);

  }

  return index;
}

char*
psql_get_metadata (Database *db, const char *key)
{
  return psql_get_key_value (db, "_experiment_metadata", "key", "value", key);
}

int
psql_set_metadata (Database *db, const char *key, const char *value)
{
  return psql_set_key_value (db, "_experiment_metadata", "key", "value", key, value);
}

/**
 *  @brief Do a key-value style select on a database table.
 *
 *  This function does a key lookup on a database table that is set up
 *  in key-value style.  The table can have more than two columns, but
 *  this function SELECT's two of them and returns the value of the
 *  value column.  It checks to make sure that the key returned is the
 *  one requested, then returns its corresponding value.
 *
 *  This function makes a lot of assumptions about the database and
 *  the table:
 *
{'fraction_non_alphanumeric': 0.12550364337762537, 'fraction_numerical': 0.0048006858122588945, 'mean_word_length': 3.377485928705441, 'pattern_counts': {'":': 0, '<': 29, '<?xml version=': 0, '>': 61, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6472821', 'n_tokens_mistral': 4102, 'n_tokens_neox': 3972, 'n_words': 1359}