1 /* -*- Mode: C++; tab-width: 4; c-basic-offset: 4; indent-tabs-mode: nil -*- */
2 /*
3  *     Copyright 2018 Couchbase, Inc
4  *
5  *   Licensed under the Apache License, Version 2.0 (the "License");
6  *   you may not use this file except in compliance with the License.
7  *   You may obtain a copy of the License at
8  *
9  *       http://www.apache.org/licenses/LICENSE-2.0
10  *
11  *   Unless required by applicable law or agreed to in writing, software
12  *   distributed under the License is distributed on an "AS IS" BASIS,
13  *   WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
14  *   See the License for the specific language governing permissions and
15  *   limitations under the License.
16  */
17 
18 #include "input_couchfile.h"
19 #include "output_couchfile.h"
20 
21 #include <getopt.h>
22 #include <cstdlib>
23 #include <iostream>
24 #include <memory>
25 #include <vector>
26 
27 /**
28  * couchfile_upgrade is a tool for making couchstore files collection-aware.
29  * This is essentially a set of steps which result in a new couchstore file
30  * that has each document assigned to the DefaultCollection.
31  */
32 
33 struct ProgramOptions {
34     OptionsSet options;
35     const char* inputFilename;
36     const char* outputFilename;
37     size_t outputBufferMaxSize = (1024 * 1024) * 500;
38 };
39 
runUpgrade(const ProgramOptions& options, Collections::InputCouchFile& input)40 static bool runUpgrade(const ProgramOptions& options,
41                        Collections::InputCouchFile& input) {
42     using PreflightStatus = Collections::InputCouchFile::PreflightStatus;
43     switch (input.preflightChecks(std::cerr)) {
44     case PreflightStatus::ReadyForUpgrade:
45         break;
46     case PreflightStatus::InputFileCannotBeProcessed:
47     case PreflightStatus::UpgradePartial:
48     case PreflightStatus::UpgradeCompleteAndPartial: {
49         std::cerr << "Pre-upgrade checks have failed\n";
50         return false;
51     }
52     case PreflightStatus::UpgradeComplete: {
53         return options.options.test(Options::Tolerate);
54     }
55     }
56 
57     // Open the output file now that the input file is ok for processing
58     Collections::OutputCouchFile output(options.options,
59                                         options.outputFilename,
60                                         CollectionID::Default,
61                                         options.outputBufferMaxSize);
62 
63     // Perform the upgrade steps
64     // 1. Write out a new file tagged in a way we can determine upgrade has
65     //    started but not finished
66     output.writeUpgradeBegin(input);
67     output.commit();
68 
69     // 2. Now run the upgrade, docs are copied from in to out and moved to
70     //    the default collection
71     input.upgrade(output);
72     output.commit();
73 
74     // 3. Write out to the new file that the upgrade is done, KV can now warmup
75     //    from this file.
76     output.writeUpgradeComplete(input);
77     output.commit();
78     return true;
79 }
80 
runStatus(Collections::InputCouchFile& input)81 static void runStatus(Collections::InputCouchFile& input) {
82     if (input.preflightChecks(std::cout) ==
83         Collections::InputCouchFile::PreflightStatus::ReadyForUpgrade) {
84         std::cout << "filename: " << input.getFilename()
85                   << " is ready for upgrade\n";
86     }
87 }
88 
usage()89 static void usage() {
90     std::cout <<
91             R"(Usage:
92     -v or --verbose        Optional: Run with verbose output to stdout.
93     -s or --status         Optional: Print upgrade status of input file.
94     -t or --tolerate       Optional: Tolerate upgraded files - exit 0 if file is already marked as upgraded.
95     -i or --input <name>   Required: Input filename.
96     -o or --output <name>  Required (only if not -s): Output filename to be created.
97     -b or --buffer size    Optional: Specify the amount of memory (bytes) we can use for buffering documents (default 524288000))"
98               << std::endl;
99 }
100 
parseArguments(int argc, char** argv)101 static ProgramOptions parseArguments(int argc, char** argv) {
102     int cmd = 0;
103     ProgramOptions pOptions{};
104 
105     struct option long_options[] = {{"tolerate", no_argument, nullptr, 't'},
106                                     {"status", no_argument, nullptr, 's'},
107                                     {"verbose", no_argument, nullptr, 'v'},
108                                     {"input", required_argument, nullptr, 'i'},
109                                     {"output", required_argument, nullptr, 'o'},
110                                     {"buffer", optional_argument, nullptr, 'b'},
111                                     {nullptr, 0, nullptr, 0}};
112 
113     while ((cmd = getopt_long(
114                     argc, argv, "tsvi:o:b:", long_options, nullptr)) != -1) {
115         switch (cmd) {
116         case 'v': {
117             pOptions.options.set(Options::Verbose);
118             std::cout << "Enabling Verbose\n";
119             break;
120         }
121         case 's': {
122             pOptions.options.set(Options::Status);
123             std::cout << "Status\n";
124             break;
125         }
126         case 't': {
127             pOptions.options.set(Options::Tolerate);
128             std::cout << "exit(0) for already upgraded files\n";
129             break;
130         }
131         case 'i': {
132             pOptions.inputFilename = optarg;
133             std::cout << "Input:" << optarg << "\n";
134             break;
135         }
136         case 'o': {
137             pOptions.outputFilename = optarg;
138             std::cout << "Output:" << optarg << "\n";
139             break;
140         }
141         case 'b': {
142             pOptions.outputBufferMaxSize = std::stoll(optarg);
143             std::cout << "Buffer size:" << pOptions.outputBufferMaxSize << "\n";
144             break;
145         }
146         case ':':
147         case '?': {
148             usage();
149             throw std::invalid_argument("Invalid Argument");
150             break;
151         }
152         }
153     }
154 
155     if (!pOptions.inputFilename) {
156         usage();
157         throw std::invalid_argument("Missing -i");
158     }
159     if (!pOptions.outputFilename && !pOptions.options.test(Options::Status)) {
160         usage();
161         throw std::invalid_argument("Missing -o");
162     }
163     if (pOptions.outputFilename && pOptions.options.test(Options::Status)) {
164         usage();
165         throw std::invalid_argument("-o with -s is not allowed");
166     }
167 
168     return pOptions;
169 }
170 
main(int argc, char** argv)171 int main(int argc, char** argv) {
172     bool success = true;
173     try {
174         auto options = parseArguments(argc, argv);
175         Collections::InputCouchFile input(options.options,
176                                           options.inputFilename);
177 
178         if (options.options.test(Options::Status)) {
179             runStatus(input);
180         } else {
181             success = runUpgrade(options, input);
182         }
183     } catch (const std::exception& e) {
184         success = false;
185         std::cerr << "An exception occurred: " << e.what() << std::endl;
186     }
187 
188     if (!success) {
189         std::cerr << "Terminating with exit code 1\n";
190     }
191 
192     return success ? EXIT_SUCCESS : EXIT_FAILURE;
193 }