bin: consolidate clients
[m6w6/libmemcached] / src / bin / memcp.cc
index 8ab28928beb2a7a23ae44d41377675d69b7bfbfc..e149ef316df876d531b4019596f6a53abd04cedd 100644 (file)
 
 #include "mem_config.h"
 
-#include <cerrno>
-#include <climits>
-#include <cstdio>
-#include <cstdlib>
-#include <cstring>
-#include <fcntl.h>
-#include <getopt.h>
-#include <iostream>
-#ifdef HAVE_STRINGS_H
-#  include <strings.h>
-#endif
-#include <sys/stat.h>
-#include <sys/types.h>
-#include <unistd.h>
-
-#include "libmemcached-1.0/memcached.h"
-
-#include "client_options.h"
-#include "utilities.h"
-
 #define PROGRAM_NAME        "memcp"
 #define PROGRAM_DESCRIPTION "Copy a set of files to a memcached cluster."
+#define PROGRAM_VERSION     "1.1"
+
+#include "common/options.hpp"
 
-/* Prototypes */
-static void options_parse(int argc, char *argv[]);
+#include <climits>
+#include <cstdlib>
+#include <fstream>
+#include <sstream>
 
-static bool opt_binary = false;
-static bool opt_udp = false;
-static bool opt_buffer = false;
-static int opt_verbose = 0;
-static char *opt_servers = NULL;
-static char *opt_hash = NULL;
-static int opt_method = OPT_SET;
-static uint32_t opt_flags = 0;
-static time_t opt_expires = 0;
-static char *opt_username;
-static char *opt_passwd;
+struct memcp_file {
+  enum key {
+    basename,
+    relative,
+    absolute
+  } type;
+  enum op {
+    SET,
+    ADD,
+    REPLACE
+  } mode;
+  const char *path;
+  uint32_t flags;
+  time_t expire;
+};
 
