source: trunk/stoposclient.cpp @ 20

/*
   Copyright 2013 Willem Vermin, SURFsara

   Licensed under the Apache License, Version 2.0 (the "License");
   you may not use this file except in compliance with the License.
   You may obtain a copy of the License at

       http://www.apache.org/licenses/LICENSE-2.0

   Unless required by applicable law or agreed to in writing, software
   distributed under the License is distributed on an "AS IS" BASIS,
   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
   See the License for the specific language governing permissions and
   limitations under the License.
 */
#include <iostream>
#include <vector>
#include <stdlib.h>
#include <getopt.h>
#include <fstream>
#include "wrequest.h"
#include "shellenv.h"
#include "stopos.h"
#define  CONNECT_ERROR_MSG "ERROR: Cannot connect to server"
#define  CONNECT_ERROR_RETRIES 10

struct returnvalue
{
  std::string msg;
  std::string key;
  std::string comm;
  std::string count;
  std::string present;
  std::string present0;
  std::string value;

  void parse(std::string &s)
  {
    // we expect somewhere in the output of the server a line like
    // STOPOS:/OK/123/beautiful
    // OK        -> msg
    // 123       -> key
    // 1         -> comm
    // beautiful -> value
    //
    // OK, 123 and beautiful are hex coded
    this->msg       = "";
    this->key       = "";
    this->comm      = "0";
    this->count     = "0";
    this->present   = "0";
    this->present0  = "0";
    this->value     = "";

    size_t p = s.find("STOPOS:");
    if (p == s.npos)
    {
      this->msg="Error: stopos server not running.";
      return;
    }

    size_t q = s.find('\n',p+1);
    std::string l = s.substr(p,q-p);
    std::vector <std::string> v;

    split(v,l,outsep);
    for (unsigned int i=1; i<v.size(); i++)
    {
      switch (i)
      {
        case 1: this->msg       = zhextostr(v[i]); break;
        case 2: this->key       = zhextostr(v[i]); break;
        case 3: this->comm      = zhextostr(v[i]); break;
        case 4: this->count     = zhextostr(v[i]); break;
        case 5: this->present   = zhextostr(v[i]); break;
        case 6: this->present0  = zhextostr(v[i]); break;
        case 7: this->value     = zhextostr(v[i]); break;
      }
      // special case for key: if key == "", we try to take
      // it from the environment.
      if (trim(this->key) == "")
        this->key = envtostr("STOPOS_KEY");
    }
  }
};

void usage()
{
  std::cerr << "Usage:" << std::endl;

  std::cerr << program_name << " -h,--help"                           << std::endl;
  std::cerr << program_name << " -v,--version"                        << std::endl;
  std::cerr << program_name << " create [-p,--pool POOL]"             << std::endl;
  std::cerr << program_name << " purge  [-p,--pool POOL]"             << std::endl;
  std::cerr << program_name << " pools"                               << std::endl;
  std::cerr << program_name << " add    [-p,--pool POOL] [FILENAME]"  << std::endl;
  std::cerr << program_name << " next   [-p,--pool POOL] [-m,--mult]" << std::endl;
  std::cerr << program_name << " remove [-p,--pool POOL] [KEY]"       << std::endl;
  std::cerr << program_name << " status [-p,--pool POOL]"             << std::endl;
  std::cerr << program_name << " dump   [-p,--pool POOL]"             << std::endl;
  std::cerr << " -q,--quiet added to any command will reduce the ouput" << std::endl;
  std::cerr << "Defaults:"                                            << std::endl;
  std::cerr << "  POOL     : pool"                                    << std::endl;
  std::cerr << "  FILENAME : standard input"                          << std::endl;

  std::cerr << "Environment:"<<std::endl;
  std::cerr << "   Defined by user:"<<std::endl;
  std::cerr << "     STOPOS_POOL    : name of pool"<<std::endl;
  std::cerr << "     STOPOS_KEY     : key"<<std::endl;
  std::cerr << "   Defined by "<<program_name<<":"<<std::endl;
  std::cerr << "     STOPOS_RC      : return message of all commands. OK == good" << std::endl;
  std::cerr << "     STOPOS_VALUE   : result of 'next' or 'pools'"<<std::endl;
  std::cerr << "     STOPOS_KEY     : result of 'next'"<<std::endl;
  std::cerr << "     STOPOS_COUNT   : result of 'status'"<<std::endl;
  std::cerr << "     STOPOS_PRESENT : result of 'status'"<<std::endl;
  std::cerr << "     STOPOS_PRESENT0: result of 'status'"<<std::endl;



}

