Not logged in

Artifact 5f2b1c090ed8d9bd204f7ac54782c244b8094183:


/* zdbs.c

   Copyright © 2011, 2012, 2013, 2016 Brandon Invergo <brandon@invergo.net>

   This file is part of zeptodb

   zeptodb is free software: you can redistribute it and/or modify
   it under the terms of the GNU General Public License as published by
   the Free Software Foundation, either version 3 of the License, or
   (at your option) any later version.

   zeptodb 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 General Public License for more details.

   You should have received a copy of the GNU General Public License
   along with zeptodb.  If not, see <http://www.gnu.org/licenses/>.
*/

#define _GNU_SOURCE

#include <config.h>
#include <string.h>
#include <stdio.h>
#include "error.h"
#include "argp.h"
#include <signal.h>
#include <stdbool.h>
#include <stdint.h>
#include "zdb.h"

const char *program_name = "zdbs";
const char *argp_program_version = PACKAGE_STRING;
const char *argp_program_bug_address = PACKAGE_BUGREPORT;
void *database;
static char doc[] =
  "zdbs -- a tool for storing records in a DBM database\v"
  "If an input file is provided, lines are read from it as new records, "
  "otherwise records are added from stdin. Records are expected in the "
  "format KEY|VALUE.";
static char args_doc[] = "DATABASE";

static struct argp_option options[] = {
  {"mmap-size",   (int)'m', "NUM", 0,
   "The size (in bytes) of the memory-mapped region to use (default=1024)", 0},
  {"cache-size", (int)'c', "NUM", 0,
   "The bucket cache size to use (default=128)", 0},
  {"block-size", (int)'b', "NUM", 0,
   "The block size to use (default=512)", 0},
  {"no-mmap", (int)'n', (const char *)NULL, 0,
   "Do not memory-map the database file", 0},
  {"no-lock", (int)'l', (const char *)NULL, 0,
   "Do not perform any locking of the database file", 0},
  {"verbose", (int)'v', (const char *)NULL, 0,
   "Print extra information.", 0},
  {"input", (int)'i', "FILE", 0,
   "Read new records from a file", 0},
  {"delim", (int)'d', "CHAR", 0,
   "Character used to separate keys from values (default '|')", 0},
  {"sync", 's', (const char *)NULL, 0,
   "Automatically synchronize all database operations to the disk", 0},
  {(const char *)NULL, 0, (const char *)NULL, 0, (const char *)NULL, 0}
};

struct arguments
{
  char *args[1];
  size_t mmap_size;
  size_t cache_size;
  size_t block_size;
  bool no_mmap;
  bool no_lock;
  bool verbose;
  char *input_file;
  const char *delim;
  bool sync;
};

static error_t
parse_opt (int key, char *arg, struct argp_state *state)
{
  struct arguments *arguments = state->input;
  unsigned long int size;
  switch (key)
    {
    case 'b':
      size = arg ? strtoul (arg, NULL, 10) : 512UL;
      if (size > SIZE_MAX)
        {
          arguments->cache_size = SIZE_MAX;
        }
      else
        {
          arguments->cache_size = (size_t)size;
        }
      break;
    case 'c':
      size = arg ? strtoul (arg, NULL, 10) : 128UL;
      if (size > SIZE_MAX)
        {
          arguments->cache_size = SIZE_MAX;
        }
      else
        {
          arguments->cache_size = (size_t)size;
        }
      break;
    case 'd':
      if (strlen (arg) > 1)
        argp_usage (state);
      arguments->delim = arg;
      break;
    case 'i':
      if (arguments->input_file)
        free (arguments->input_file);
      if (!arg)
        {
          arguments->input_file = NULL;
        }
      else
        {
          arguments->input_file =
            (char *) malloc (sizeof (char) * (strlen (arg)+1));
          if (!arguments->input_file)
            error (EXIT_FAILURE, errno,
                   "Failed to allocate memory");
          strncpy (arguments->input_file, arg, strlen (arg));
          arguments->input_file[strlen (arg)] = '\0';
        }
      break;
    case 'l':
      arguments->no_lock = true;
      break;
    case 'm':
      size = arg ? strtoul (arg, NULL, 10) : 131072UL;
      if (size > SIZE_MAX)
        {
          arguments->mmap_size = SIZE_MAX;
        }
      else
        {
          arguments->mmap_size = (size_t)size;
        }
      break;
    case 'n':
      arguments->no_mmap = true;
      break;
    case 's':
      arguments->sync = true;
      break;
    case 'v':
      arguments->verbose = true;
      break;
    case ARGP_KEY_ARG:
      if (state->arg_num >= 1)
        argp_usage (state);
      arguments->args[state->arg_num] = arg;
      break;
    case ARGP_KEY_END:
      if (state->arg_num == 0)
        argp_usage (state);
      break;
    default:
      return ARGP_ERR_UNKNOWN;
    }
  return 0;
}