-static long strtol_wrapper(const char *nptr, int base, bool *error) {
-  long val;
-  char *endptr;
+static void add_file(std::vector<memcp_file> &files, const client_options &opt, const char *file) {
+  memcp_file::key type = memcp_file::basename;
+  memcp_file::op mode = memcp_file::SET;
+  uint32_t flags = 0;
+  time_t expire = 0;
 
-  errno = 0; /* To distinguish success/failure after call */
-  val = strtol(nptr, &endptr, base);
+  if (opt.isset("absolute")) {
+    type = memcp_file::absolute;
+  } else if (opt.isset("relative")) {
+    type = memcp_file::relative;
+  }
 
-  /* Check for various possible errors */
+  if (opt.isset("replace")) {
+    mode = memcp_file::REPLACE;
+  } else if (opt.isset("add")) {
+    mode = memcp_file::ADD;
+  }
 
-  if ((errno == ERANGE and (val == LONG_MAX or val == LONG_MIN)) or (errno && val == 0)) {
-    *error = true;
-    return 0;
+  if (auto flags_str = opt.argof("flags")) {
+    flags = std::stoul(flags_str);
+  }
+  if (auto expire_str = opt.argof("expire")) {
+    expire = std::stoul(expire_str);
   }
 
-  if (endptr == nptr) {
-    *error = true;
-    return 0;
+  if (opt.isset("debug")) {
+    auto mode_str = mode == memcp_file::REPLACE ? "REPLACE" : mode == memcp_file::ADD ? "ADD" : "SET";
+    std::cerr << "Scheduling " << mode_str << " '" << file << "' (expire=" << expire << ", flags=" << flags << ").\n";
   }
 
-  *error = false;
-  return val;
+  files.emplace_back(memcp_file{type, mode, file, flags, expire});
 }
 
 int main(int argc, char *argv[]) {
-  options_parse(argc, argv);
+  memcached_st memc;
+  std::vector<memcp_file> files{};
+  client_options opt{PROGRAM_NAME, PROGRAM_VERSION, PROGRAM_DESCRIPTION, "file [file ...]"};
+
+  opt.add(nullptr, '-', no_argument, "GNU argv extension")
+      .parse = [&files](client_options &opt_, client_options::extended_option &ext) {
+    add_file(files, opt_, ext.arg);
+    return true;
+  };
 
-  if (optind >= argc) {
-    fprintf(stderr, "Expected argument after options\n");
-    exit(EXIT_FAILURE);
+  for (const auto &def : opt.defaults) {
+    opt.add(def);
   }
 
-  initialize_sockets();
-
-  memcached_st *memc = memcached_create(NULL);
+  opt.add("set", 'S', no_argument, "Perform SET operations.")
+      .parse = [](client_options &opt_, client_options::extended_option &) {
+    opt_.unset("add");
+    opt_.unset("replace");
+    return true;
+  };
+  opt.add("add", 'A', no_argument, "Perform ADD operations.")
+      .parse = [](client_options &opt_, client_options::extended_option &) {
+    opt_.unset("set");
+    opt_.unset("replace");
+    return true;
+  };
+  opt.add("replace", 'R', no_argument, "Perform REPLACE operations.")
+      .parse = [](client_options &opt_, client_options::extended_option &) {
+    opt_.unset("set");
+    opt_.unset("add");
+    return true;
+  };
 
-  if (opt_udp) {
-    if (opt_verbose) {
-      std::cout << "Enabling UDP" << std::endl;
+  opt.add("udp", 'U', no_argument, "Use UDP.")
+      .apply = [](const client_options &opt, const client_options::extended_option &ext, memcached_st *memc) {
+    if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, ext.set)) {
+      if (!opt.isset("quiet")) {
+        std::cerr << memcached_last_error_message(memc) << "\n";
+      }
+      return false;
     }
+    return true;
+  };
+  opt.add("flags", 'F', required_argument, "Set key flags, too.");
+  opt.add("expire", 'e', required_argument, "Set expire time, too.");
 
-    if (memcached_failed(memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, opt_udp))) {
-      memcached_free(memc);
-      std::cerr << "Could not enable UDP protocol." << std::endl;
-      return EXIT_FAILURE;
-    }
-  }
+  opt.add("basename", '.', no_argument, "Use basename of path as key (default).");
+  opt.add("relative", '+', no_argument, "Use relative path (as passed), instead of basename only.");
+  opt.add("absolute", '/', no_argument, "Use absolute path (real path), instead of basename only.");
 
-  if (opt_buffer) {
-    if (opt_verbose) {
-      std::cout << "Enabling MEMCACHED_BEHAVIOR_BUFFER_REQUESTS" << std::endl;
-    }
+  // defaults
+  opt.set("set");
+  opt.set("basename");
 
-    if (memcached_failed(
-            memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BUFFER_REQUESTS, opt_buffer))) {
-      memcached_free(memc);
-      std::cerr << "Could not enable MEMCACHED_BEHAVIOR_BUFFER_REQUESTS." << std::endl;
-      return EXIT_FAILURE;
-    }
+  char **argp = nullptr;
+  if (!opt.parse(argc, argv, &argp)) {
+    exit(EXIT_FAILURE);
   }
 
-  process_hash_option(memc, opt_hash);
-
-  if (opt_servers == NULL) {
-    char *temp;
-
-    if ((temp = getenv("MEMCACHED_SERVERS"))) {
-      opt_servers = strdup(temp);
-    }
-#if 0
-    else if (argc >= 1 and argv[--argc])
-    {
-      opt_servers= strdup(argv[argc]);
-    }
-#endif
-
-    if (opt_servers == NULL) {
-      std::cerr << "No Servers provided" << std::endl;
-      exit(EXIT_FAILURE);
+  if (!memcached_create(&memc)) {
+    if (!opt.isset("quiet")) {
+      std::cerr << "Failed to initialize memcached client.\n";
     }
+    exit(EXIT_FAILURE);
   }
 
-  memcached_server_st *servers = memcached_servers_parse(opt_servers);
-  if (servers == NULL or memcached_server_list_count(servers) == 0) {
-    std::cerr << "Invalid server list provided:" << opt_servers << std::endl;
-    return EXIT_FAILURE;
-  }
-
-  memcached_server_push(memc, servers);
-  memcached_server_list_free(servers);
-  memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_BINARY_PROTOCOL, opt_binary);
-  if (opt_username and LIBMEMCACHED_WITH_SASL_SUPPORT == 0) {
-    memcached_free(memc);
-    std::cerr << "--username was supplied, but binary was not built with SASL support."
-              << std::endl;
-    return EXIT_FAILURE;
+  if (!opt.apply(&memc)) {
+    exit(EXIT_FAILURE);
   }
 
-  if (opt_username) {
-    memcached_return_t ret;
-    if (memcached_failed(ret = memcached_set_sasl_auth_data(memc, opt_username, opt_passwd))) {
-      std::cerr << memcached_last_error_message(memc) << std::endl;
-      memcached_free(memc);
-      return EXIT_FAILURE;
+  if (files.empty()) {
+    if (!*argp) {
+      if (!opt.isset("quiet")) {
+        std::cerr << "No file(s) provided.\n";
+      }
+      memcached_free(&memc);
+      exit(EXIT_FAILURE);
+    }
+    for (auto arg = argp; *arg; ++arg) {
+      add_file(files, opt, *arg);
     }
   }
 
-  int exit_code = EXIT_SUCCESS;
-  while (optind < argc) {
-    int fd = open(argv[optind], O_RDONLY);
-    if (fd < 0) {
-      std::cerr << "memcp " << argv[optind] << " " << strerror(errno) << std::endl;
-      optind++;
-      exit_code = EXIT_FAILURE;
-      continue;
-    }
+  auto exit_code = EXIT_SUCCESS;
+  for (const auto &file : files) {
+    auto filename = file.path;
+    std::ifstream filestream{filename, std::ios::in|std::ios::binary};
 
-    struct stat sbuf;
-    if (fstat(fd, &sbuf) == -1) {
-      std::cerr << "memcp " << argv[optind] << " " << strerror(errno) << std::endl;
-      optind++;
+    if (!filestream) {
+      if (!opt.isset("quiet")) {
+        std::cerr << "Could not open file '" << filename << "'.\n";
+      }
       exit_code = EXIT_FAILURE;
-      continue;
-    }
-
-    char *ptr = rindex(argv[optind], '/');
-    if (ptr) {
-      ptr++;
+      // continue;
     } else {
-      ptr = argv[optind];
-    }
-
-    if (opt_verbose) {
-      static const char *opstr[] = {"set", "add", "replace"};
-      printf("op: %s\nsource file: %s\nlength: %lu\n"
-             "key: %s\nflags: %x\nexpires: %lu\n",
-             opstr[opt_method - OPT_SET], argv[optind], (unsigned long) sbuf.st_size, ptr,
-             opt_flags, (unsigned long) opt_expires);
-    }
-
-    // The file may be empty
-    char *file_buffer_ptr = NULL;
-    if (sbuf.st_size > 0) {
-      if ((file_buffer_ptr = (char *) malloc(sizeof(char) * (size_t) sbuf.st_size)) == NULL) {
-        std::cerr << "Error allocating file buffer(" << strerror(errno) << ")" << std::endl;
-        close(fd);
-        exit(EXIT_FAILURE);
+      const char *path;
+      char rpath[PATH_MAX+1];
+
+      if (file.type == memcp_file::relative) {
+        path = filename;
+      } else if (file.type == memcp_file::absolute) {
+        path = realpath(filename, rpath);
+        if (!path) {
+          if (!opt.isset("quiet")) {
+            perror(filename);
+          }
+          exit_code = EXIT_FAILURE;
+          continue;
+        }
+      } else {
+        path = basename(filename);
       }
 
-      ssize_t read_length;
-      if ((read_length = ::read(fd, file_buffer_ptr, (size_t) sbuf.st_size)) == -1) {
-        std::cerr << "Error while reading file " << file_buffer_ptr << " (" << strerror(errno)
-                  << ")" << std::endl;
-        close(fd);
-        free(file_buffer_ptr);
-        exit(EXIT_FAILURE);
+      std::ostringstream data{};
+      data << filestream.rdbuf();
+
+      memcached_return_t rc;
+      const char *mode;
+      if (file.mode == memcp_file::REPLACE) {
+        mode = "replace";
+        rc = memcached_replace(&memc, path, strlen(path), data.str().c_str(), data.str().length(),
+                               file.expire, file.flags);
+      } else if (file.mode == memcp_file::ADD) {
+        mode = "add";
+        rc = memcached_add(&memc, path, strlen(path), data.str().c_str(), data.str().length(),
+                           file.expire, file.flags);
+      } else {
+        mode = "set";
+        rc = memcached_set(&memc, path, strlen(path), data.str().c_str(), data.str().length(),
+                           file.expire, file.flags);
       }
 
-      if (read_length != sbuf.st_size) {
-        std::cerr << "Failure while reading file. Read length was not equal to stat() length"
-                  << std::endl;
-        close(fd);
-        free(file_buffer_ptr);
-        exit(EXIT_FAILURE);
-      }
-    }
-
-    memcached_return_t rc;
-    if (opt_method == OPT_ADD) {
-      rc = memcached_add(memc, ptr, strlen(ptr), file_buffer_ptr, (size_t) sbuf.st_size,
-                         opt_expires, opt_flags);
-    } else if (opt_method == OPT_REPLACE) {
-      rc = memcached_replace(memc, ptr, strlen(ptr), file_buffer_ptr, (size_t) sbuf.st_size,
-                             opt_expires, opt_flags);
-    } else {
-      rc = memcached_set(memc, ptr, strlen(ptr), file_buffer_ptr, (size_t) sbuf.st_size,
-                         opt_expires, opt_flags);
-    }
-
-    if (memcached_failed(rc)) {
-      std::cerr << "Error occrrured during memcached_set(): " << memcached_last_error_message(memc)
-                << std::endl;
-      exit_code = EXIT_FAILURE;
-    }
-
-    ::free(file_buffer_ptr);
-    ::close(fd);
-    optind++;
-  }
-
-  if (opt_verbose) {
-    std::cout << "Calling memcached_free()" << std::endl;
-  }
-
-  memcached_free(memc);
-
-  if (opt_servers) {
-    free(opt_servers);
-  }
-
-  if (opt_hash) {
-    free(opt_hash);
-  }
-
-  return exit_code;
-}
-
-static void options_parse(int argc, char *argv[]) {
-  memcached_programs_help_st help_options[] = {
-      {0},
-  };
-
-  static struct option long_options[] = {
-      {(OPTIONSTRING) "version", no_argument, NULL, OPT_VERSION},
-      {(OPTIONSTRING) "help", no_argument, NULL, OPT_HELP},
-      {(OPTIONSTRING) "quiet", no_argument, NULL, OPT_QUIET},
-      {(OPTIONSTRING) "udp", no_argument, NULL, OPT_UDP},
-      {(OPTIONSTRING) "buffer", no_argument, NULL, OPT_BUFFER},
-      {(OPTIONSTRING) "verbose", no_argument, &opt_verbose, OPT_VERBOSE},
-      {(OPTIONSTRING) "debug", no_argument, &opt_verbose, OPT_DEBUG},
-      {(OPTIONSTRING) "servers", required_argument, NULL, OPT_SERVERS},
-      {(OPTIONSTRING) "flag", required_argument, NULL, OPT_FLAG},
-      {(OPTIONSTRING) "expire", required_argument, NULL, OPT_EXPIRE},
-      {(OPTIONSTRING) "set", no_argument, NULL, OPT_SET},
-      {(OPTIONSTRING) "add", no_argument, NULL, OPT_ADD},
-      {(OPTIONSTRING) "replace", no_argument, NULL, OPT_REPLACE},
-      {(OPTIONSTRING) "hash", required_argument, NULL, OPT_HASH},
-      {(OPTIONSTRING) "binary", no_argument, NULL, OPT_BINARY},
-      {(OPTIONSTRING) "username", required_argument, NULL, OPT_USERNAME},
-      {(OPTIONSTRING) "password", required_argument, NULL, OPT_PASSWD},
-      {0, 0, 0, 0},
-  };
+      if (MEMCACHED_SUCCESS != rc) {
+        exit_code = EXIT_FAILURE;
 
-  bool opt_version = false;
-  bool opt_help = false;
-  int option_index = 0;
-
-  while (1) {
-    int option_rv = getopt_long(argc, argv, "Vhvds:", long_options, &option_index);
-
-    if (option_rv == -1)
-      break;
-
-    switch (option_rv) {
-    case 0:
-      break;
-
-    case OPT_BINARY:
-      opt_binary = true;
-      break;
-
-    case OPT_VERBOSE: /* --verbose or -v */
-      opt_verbose = OPT_VERBOSE;
-      break;
-
-    case OPT_DEBUG: /* --debug or -d */
-      opt_verbose = OPT_DEBUG;
-      break;
-
-    case OPT_VERSION: /* --version or -V */
-      opt_version = true;
-      break;
-
-    case OPT_HELP: /* --help or -h */
-      opt_help = true;
-      break;
-
-    case OPT_SERVERS: /* --servers or -s */
-      opt_servers = strdup(optarg);
-      break;
-
-    case OPT_FLAG: /* --flag */
-    {
-      bool strtol_error;
-      opt_flags = (uint32_t) strtol_wrapper(optarg, 16, &strtol_error);
-      if (strtol_error == true) {
-        fprintf(stderr, "Bad value passed via --flag\n");
-        exit(1);
+        auto error = memcached_last_error(&memc)
+            ? memcached_last_error_message(&memc)
+            : memcached_strerror(&memc, rc);
+        std::cerr << "Error occurred during memcached_" << mode <<"('" << path << "'): " << error << "\n";
+        break;
       }
-    } break;
 
-    case OPT_EXPIRE: /* --expire */
-    {
-      bool strtol_error;
-      opt_expires = (time_t) strtol_wrapper(optarg, 10, &strtol_error);
-      if (strtol_error == true) {
-        fprintf(stderr, "Bad value passed via --expire\n");
-        exit(1);
+      if (opt.isset("verbose")) {
+        std::cerr << path << "\n";
       }
-    } break;
-
-    case OPT_SET:
-      opt_method = OPT_SET;
-      break;
-
-    case OPT_REPLACE:
-      opt_method = OPT_REPLACE;
-      break;
-
-    case OPT_ADD:
-      opt_method = OPT_ADD;
-      break;
-
-    case OPT_HASH:
-      opt_hash = strdup(optarg);
-      break;
-
-    case OPT_USERNAME:
-      opt_username = optarg;
-      break;
-
-    case OPT_PASSWD:
-      opt_passwd = optarg;
-      break;
-
-    case OPT_QUIET:
-      close_stdio();
-      break;
-
-    case OPT_UDP:
-      opt_udp = true;
-      break;
-
-    case OPT_BUFFER:
-      opt_buffer = true;
-      break;
-
-    case '?':
-      /* getopt_long already printed an error message. */
-      exit(1);
-    default:
-      abort();
     }
   }
 
-  if (opt_version) {
-    version_command(PROGRAM_NAME);
-    exit(EXIT_SUCCESS);
-  }
-
-  if (opt_help) {
-    help_command(PROGRAM_NAME, PROGRAM_DESCRIPTION, long_options, help_options);
-    exit(EXIT_SUCCESS);
-  }
+  memcached_free(&memc);
+  exit(exit_code);
 }