1 /* Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2 
3 Licensed under the Apache License, Version 2.0 (the "License");
4 you may not use this file except in compliance with the License.
5 You may obtain a copy of the License at
6 
7     http://www.apache.org/licenses/LICENSE-2.0
8 
9 Unless required by applicable law or agreed to in writing, software
10 distributed under the License is distributed on an "AS IS" BASIS,
11 WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
12 See the License for the specific language governing permissions and
13 limitations under the License.
14 ==============================================================================*/
15 
16 #ifndef TENSORFLOW_CONTRIB_BIGTABLE_KERNELS_BIGTABLE_LIB_H_
17 #define TENSORFLOW_CONTRIB_BIGTABLE_KERNELS_BIGTABLE_LIB_H_
18 
19 // Note: we use bigtable/client/internal/table.h as this is the no-exception API
20 
21 #include "google/cloud/bigtable/data_client.h"
22 #include "google/cloud/bigtable/internal/table.h"
23 #include "tensorflow/core/framework/dataset.h"
24 #include "tensorflow/core/framework/resource_mgr.h"
25 
26 namespace tensorflow {
27 
28 Status GrpcStatusToTfStatus(const ::grpc::Status& status);
29 Status GcpStatusToTfStatus(const ::google::cloud::Status& status);
30 
31 string RegexFromStringSet(const std::vector<string>& strs);
32 
33 class BigtableClientResource : public ResourceBase {
34  public:
BigtableClientResource(string project_id,string instance_id,std::shared_ptr<google::cloud::bigtable::DataClient> client)35   BigtableClientResource(
36       string project_id, string instance_id,
37       std::shared_ptr<google::cloud::bigtable::DataClient> client)
38       : project_id_(std::move(project_id)),
39         instance_id_(std::move(instance_id)),
40         client_(std::move(client)) {}
41 
get_client()42   std::shared_ptr<google::cloud::bigtable::DataClient> get_client() {
43     return client_;
44   }
45 
DebugString()46   string DebugString() const override {
47     return strings::StrCat("BigtableClientResource(project_id: ", project_id_,
48                            ", instance_id: ", instance_id_, ")");
49   }
50 
51  private:
52   const string project_id_;
53   const string instance_id_;
54   std::shared_ptr<google::cloud::bigtable::DataClient> client_;
55 };
56 
57 class BigtableTableResource : public ResourceBase {
58  public:
BigtableTableResource(BigtableClientResource * client,string table_name)59   BigtableTableResource(BigtableClientResource* client, string table_name)
60       : client_(client),
61         table_name_(std::move(table_name)),
62         table_(client->get_client(), table_name_,
63                google::cloud::bigtable::AlwaysRetryMutationPolicy()) {
64     client_->Ref();
65   }
66 
~BigtableTableResource()67   ~BigtableTableResource() override { client_->Unref(); }
68 
table()69   ::google::cloud::bigtable::noex::Table& table() { return table_; }
70 
DebugString()71   string DebugString() const override {
72     return strings::StrCat(
73         "BigtableTableResource(client: ", client_->DebugString(),
74         ", table: ", table_name_, ")");
75   }
76 
77  private:
78   BigtableClientResource* client_;  // Ownes one ref.
79   const string table_name_;
80   ::google::cloud::bigtable::noex::Table table_;
81 };
82 
83 namespace data {
84 
85 // BigtableReaderDatasetIterator is an abstract class for iterators from
86 // datasets that are "readers" (source datasets, not transformation datasets)
87 // that read from Bigtable.
88 template <typename Dataset>
89 class BigtableReaderDatasetIterator : public DatasetIterator<Dataset> {
90  public:
BigtableReaderDatasetIterator(const typename DatasetIterator<Dataset>::Params & params)91   explicit BigtableReaderDatasetIterator(
92       const typename DatasetIterator<Dataset>::Params& params)
93       : DatasetIterator<Dataset>(params) {}
94 
GetNextInternal(IteratorContext * ctx,std::vector<Tensor> * out_tensors,bool * end_of_sequence)95   Status GetNextInternal(IteratorContext* ctx, std::vector<Tensor>* out_tensors,
96                          bool* end_of_sequence) override {
97     mutex_lock l(mu_);
98     TF_RETURN_IF_ERROR(EnsureIteratorInitialized());
99     if (iterator_ == reader_->end()) {
100       *end_of_sequence = true;
101       return Status::OK();
102     }
103     if (!*iterator_) {
104       return GcpStatusToTfStatus(iterator_->status());
105     }
106     *end_of_sequence = false;
107     google::cloud::bigtable::Row& row = **iterator_;
108     Status s = ParseRow(ctx, row, out_tensors);
109     // Ensure we always advance.
110     ++iterator_;
111     return s;
112   }
113 
114  protected:
115   virtual ::google::cloud::bigtable::RowRange MakeRowRange() = 0;
116   virtual ::google::cloud::bigtable::Filter MakeFilter() = 0;
117   virtual Status ParseRow(IteratorContext* ctx,
118                           const ::google::cloud::bigtable::Row& row,
119                           std::vector<Tensor>* out_tensors) = 0;
120 
121  private:
EnsureIteratorInitialized()122   Status EnsureIteratorInitialized() EXCLUSIVE_LOCKS_REQUIRED(mu_) {
123     if (reader_) {
124       return Status::OK();
125     }
126 
127     auto rows = MakeRowRange();
128     auto filter = MakeFilter();
129 
130     // Note: the this in `this->dataset()` below is necessary due to namespace
131     // name conflicts.
132     reader_.reset(new ::google::cloud::bigtable::RowReader(
133         this->dataset()->table()->table().ReadRows(rows, filter)));
134     iterator_ = reader_->begin();
135     return Status::OK();
136   }
137 
138   mutex mu_;
139   std::unique_ptr<::google::cloud::bigtable::RowReader> reader_ GUARDED_BY(mu_);
140   ::google::cloud::bigtable::RowReader::iterator iterator_ GUARDED_BY(mu_);
141 };
142 
143 }  // namespace data
144 
145 }  // namespace tensorflow
146 
147 #endif  // TENSORFLOW_CONTRIB_BIGTABLE_KERNELS_BIGTABLE_LIB_H_
148