git.m6w6.name Git - awesomized/libmemcached/blob - src/bin/memcp.cc

   1 /*
   2     +--------------------------------------------------------------------+
   3     | libmemcached - C/C++ Client Library for memcached                  |
   4     +--------------------------------------------------------------------+
   5     | Redistribution and use in source and binary forms, with or without |
   6     | modification, are permitted under the terms of the BSD license.    |
   7     | You should have received a copy of the license in a bundled file   |
   8     | named LICENSE; in case you did not receive a copy you can review   |
   9     | the terms online at: https://opensource.org/licenses/BSD-3-Clause  |
  10     +--------------------------------------------------------------------+
  11     | Copyright (c) 2006-2014 Brian Aker   https://datadifferential.com/ |
  12     | Copyright (c) 2020 Michael Wallner   <mike@php.net>                |
  13     +--------------------------------------------------------------------+
  14 */
  15
  16 #include "mem_config.h"
  17
  18 #define PROGRAM_NAME        "memcp"
  19 #define PROGRAM_DESCRIPTION "Copy a set of files to a memcached cluster."
  20 #define PROGRAM_VERSION     "1.1"
  21
  22 #include "common/options.hpp"
  23
  24 #include <climits>
  25 #include <cstdlib>
  26 #include <fstream>
  27 #include <sstream>
  28
  29 struct memcp_file {
  30   enum key {
  31     basename,
  32     relative,
  33     absolute
  34   } type;
  35   enum op {
  36     SET,
  37     ADD,
  38     REPLACE
  39   } mode;
  40   const char *path;
  41   uint32_t flags;
  42   time_t expire;
  43 };
  44
  45 static void add_file(std::vector<memcp_file> &files, const client_options &opt, const char *file) {
  46   memcp_file::key type = memcp_file::basename;
  47   memcp_file::op mode = memcp_file::SET;
  48   uint32_t flags = 0;
  49   time_t expire = 0;
  50
  51   if (opt.isset("absolute")) {
  52     type = memcp_file::absolute;
  53   } else if (opt.isset("relative")) {
  54     type = memcp_file::relative;
  55   }
  56
  57   if (opt.isset("replace")) {
  58     mode = memcp_file::REPLACE;
  59   } else if (opt.isset("add")) {
  60     mode = memcp_file::ADD;
  61   }
  62
  63   if (auto flags_str = opt.argof("flags")) {
  64     flags = std::stoul(flags_str);
  65   }
  66   if (auto expire_str = opt.argof("expire")) {
  67     expire = std::stoul(expire_str);
  68   }
  69
  70   if (opt.isset("debug")) {
  71     auto mode_str = mode == memcp_file::REPLACE ? "REPLACE" : mode == memcp_file::ADD ? "ADD" : "SET";
  72     std::cerr << "Scheduling " << mode_str << " '" << file << "' (expire=" << expire << ", flags=" << flags << ").\n";
  73   }
  74
  75   files.emplace_back(memcp_file{type, mode, file, flags, expire});
  76 }
  77
  78 int main(int argc, char *argv[]) {
  79   memcached_st memc;
  80   std::vector<memcp_file> files{};
  81   client_options opt{PROGRAM_NAME, PROGRAM_VERSION, PROGRAM_DESCRIPTION, "file [file ...]"};
  82
  83   opt.add(nullptr, '-', no_argument, "GNU argv extension")
  84       .parse = [&files](client_options &opt_, client_options::extended_option &ext) {
  85     add_file(files, opt_, ext.arg);
  86     return true;
  87   };
  88
  89   for (const auto &def : opt.defaults) {
  90     opt.add(def);
  91   }
  92
  93   opt.add("set", 'S', no_argument, "Perform SET operations.")
  94       .parse = [](client_options &opt_, client_options::extended_option &) {
  95     opt_.unset("add");
  96     opt_.unset("replace");
  97     return true;
  98   };
  99   opt.add("add", 'A', no_argument, "Perform ADD operations.")
 100       .parse = [](client_options &opt_, client_options::extended_option &) {
 101     opt_.unset("set");
 102     opt_.unset("replace");
 103     return true;
 104   };
 105   opt.add("replace", 'R', no_argument, "Perform REPLACE operations.")
 106       .parse = [](client_options &opt_, client_options::extended_option &) {
 107     opt_.unset("set");
 108     opt_.unset("add");
 109     return true;
 110   };
 111
 112   opt.add("udp", 'U', no_argument, "Use UDP.")
 113       .apply = [](const client_options &opt, const client_options::extended_option &ext, memcached_st *memc) {
 114     if (MEMCACHED_SUCCESS != memcached_behavior_set(memc, MEMCACHED_BEHAVIOR_USE_UDP, ext.set)) {
 115       if (!opt.isset("quiet")) {
 116         std::cerr << memcached_last_error_message(memc) << "\n";
 117       }
 118       return false;
 119     }
 120     return true;
 121   };
 122   opt.add("flags", 'F', required_argument, "Set key flags, too.");
 123   opt.add("expire", 'e', required_argument, "Set expire time, too.");
 124
 125   opt.add("basename", '.', no_argument, "Use basename of path as key (default).");
 126   opt.add("relative", '+', no_argument, "Use relative path (as passed), instead of basename only.");
 127   opt.add("absolute", '/', no_argument, "Use absolute path (real path), instead of basename only.");
 128
 129   // defaults
 130   opt.set("set");
 131   opt.set("basename");
 132
 133   char **argp = nullptr;
 134   if (!opt.parse(argc, argv, &argp)) {
 135     exit(EXIT_FAILURE);
 136   }
 137
 138   if (!memcached_create(&memc)) {
 139     if (!opt.isset("quiet")) {
 140       std::cerr << "Failed to initialize memcached client.\n";
 141     }
 142     exit(EXIT_FAILURE);
 143   }
 144
 145   if (!opt.apply(&memc)) {
 146     exit(EXIT_FAILURE);
 147   }
 148
 149   if (files.empty()) {
 150     if (!*argp) {
 151       if (!opt.isset("quiet")) {
 152         std::cerr << "No file(s) provided.\n";
 153       }
 154       memcached_free(&memc);
 155       exit(EXIT_FAILURE);
 156     }
 157     for (auto arg = argp; *arg; ++arg) {
 158       add_file(files, opt, *arg);
 159     }
 160   }
 161
 162   auto exit_code = EXIT_SUCCESS;
 163   for (const auto &file : files) {
 164     auto filename = file.path;
 165     std::ifstream filestream{filename, std::ios::in|std::ios::binary};
 166
 167     if (!filestream) {
 168       if (!opt.isset("quiet")) {
 169         std::cerr << "Could not open file '" << filename << "'.\n";
 170       }
 171       exit_code = EXIT_FAILURE;
 172       // continue;
 173     } else {
 174       const char *path;
 175       char rpath[PATH_MAX+1];
 176
 177       if (file.type == memcp_file::relative) {
 178         path = filename;
 179       } else if (file.type == memcp_file::absolute) {
 180         path = realpath(filename, rpath);
 181         if (!path) {
 182           if (!opt.isset("quiet")) {
 183             perror(filename);
 184           }
 185           exit_code = EXIT_FAILURE;
 186           continue;
 187         }
 188       } else {
 189         path = basename(filename);
 190       }
 191
 192       std::ostringstream data{};
 193       data << filestream.rdbuf();
 194
 195       memcached_return_t rc;
 196       const char *mode;
 197       if (file.mode == memcp_file::REPLACE) {
 198         mode = "replace";
 199         rc = memcached_replace(&memc, path, strlen(path), data.str().c_str(), data.str().length(),
 200                                file.expire, file.flags);
 201       } else if (file.mode == memcp_file::ADD) {
 202         mode = "add";
 203         rc = memcached_add(&memc, path, strlen(path), data.str().c_str(), data.str().length(),
 204                            file.expire, file.flags);
 205       } else {
 206         mode = "set";
 207         rc = memcached_set(&memc, path, strlen(path), data.str().c_str(), data.str().length(),
 208                            file.expire, file.flags);
 209       }
 210
 211       if (MEMCACHED_SUCCESS != rc) {
 212         exit_code = EXIT_FAILURE;
 213
 214         auto error = memcached_last_error(&memc)
 215             ? memcached_last_error_message(&memc)
 216             : memcached_strerror(&memc, rc);
 217         std::cerr << "Error occurred during memcached_" << mode <<"('" << path << "'): " << error << "\n";
 218         break;
 219       }
 220
 221       if (opt.isset("verbose")) {
 222         std::cerr << path << "\n";
 223       }
 224     }
 225   }
 226
 227   memcached_free(&memc);
 228   exit(exit_code);
 229 }