1 /* 2 * Copyright 2022 Google LLC 3 * 4 * Licensed under the Apache License, Version 2.0 (the "License"); 5 * you may not use this file except in compliance with the License. 6 * You may obtain a copy of the License at 7 * 8 * http://www.apache.org/licenses/LICENSE-2.0 9 * 10 * Unless required by applicable law or agreed to in writing, software 11 * distributed under the License is distributed on an "AS IS" BASIS, 12 * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. 13 * See the License for the specific language governing permissions and 14 * limitations under the License. 15 */ 16 #ifndef FCP_CLIENT_ENGINE_EXAMPLE_ITERATOR_FACTORY_H_ 17 #define FCP_CLIENT_ENGINE_EXAMPLE_ITERATOR_FACTORY_H_ 18 19 #include <functional> 20 #include <memory> 21 22 #include "absl/status/statusor.h" 23 #include "fcp/client/simple_task_environment.h" 24 #include "fcp/protos/plan.pb.h" 25 26 namespace fcp { 27 namespace client { 28 namespace engine { 29 30 // An interface for engine-internal usage, describing a handle with which 31 // ExampleIterator instances can be created, based on a given ExampleSelector. 32 // Each handle indicates which type of ExampleSelector it is actually able to 33 // handle. 34 class ExampleIteratorFactory { 35 public: 36 // Whether this factory can create iterators that serve the given 37 // `ExampleSelector`-described query. 38 virtual bool CanHandle( 39 const google::internal::federated::plan::ExampleSelector& 40 example_selector) = 0; 41 42 // Creates an iterator for the given `ExampleSelector`-described query. 43 virtual absl::StatusOr<std::unique_ptr<ExampleIterator>> 44 CreateExampleIterator( 45 const google::internal::federated::plan::ExampleSelector& 46 example_selector) = 0; 47 48 // Whether stats should be generated and logged into the OpStats database for 49 // iterators created by this factory. 50 virtual bool ShouldCollectStats() = 0; 51 ~ExampleIteratorFactory()52 virtual ~ExampleIteratorFactory() {} 53 }; 54 55 // A utility ExampleIteratorFactory implementation that can use simple 56 // std::function objects and wrap them. 57 class FunctionalExampleIteratorFactory : public ExampleIteratorFactory { 58 public: 59 // Creates an `ExampleIteratorFactory` that can handle all queries and for 60 // which stats are collected, and delegates the creation of the 61 // `ExampleIterator` to an std::function. FunctionalExampleIteratorFactory(std::function<absl::StatusOr<std::unique_ptr<ExampleIterator>> (const google::internal::federated::plan::ExampleSelector &)> create_iterator_func)62 explicit FunctionalExampleIteratorFactory( 63 std::function<absl::StatusOr<std::unique_ptr<ExampleIterator>>( 64 const google::internal::federated::plan::ExampleSelector& 65 66 )> 67 create_iterator_func) 68 : can_handle_func_( 69 [](const google::internal::federated::plan::ExampleSelector&) { 70 return true; 71 }), 72 create_iterator_func_(create_iterator_func), 73 should_collect_stats_(true) {} 74 75 // Creates an `ExampleIteratorFactory` that delegates to an std::function to 76 // determine if a given query can be handled, and delegates the creation of 77 // the `ExampleIterator` to an std::function as well. FunctionalExampleIteratorFactory(std::function<bool (const google::internal::federated::plan::ExampleSelector &)> can_handle_func,std::function<absl::StatusOr<std::unique_ptr<ExampleIterator>> (const google::internal::federated::plan::ExampleSelector &)> create_iterator_func,bool should_collect_stats)78 FunctionalExampleIteratorFactory( 79 std::function< 80 bool(const google::internal::federated::plan::ExampleSelector&)> 81 can_handle_func, 82 std::function<absl::StatusOr<std::unique_ptr<ExampleIterator>>( 83 const google::internal::federated::plan::ExampleSelector& 84 85 )> 86 create_iterator_func, 87 bool should_collect_stats) 88 : can_handle_func_(can_handle_func), 89 create_iterator_func_(create_iterator_func), 90 should_collect_stats_(should_collect_stats) {} 91 CanHandle(const google::internal::federated::plan::ExampleSelector & example_selector)92 bool CanHandle(const google::internal::federated::plan::ExampleSelector& 93 example_selector) override { 94 return can_handle_func_(example_selector); 95 } 96 CreateExampleIterator(const google::internal::federated::plan::ExampleSelector & example_selector)97 absl::StatusOr<std::unique_ptr<ExampleIterator>> CreateExampleIterator( 98 const google::internal::federated::plan::ExampleSelector& 99 example_selector) override { 100 return create_iterator_func_(example_selector); 101 } 102 ShouldCollectStats()103 bool ShouldCollectStats() override { return should_collect_stats_; } 104 105 private: 106 std::function<bool(const google::internal::federated::plan::ExampleSelector&)> 107 can_handle_func_; 108 std::function<absl::StatusOr<std::unique_ptr<ExampleIterator>>( 109 const google::internal::federated::plan::ExampleSelector&)> 110 create_iterator_func_; 111 bool should_collect_stats_; 112 }; 113 114 } // namespace engine 115 } // namespace client 116 } // namespace fcp 117 118 #endif // FCP_CLIENT_ENGINE_EXAMPLE_ITERATOR_FACTORY_H_ 119