int handle_add(wrequest &w,const std::string &fname,returnvalue &r, const bool quiet)
{
  size_t number_of_lines=0;

  bool read_from_stdin = 0;
  int rc;
  std::ifstream inp;
  std::istream *pinp;
  if(fname == "")  // read from stdin
  {
    read_from_stdin = 1;
    pinp = &std::cin;
  }
  else
  {
    inp.open(fname.c_str());
    if (!inp.is_open())
    {
      r.msg = "ERROR: Cannot open '"+fname+"'";
      return 1;
    }

    number_of_lines = count_lines(inp);
    pinp = &inp;
  }
  std::string line;
  size_t n = 0;
  w.set_command("add");
  std::string body,header;

  progress_bar bar;

  while (getline(*pinp,line))
  {
    if (! quiet)
    {
      bool c;
      std::string s;
      if (read_from_stdin)
        s = bar.show(n,c);
      else
        s = bar.show(n,(double)n/number_of_lines,c);
      if (c)
        std::cerr << '\r' << s;
    }
    if (line.size() > maxlinelength)
    {
      r.msg =  "Error: line length exeeds "+ NumberToString(maxlinelength);
      break;
    }
    w.set_value(line);

    for (int i=0; i<CONNECT_ERROR_RETRIES; i++)
    {
      rc = w.send_message(body,header);
      if (rc == 0) break;
      sleep(1);
    }
    if (rc != 0)
    {
      r.msg = CONNECT_ERROR_MSG;
      return 1;
    }

    r.parse(body);
    if (r.msg != "OK")
      break;
    n++;
  }
  if (! quiet)
  {
    bar.clear();
    if (read_from_stdin)
      std::cerr << '\r' << bar.show(n) << std::endl;
    else
      std::cerr << '\r' << bar.show(n,1.0) << std::endl;

    std::cerr << program_name << ": Number of lines added = "<<n<<std::endl;
  }
  if (r.msg != "OK")
    return 1;
  return 0;
}

