| 145 | } |
| 146 | |
| 147 | Status BigQueryTableAccessor::ReadRow(int64* row_id, Example* example) { |
| 148 | if (Done()) { |
| 149 | return errors::OutOfRange("Reached end of table ", FullTableName()); |
| 150 | } |
| 151 | |
| 152 | // If the next row is already fetched and cached, return the row from the |
| 153 | // buffer. Otherwise, fill up the row buffer from BigQuery and return a row. |
| 154 | if (next_row_in_buffer_ != -1 && |
| 155 | next_row_in_buffer_ < ComputeMaxResultsArg()) { |
| 156 | *row_id = first_buffered_row_index_ + next_row_in_buffer_; |
| 157 | *example = row_buffer_[next_row_in_buffer_]; |
| 158 | next_row_in_buffer_++; |
| 159 | } else { |
| 160 | string auth_token; |
| 161 | TF_RETURN_IF_ERROR( |
| 162 | AuthProvider::GetToken(auth_provider_.get(), &auth_token)); |
| 163 | |
| 164 | std::unique_ptr<HttpRequest> request(http_request_factory_->Create()); |
| 165 | std::vector<char> output_buffer; |
| 166 | output_buffer.reserve(kBufferSize); |
| 167 | |
| 168 | // The first time that we access BigQuery there is no page token. After that |
| 169 | // we use the page token (which returns rows faster). |
| 170 | if (!next_page_token_.empty()) { |
| 171 | request->SetUri(strings::StrCat( |
| 172 | BigQueryUriPrefix(), "data?maxResults=", ComputeMaxResultsArg(), |
| 173 | "&pageToken=", request->EscapeString(next_page_token_))); |
| 174 | first_buffered_row_index_ += row_buffer_.size(); |
| 175 | } else { |
| 176 | request->SetUri(strings::StrCat( |
| 177 | BigQueryUriPrefix(), "data?maxResults=", ComputeMaxResultsArg(), |
| 178 | "&startIndex=", first_buffered_row_index_)); |
| 179 | } |
| 180 | request->AddAuthBearerHeader(auth_token); |
| 181 | request->SetResultBuffer(&output_buffer); |
| 182 | TF_RETURN_WITH_CONTEXT_IF_ERROR(request->Send(), " when reading rows from ", |
| 183 | FullTableName()); |
| 184 | |
| 185 | // Parse the returned row. |
| 186 | StringPiece response_piece = |
| 187 | StringPiece(&output_buffer[0], output_buffer.size()); |
| 188 | Json::Value root; |
| 189 | TF_RETURN_IF_ERROR(ParseJson(response_piece, &root)); |
| 190 | for (unsigned int i = 0; i < root["rows"].size(); ++i) { |
| 191 | row_buffer_[i].Clear(); |
| 192 | TF_RETURN_IF_ERROR( |
| 193 | ParseColumnValues(root["rows"][i], schema_root_, &row_buffer_[i])); |
| 194 | } |
| 195 | |
| 196 | next_page_token_ = root["pageToken"].asString(); |
| 197 | *row_id = first_buffered_row_index_; |
| 198 | *example = row_buffer_[0]; |
| 199 | next_row_in_buffer_ = 1; |
| 200 | } |
| 201 | return Status::OK(); |
| 202 | } |
| 203 | |
| 204 | int64 BigQueryTableAccessor::ComputeMaxResultsArg() { |