static struct argp argp = {options, parse_opt, args_doc, doc,
                           (const struct argp_child *)NULL,
                           NULL, (const char *)NULL};


int
addrecords (void *db, FILE *input, const char *delim, bool verbose)
{
  size_t linelen = 512;
  char *lineptr = (char *) malloc (linelen);
  char *linecpy;
  char *key_str, *value_str;
  size_t line_size, key_size, value_size;
  long int count = 1;

  if (!lineptr)
    error (EXIT_FAILURE, errno,
           "Failed to allocate memory");

  /* add each line of the input to the database */
  while (getline (&lineptr, &linelen, input) != EOF)
    {
      if (!lineptr)
          error (EXIT_FAILURE, errno, "Failed to read line");
      line_size = strlen (lineptr);
      if (lineptr[line_size - 1] == '\n')
        line_size -= 1;
      if (line_size == 0)
        error (0, errno, "Empty record");
      linecpy = (char *) malloc (sizeof (char) * (line_size+1));
      if (!linecpy)
        {
          free (lineptr);
          error (EXIT_FAILURE, errno, "Failed to allocate memory");
        }
      strncpy (linecpy, lineptr, sizeof (char) * (line_size));
      linecpy[line_size] = '\0';
      key_str = strtok (linecpy, delim);
      if (!key_str)
        {
          error (0, errno, "Invalid record format (key%svalue)", delim);
          continue;
        }
      value_str = strtok (NULL, delim);
      if (!value_str)
        {
          error (0, errno, "Invalid record format (key%svalue)", delim);
          continue;
        }
      value_size = strlen (value_str);
      key_size = strlen (key_str);
      if (key_size == 0)
        {
          error (0, errno, "Empty key value");
          free (linecpy);
          continue;
        }
      if (value_size == 0)
        {
          error (0, errno, "Empty value");
          free (linecpy);
          continue;
        }
      if (verbose)
        printf ("%ld: ", count);
      if (zdb_store (db, key_str, key_size, value_str, value_size, verbose))
        error (0, errno, "Could not store record in database: %s", key_str);
      count++;
      free (linecpy);
    }
  free (lineptr);
  return (0);
}

void
termination_handler (int signum __attribute__ ((unused)))
{
  printf ("Interrupt caught, closing database\n");
  if (database && zdb_close (database, false))
    error (EXIT_FAILURE, errno, "Failed to close database");
  else
    exit (EXIT_SUCCESS);
}

int
main (int argc, char **argv)
{
  FILE *input;
  struct arguments arguments;
  int mode = ZDB_WRITER;
  error_t argp_result;
  int result;

  arguments.mmap_size = 1024;
  arguments.cache_size = 128;
  arguments.block_size = 512;
  arguments.no_mmap = false;
  arguments.no_lock = false;
  arguments.verbose = false;
  arguments.sync = false;
  arguments.input_file = NULL;
  arguments.delim = "|";

  argp_result = argp_parse (&argp, argc, argv, 0, 0, &arguments);
  if (argp_result)
    error (EXIT_FAILURE, argp_result, "Failed to parse arguments");

  if (signal (SIGINT, termination_handler) == SIG_IGN)
    signal (SIGINT, SIG_IGN);
  if (signal (SIGHUP, termination_handler) == SIG_IGN)
    signal (SIGHUP, SIG_IGN);
  if (signal (SIGTERM, termination_handler) == SIG_IGN)
    signal (SIGTERM, SIG_IGN);

  if (arguments.sync)
    mode ^= ZDB_SYNC;
  if (arguments.no_lock)
    mode ^= ZDB_NOLOCK;
  if (arguments.no_mmap)
    mode ^= ZDB_NOMMAP;

  /* open the database */
  database = zdb_open (arguments.args[0], mode, arguments.mmap_size,
                       arguments.cache_size, arguments.block_size,
                       arguments.verbose);
  if (!database)
    error (EXIT_FAILURE, errno,
           "Failed to open database %s", arguments.args[0]);

  /* if an input file was provided, use it, otherwise use stdin */
  if (!arguments.input_file || strcmp (arguments.input_file, "-") == 0)
    input = stdin;
  else
    input = fopen (arguments.input_file, "r");
  if (!input)
    {
      if (arguments.input_file)
        free (arguments.input_file);
      error (EXIT_FAILURE, errno,
             "Failed to open input file %s", arguments.input_file);
    }
  result = addrecords (database, input, arguments.delim, arguments.verbose);

  if (fclose (input))
    {
      if (arguments.input_file)
        error (EXIT_FAILURE, errno,
               "Failed to close input file %s", arguments.input_file);
      else
        error (EXIT_FAILURE, errno,
               "Failed to close stdin");
    }
  if (arguments.input_file)
    free (arguments.input_file);
  if (zdb_close (database, arguments.verbose))
    error (EXIT_FAILURE, errno,
           "Failed to close database %s", arguments.args[0]);
  if (result != 0)
    exit (EXIT_FAILURE);
  exit (EXIT_SUCCESS);
}