int main(int argc, char*argv[])
{
  wrequest w;
  
  int rc;

  rc = w.read_config();
  if (rc != 0 || w.get_config().size() != w.config_size)
  {
    std::cerr << "No valid config file found, creating it ..." << std::endl;
    rc = w.write_config();
    rc = rc | w.read_config();
    if (rc != 0)
    {
      std::cerr << "Cannot write config file, exiting." << std::endl;
      return 1;
    }
  }

  bool addflag=0;
  std::string inputfilename;
  std::vector <std::string> parms;

  static struct option long_options[] =
  {
    {"help",     0, 0, 'h'},
    {"pool",     1, 0, 'p'},
    {"version",  0, 0, 'v'},
    {"multi",    0, 0, 'm'},
    {"quiet",    0, 0, 'q'},
    {0,          0, 0, 0}
  };

  // determine name of pool, in case it is not set on commandline

  std::string pool = default_pool;
  std::string p = envtostr("STOPOS_POOL");

  if (p != "")
    pool = p;

  if (trim(w.get_pool()) != "")
    pool = w.get_pool();

  bool quiet = 0;

  while(1)
  {
    int c = getopt_long(argc,argv,"-hp:vmq",long_options,0);
    if (c == -1)
    {
      if (optind < argc)  // unexpected extra argument
      {
        std::cerr << program_name << ": Unexpected argument:" << argv[optind] << std::endl;
        return 1;
      }
      break;
    }
    if ( c == '?' || c == ':' )
    {
      std::cerr << program_name << ": Invalid parameter, try "<< program_name<<" --help " <<std::endl;
      return 1;
    }
    switch(c)
    {
      case 'h':
        usage();
        return 0;
        break;
      case 'p':
        pool = optarg;
        break;
      case 'f':
        inputfilename = optarg;
        break;
      case 'v':
        std::cerr << program_version << std::endl;
        return 0;
        break;
      case 'm':
        w.set_multi("yes");
        break;
      case 'q':
        quiet = 1;
        break;
      case 1:
        parms.push_back(optarg);
    }
  }

  w.set_pool(pool);
  std::string serverurl=envtostr("STOPOS_SERVER_URL");
  if (serverurl == "")
    serverurl = SERVER_URL;
  w.set_serverurl(serverurl);
#ifdef DEFAULT_FLAT
  std::string prot = "flat";
#endif
#ifdef DEFAULT_GDBM
  std::string prot = "gdbm";
#endif
#ifdef DEFAULT_FILES
  std::string prot = "files";
#endif
#ifdef DEFAULT_MYSQL
  std::string prot = "mysql";
#endif
  std::string protenv;
  protenv = envtostr("STOPOS_PROT");
  if (protenv.size() !=0)
    prot = protenv;
  w.set_protocol(prot);
  w.set_whoami(envtostr("USER"));

  rc = 0;
  returnvalue r;

  if (parms.size() == 0)
  {
    std::cerr << program_name << ": No command given." << std::endl;
    return 1;
  }

  std::string command = parms[0];

  // only the remove and add commands can accept an extra parameter
  //

  if (parms.size() >1)
    if (command !="remove" && command != "add")
    {
      std::cerr << program_name << ": Extraneous parameter:'"<<parms[1]<<"'."<<std::endl;
      return 1;
    }

  bool count_set     = 0;
  bool value_set     = 0;
  bool committed_set = 0;
  bool present_set   = 0;
  bool present0_set  = 0;
  bool key_set       = 0;
  bool rc_set        = 0;

  if(command == "create")
  {
    rc_set = 1;
    w.set_command("create");
  }
  else if(command == "purge")
  {
    rc_set = 1;
    w.set_command("purge");
  }
  else if(command == "status")
  {
    rc_set       = 1;
    present_set  = 1;
    present0_set = 1;
    count_set    = 1;
    w.set_command("status");
  }
  else if(command == "next")
  {
    rc_set        = 1;
    value_set     = 1;
    key_set       = 1;
    committed_set = 1;
    w.set_command("next");
  }
  else if(command == "dump")
  {
    rc_set        = 1;
    value_set     = 1;
    key_set       = 1;
    committed_set = 1;
    w.set_command("dump");
  }
  else if(command == "pools")
  {
    rc_set    = 1;
    value_set = 1;
    w.set_command("pools");
  }
  else if(command == "remove")
  {
    rc_set    = 1;
    w.set_command("remove");
    if (parms.size() < 2)
    {
      std::string a=envtostr("STOPOS_KEY");
      if (a == "")
      {
        std::cerr << program_name << ": No key found, exit.\n";
        return 1;
      }
      w.set_value(a);
    }
    else
      w.set_value(parms[1]);
  }
  else if(command == "add")
  {
    rc_set  = 1;
    key_set = 1;
    addflag = 1;
  }
  else
  {
    std::cerr << program_name << ": Invalid command: '"<<command<<"'"<<std::endl;
    return 1;
  }

  if(addflag)
  {
    if (parms.size() > 1)
      inputfilename = parms[1];
    rc |= handle_add(w,inputfilename,r,quiet);
  }
  else
  {
    std::string body,header;

    for (int i=0; i<CONNECT_ERROR_RETRIES; i++)
    {
      rc = w.send_message(body,header);
      if (rc == 0) break;
      sleep(1);
    }
    if (rc != 0)
      r.msg = CONNECT_ERROR_MSG;
    else
      r.parse(body);
  }

  int shelltype;

  char *shell = getenv("SHELL");
  std::string sh;
  if (shell == 0)
    sh = "bash";
  else
    sh = shell;

  if (sh.find("csh") != sh.npos)
    shelltype=SHELL_CSH;
  else
    shelltype = SHELL_SH;

  if(rc_set)        std::cout << shellenv(r.msg,      "STOPOS_RC",        shelltype);
  if(key_set)       std::cout << shellenv(r.key,      "STOPOS_KEY",       shelltype);
  if(committed_set) std::cout << shellenv(r.comm,     "STOPOS_COMMITTED", shelltype);
  if(count_set)     std::cout << shellenv(r.count,    "STOPOS_COUNT",     shelltype);
  if(present_set)   std::cout << shellenv(r.present,  "STOPOS_PRESENT",   shelltype);
  if(present0_set)  std::cout << shellenv(r.present0, "STOPOS_PRESENT0",  shelltype);
  if(value_set)     std::cout << shellenv(r.value,    "STOPOS_VALUE",     shelltype);

  if (command == "status" && r.msg == "OK" && ! quiet)
  {
    std::cerr << "lines added:           " << r.count    <<std::endl;
    std::cerr << "lines present:         " << r.present  <<std::endl;
    std::cerr << "lines never committed: " << r.present0 << std::endl;
  }
  else if (command == "pools" && ! quiet)
  {
    if(r.value.size()>0) std::cerr << r.value << std::endl;
  }
  if (r.msg != "OK" && ! quiet)
  {
    std::cerr << program_name<<": "<<r.msg << std::endl;
    rc = 1;
  }

  return rc;
}