[tf.data service] Add util for handling grpc errors.

PiperOrigin-RevId: 302090444
Change-Id: I8aa9c1427c55e8e956225ab0a7f97c97bb96fc32

tensorflow/core/data/service/BUILD

@@ -45,6 +45,30 @@ tf_proto_library(
     ],
 )
 
+cc_library(
+    name = "grpc_util",
+    srcs = ["grpc_util.cc"],
+    hdrs = [
+        "grpc_util.h",
+    ],
+    deps = [
+        "//tensorflow:grpc++",
+        "//tensorflow/core:lib",
+    ],
+)
+
+tf_cc_test(
+    name = "grpc_util_test",
+    srcs = ["grpc_util_test.cc"],
+    deps = [
+        ":grpc_util",
+        "//tensorflow/core:lib",
+        "//tensorflow/core:test",
+        "//tensorflow/core:test_main",
+        "//tensorflow/core:testlib",
+    ],
+)
+
 cc_library(
     name = "compression_utils",
     srcs = ["compression_utils.cc"],

tensorflow/core/data/service/grpc_util.cc

@@ -0,0 +1,37 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#include "tensorflow/core/data/service/grpc_util.h"
+
+#include "tensorflow/core/lib/core/errors.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace data {
+namespace grpc_util {
+
+Status WrapError(const std::string& message, const grpc::Status& status) {
+  if (status.ok()) {
+    return errors::Internal("Expected a non-ok grpc status. Wrapping message: ",
+                            message);
+  } else {
+    return Status(static_cast<tensorflow::error::Code>(status.error_code()),
+                  absl::StrCat(message, ": ", status.error_message()));
+  }
+}
+
+}  // namespace grpc_util
+}  // namespace data
+}  // namespace tensorflow

tensorflow/core/data/service/grpc_util.h

@@ -0,0 +1,33 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+
+#ifndef TENSORFLOW_CORE_DATA_SERVICE_GRPC_UTIL_H_
+#define TENSORFLOW_CORE_DATA_SERVICE_GRPC_UTIL_H_
+
+#include "grpcpp/grpcpp.h"
+#include "tensorflow/core/platform/status.h"
+
+namespace tensorflow {
+namespace data {
+namespace grpc_util {
+
+// Wraps a grpc::Status in a tensorflow::Status with the given message.
+Status WrapError(const std::string& message, const grpc::Status& status);
+
+}  // namespace grpc_util
+}  // namespace data
+}  // namespace tensorflow
+
+#endif  // TENSORFLOW_CORE_DATA_SERVICE_GRPC_UTIL_H_

tensorflow/core/data/service/grpc_util_test.cc

@@ -0,0 +1,39 @@
+/* Copyright 2020 The TensorFlow Authors. All Rights Reserved.
+
+Licensed under the Apache License, Version 2.0 (the "License");
+you may not use this file except in compliance with the License.
+You may obtain a copy of the License at
+
+    http://www.apache.org/licenses/LICENSE-2.0
+
+Unless required by applicable law or agreed to in writing, software
+distributed under the License is distributed on an "AS IS" BASIS,
+WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+See the License for the specific language governing permissions and
+limitations under the License.
+==============================================================================*/
+#include "tensorflow/core/data/service/grpc_util.h"
+
+#include "tensorflow/core/platform/errors.h"
+#include "tensorflow/core/platform/test.h"
+
+namespace tensorflow {
+namespace data {
+namespace grpc_util {
+
+TEST(GrpcUtil, WrapInvalidArgument) {
+  grpc::Status s(grpc::StatusCode::INVALID_ARGUMENT, "test message");
+  Status wrapped = WrapError("wrapping message", s);
+  ASSERT_EQ(wrapped, errors::InvalidArgument("wrapping message: test message"));
+}
+
+TEST(GrpcUtil, WrapOk) {
+  grpc::Status s;
+  Status wrapped = WrapError("wrapping message", s);
+  ASSERT_EQ(wrapped, errors::Internal("Expected a non-ok grpc status. Wrapping "
+                                      "message: wrapping message"));
+}
+
+}  // namespace grpc_util
+}  // namespace data
+}  // namespace tensorflow