/Users/deen/code/yugabyte-db/src/yb/tools/yb-generate_partitions.cc
Line | Count | Source (jump to first uncovered line) |
1 | | // Copyright (c) YugaByte, Inc. |
2 | | // |
3 | | // Licensed under the Apache License, Version 2.0 (the "License"); you may not use this file except |
4 | | // in compliance with the License. You may obtain a copy of the License at |
5 | | // |
6 | | // http://www.apache.org/licenses/LICENSE-2.0 |
7 | | // |
8 | | // Unless required by applicable law or agreed to in writing, software distributed under the License |
9 | | // is distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express |
10 | | // or implied. See the License for the specific language governing permissions and limitations |
11 | | // under the License. |
12 | | // |
13 | | |
14 | | #include "yb/tools/yb-generate_partitions.h" |
15 | | |
16 | | #include <map> |
17 | | |
18 | | #include "yb/client/client.h" |
19 | | #include "yb/client/table.h" |
20 | | #include "yb/client/yb_op.h" |
21 | | |
22 | | #include "yb/common/common.pb.h" |
23 | | #include "yb/common/ql_protocol.pb.h" |
24 | | #include "yb/common/ql_value.h" |
25 | | #include "yb/common/schema.h" |
26 | | |
27 | | #include "yb/master/master_client.pb.h" |
28 | | |
29 | | #include "yb/util/enums.h" |
30 | | #include "yb/util/status.h" |
31 | | #include "yb/util/status_format.h" |
32 | | #include "yb/util/stol_utils.h" |
33 | | #include "yb/util/timestamp.h" |
34 | | |
35 | | namespace yb { |
36 | | namespace tools { |
37 | | |
38 | | using client::YBSchema; |
39 | | using client::YBTableName; |
40 | | using google::protobuf::RepeatedPtrField; |
41 | | using master::TabletLocationsPB; |
42 | | using std::map; |
43 | | using std::string; |
44 | | using std::vector; |
45 | | |
46 | | YBPartitionGenerator::YBPartitionGenerator(const YBTableName& table_name, |
47 | | const vector<string>& master_addresses) : |
48 | | table_name_(table_name), |
49 | 0 | master_addresses_(master_addresses) { |
50 | 0 | } |
51 | | |
52 | 0 | YBPartitionGenerator::~YBPartitionGenerator() { |
53 | 0 | } |
54 | | |
55 | 0 | Status YBPartitionGenerator::Init() { |
56 | 0 | client::YBClientBuilder builder; |
57 | 0 | for (const string& master_address : master_addresses_) { |
58 | 0 | builder.add_master_server_addr(master_address); |
59 | 0 | } |
60 | 0 | client_ = VERIFY_RESULT(builder.Build()); |
61 | 0 | RETURN_NOT_OK(client_->OpenTable(table_name_, &table_)); |
62 | 0 | RepeatedPtrField<TabletLocationsPB> tablets; |
63 | 0 | RETURN_NOT_OK(client_->GetTablets( |
64 | 0 | table_name_, /* max_tablets */ 0, &tablets, /* partition_list_version =*/ nullptr)); |
65 | 0 | RETURN_NOT_OK(BuildTabletMap(tablets)); |
66 | 0 | return Status::OK(); |
67 | 0 | } |
68 | | |
69 | 0 | Status YBPartitionGenerator::BuildTabletMap(const RepeatedPtrField<TabletLocationsPB>& tablets) { |
70 | 0 | for (const TabletLocationsPB& tablet : tablets) { |
71 | 0 | tablet_map_.emplace(tablet.partition().partition_key_start(), tablet); |
72 | 0 | } |
73 | 0 | return Status::OK(); |
74 | 0 | } |
75 | | |
76 | | Status YBPartitionGenerator::LookupTabletId(const string& row, |
77 | | string* tablet_id, |
78 | 0 | string* partition_key) { |
79 | 0 | return LookupTabletId(row, {}, tablet_id, partition_key); |
80 | 0 | } |
81 | | |
82 | | Status YBPartitionGenerator::LookupTabletId(const string& row, |
83 | | const std::set<int>& skipped_cols, |
84 | | string* tablet_id, |
85 | 0 | string* partition_key) { |
86 | 0 | CsvTokenizer tokenizer = Tokenize(row); |
87 | 0 | return LookupTabletIdWithTokenizer(tokenizer, skipped_cols, tablet_id, partition_key); |
88 | 0 | } |
89 | | |
90 | | Status YBPartitionGenerator::LookupTabletIdWithTokenizer(const CsvTokenizer& tokenizer, |
91 | | const std::set<int>& skipped_cols, |
92 | 0 | string* tablet_id, string* partition_key) { |
93 | 0 | const Schema &schema = table_->InternalSchema(); |
94 | 0 | size_t ncolumns = std::distance(tokenizer.begin(), tokenizer.end()); |
95 | 0 | if (ncolumns < schema.num_hash_key_columns()) { |
96 | 0 | return STATUS_SUBSTITUTE(InvalidArgument, "row doesn't have enough columns for primary " |
97 | 0 | "key, found: $0 need atleast $1", ncolumns, schema.num_hash_key_columns()); |
98 | 0 | } |
99 | | |
100 | 0 | std::unique_ptr<client::YBqlReadOp> yb_op(table_->NewQLRead()); |
101 | 0 | QLReadRequestPB* ql_read = yb_op->mutable_request(); |
102 | | |
103 | | // Set the hash column values to compute the partition key. |
104 | 0 | auto it = tokenizer.begin(); |
105 | 0 | int col_id = 0; |
106 | 0 | for (size_t i = 0; i < schema.num_hash_key_columns(); col_id++, it++) { |
107 | 0 | if (skipped_cols.find(col_id) != skipped_cols.end()) { |
108 | 0 | continue; |
109 | 0 | } |
110 | 0 | if (IsNull(*it)) { |
111 | 0 | return STATUS_SUBSTITUTE(IllegalState, "Primary key cannot be null: $0", *it); |
112 | 0 | } |
113 | | |
114 | 0 | DataType column_type = schema.column(i).type_info()->type; |
115 | 0 | auto* value_pb = ql_read->add_hashed_column_values()->mutable_value(); |
116 | |
|
117 | 0 | switch(column_type) { |
118 | 0 | YB_SET_INT_VALUE(value_pb, *it, 8); |
119 | 0 | YB_SET_INT_VALUE(value_pb, *it, 16); |
120 | 0 | YB_SET_INT_VALUE(value_pb, *it, 32); |
121 | 0 | YB_SET_INT_VALUE(value_pb, *it, 64); |
122 | 0 | case DataType::STRING: |
123 | 0 | value_pb->set_string_value(*it); |
124 | 0 | break; |
125 | 0 | case DataType::TIMESTAMP: { |
126 | 0 | auto ts = TimestampFromString(*it); |
127 | 0 | RETURN_NOT_OK(ts); |
128 | 0 | value_pb->set_timestamp_value(ts->ToInt64()); |
129 | 0 | break; |
130 | 0 | } |
131 | 0 | case DataType::BOOL: FALLTHROUGH_INTENDED; |
132 | 0 | case DataType::FLOAT: FALLTHROUGH_INTENDED; |
133 | 0 | case DataType::DOUBLE: FALLTHROUGH_INTENDED; |
134 | 0 | case DataType::MAP: FALLTHROUGH_INTENDED; |
135 | 0 | case DataType::SET: FALLTHROUGH_INTENDED; |
136 | 0 | case DataType::LIST: |
137 | 0 | LOG(FATAL) << "Invalid datatype for partition column: " << column_type; |
138 | 0 | default: |
139 | 0 | FATAL_INVALID_ENUM_VALUE(DataType, column_type); |
140 | 0 | } |
141 | 0 | i++; // Avoid incrementing if we are skipping the column. |
142 | 0 | } |
143 | | |
144 | | // Compute the hash function. |
145 | 0 | RETURN_NOT_OK(yb_op->GetPartitionKey(partition_key)); |
146 | | |
147 | | // Find the appropriate table. |
148 | 0 | auto iter = tablet_map_.upper_bound(*partition_key); |
149 | 0 | if (iter == tablet_map_.begin()) { |
150 | 0 | return STATUS_SUBSTITUTE(IllegalState, "Couldn't find partition key $0 in tablet map", |
151 | 0 | *partition_key); |
152 | 0 | } |
153 | 0 | --iter; |
154 | 0 | *tablet_id = iter->second.tablet_id(); |
155 | 0 | return Status::OK(); |
156 | 0 | } |
157 | | |
158 | | } // namespace tools |
159 | | } // namespace yb |