1# Cloud Bigtable client for TensorFlow
2
3package(
4    default_visibility = ["//tensorflow:internal"],
5)
6
7licenses(["notice"])  # Apache 2.0
8
9load("//tensorflow:tensorflow.bzl", "tf_custom_op_py_library")
10load(
11    "//tensorflow:tensorflow.bzl",
12    "tf_copts",
13    "tf_custom_op_library",
14    "tf_gen_op_libs",
15    "tf_gen_op_wrapper_py",
16    "tf_kernel_library",
17    "tf_cc_test",
18    "tf_py_test",
19)
20
21tf_custom_op_py_library(
22    name = "bigtable",
23    srcs = ["__init__.py"] + glob(["python/ops/*.py"]),
24    dso = [
25        ":python/ops/_bigtable.so",
26    ],
27    kernels = [
28        ":bigtable_kernels",
29        ":bigtable_ops_op_lib",
30    ],
31    srcs_version = "PY2AND3",
32    deps = [
33        ":bigtable_ops",
34        "//tensorflow/contrib/data/python/ops:interleave_ops",
35        "//tensorflow/contrib/util:util_py",
36        "//tensorflow/python:framework_for_generated_wrappers",
37        "//tensorflow/python:platform",
38        "//tensorflow/python:util",
39        "//tensorflow/python/data",
40    ],
41)
42
43KERNEL_FILES = [
44    "kernels/bigtable_kernels.cc",
45    "kernels/bigtable_lookup_dataset_op.cc",
46    "kernels/bigtable_prefix_key_dataset_op.cc",
47    "kernels/bigtable_range_key_dataset_op.cc",
48    "kernels/bigtable_sample_keys_dataset_op.cc",
49    "kernels/bigtable_sample_key_pairs_dataset_op.cc",
50    "kernels/bigtable_scan_dataset_op.cc",
51]
52
53tf_custom_op_library(
54    name = "python/ops/_bigtable.so",
55    srcs = KERNEL_FILES + [
56        "ops/bigtable_ops.cc",
57    ],
58    deps = [
59        ":bigtable_lib_cc",
60        ":bigtable_range_helpers",
61        "@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
62    ],
63)
64
65tf_gen_op_wrapper_py(
66    name = "bigtable_ops",
67    deps = [":bigtable_ops_op_lib"],
68)
69
70tf_gen_op_libs(
71    op_lib_names = [
72        "bigtable_ops",
73        "bigtable_test_ops",
74    ],
75)
76
77tf_kernel_library(
78    name = "bigtable_kernels",
79    srcs = KERNEL_FILES,
80    deps = [
81        ":bigtable_lib_cc",
82        ":bigtable_range_helpers",
83        "//tensorflow/core:framework_headers_lib",
84        "//third_party/eigen3",
85        "@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
86    ],
87)
88
89# A library for use in the bigtable kernels.
90cc_library(
91    name = "bigtable_lib_cc",
92    srcs = ["kernels/bigtable_lib.cc"],
93    hdrs = ["kernels/bigtable_lib.h"],
94    deps = [
95        "//tensorflow/core:framework_headers_lib",
96        "//third_party/eigen3",
97        "@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
98    ],
99)
100
101cc_library(
102    name = "bigtable_range_helpers",
103    srcs = ["kernels/bigtable_range_helpers.cc"],
104    hdrs = ["kernels/bigtable_range_helpers.h"],
105    deps = [
106        "//tensorflow/core:framework_headers_lib",
107    ],
108)
109
110cc_library(
111    name = "bigtable_test_client",
112    srcs = ["kernels/test_kernels/bigtable_test_client.cc"],
113    hdrs = ["kernels/test_kernels/bigtable_test_client.h"],
114    deps = [
115        "//tensorflow/core:framework_headers_lib",
116        "@com_github_googleapis_googleapis//:bigtable_protos",
117        "@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
118        "@com_googlesource_code_re2//:re2",
119    ],
120)
121
122tf_cc_test(
123    name = "bigtable_test_client_test",
124    srcs = ["kernels/test_kernels/bigtable_test_client_test.cc"],
125    tags = ["manual"],
126    deps = [
127        ":bigtable_test_client",
128        "//tensorflow/core:test",
129        "//tensorflow/core:test_main",
130        "@com_github_googlecloudplatform_google_cloud_cpp//google/cloud/bigtable:bigtable_client",
131    ],
132)
133
134tf_cc_test(
135    name = "bigtable_range_helpers_test",
136    size = "small",
137    srcs = ["kernels/bigtable_range_helpers_test.cc"],
138    deps = [
139        ":bigtable_range_helpers",
140        "//tensorflow/core:test",
141        "//tensorflow/core:test_main",
142    ],
143)
144
145tf_gen_op_wrapper_py(
146    name = "bigtable_test_ops",
147    deps = [":bigtable_test_ops_op_lib"],
148)
149
150tf_custom_op_library(
151    name = "python/kernel_tests/_bigtable_test.so",
152    srcs = [
153        "kernels/test_kernels/bigtable_test_client_op.cc",
154        "ops/bigtable_test_ops.cc",
155    ],
156    deps = [
157        ":bigtable_lib_cc",
158        ":bigtable_test_client",
159        "@com_googlesource_code_re2//:re2",
160    ],
161)
162
163# Don't use tf_kernel_library because it prevents access to strings/stringprintf.h
164cc_library(
165    name = "bigtable_test_kernels",
166    srcs = [
167        "kernels/test_kernels/bigtable_test_client_op.cc",
168    ],
169    copts = tf_copts(),
170    linkstatic = 1,
171    deps = [
172        ":bigtable_lib_cc",
173        ":bigtable_test_client",
174        "//tensorflow/core:framework_headers_lib",
175        "//third_party/eigen3",
176        "@com_googlesource_code_re2//:re2",
177    ],
178    alwayslink = 1,
179)
180
181tf_custom_op_py_library(
182    name = "bigtable_test_py",
183    dso = [
184        ":python/kernel_tests/_bigtable_test.so",
185    ],
186    kernels = [
187        ":bigtable_test_kernels",
188        ":bigtable_test_ops_op_lib",
189    ],
190    srcs_version = "PY2AND3",
191    deps = [
192        ":bigtable_test_ops",
193    ],
194)
195
196tf_py_test(
197    name = "bigtable_ops_test",
198    size = "small",
199    srcs = ["python/kernel_tests/bigtable_ops_test.py"],
200    additional_deps = [
201        ":bigtable",
202        ":bigtable_test_py",
203        "//tensorflow/core:protos_all_py",
204        "//tensorflow/contrib/util:util_py",
205        "//tensorflow/python:array_ops",
206        "//tensorflow/python:client_testlib",
207        "//tensorflow/python:dtypes",
208        "//tensorflow/python:framework_for_generated_wrappers",
209        "//tensorflow/python:platform",
210        "//tensorflow/python:util",
211    ],
212    tags = ["manual"],
213)
214