Skip to content

Commit 72e19f3

Browse files
authored
Add initial bigtable stub test (#1286)
* Add initial bigtable stub test Signed-off-by: Yong Tang <yong.tang.github@outlook.com> * Fix kokoro test Signed-off-by: Yong Tang <yong.tang.github@outlook.com>
1 parent 8d63a0b commit 72e19f3

File tree

5 files changed

+128
-5
lines changed

5 files changed

+128
-5
lines changed

.github/workflows/build.wheel.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -6,7 +6,7 @@ run_test() {
66
entry=$1
77
CPYTHON_VERSION=$($entry -c 'import sys; print(str(sys.version_info[0])+str(sys.version_info[1]))')
88
(cd wheelhouse && $entry -m pip install tensorflow_io-*-cp${CPYTHON_VERSION}-*.whl)
9-
$entry -m pip install -q pytest pytest-benchmark boto3 fastavro avro-python3 scikit-image pandas pyarrow==3.0.0 google-cloud-pubsub==2.1.0 google-cloud-bigquery-storage==1.1.0 google-cloud-bigquery==2.3.1 google-cloud-storage==1.32.0
9+
$entry -m pip install -q pytest pytest-benchmark boto3 fastavro avro-python3 scikit-image pandas pyarrow==3.0.0 google-cloud-pubsub==2.1.0 google-cloud-bigtable==1.6.0 google-cloud-bigquery-storage==1.1.0 google-cloud-bigquery==2.3.1 google-cloud-storage==1.32.0
1010
(cd tests && $entry -m pytest --benchmark-disable -v --import-mode=append $(find . -type f \( -iname "test_*.py" ! \( -iname "test_*_eager.py" \) \)))
1111
(cd tests && $entry -m pytest --benchmark-disable -v --import-mode=append $(find . -type f \( -iname "test_*_eager.py" ! \( -iname "test_bigquery_eager.py" \) \)))
1212
# GRPC and test_bigquery_eager tests have to be executed separately because of https://github.com/grpc/grpc/issues/20034

.github/workflows/build.yml

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -198,7 +198,7 @@ jobs:
198198
set -x -e
199199
bash -x -e tests/test_kafka/kafka_test.sh
200200
bash -x -e tests/test_azure/start_azure.sh
201-
bash -x -e tests/test_pubsub/pubsub_test.sh
201+
bash -x -e tests/test_gcloud/test_pubsub_bigtable.sh
202202
bash -x -e tests/test_pulsar/pulsar_test.sh
203203
- name: Install ${{ matrix.python }} macOS
204204
run: |
@@ -296,7 +296,7 @@ jobs:
296296
bash -x -e .github/workflows/build.space.sh
297297
bash -x -e tests/test_kafka/kafka_test.sh
298298
bash -x -e tests/test_aws/aws_test.sh
299-
bash -x -e tests/test_pubsub/pubsub_test.sh
299+
bash -x -e tests/test_gcloud/test_pubsub_bigtable.sh
300300
bash -x -e tests/test_prometheus/prometheus_test.sh start
301301
bash -x -e tests/test_elasticsearch/elasticsearch_test.sh start
302302
bash -x -e tests/test_mongodb/mongodb_test.sh start

.kokorun/io_cpu.sh

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -75,7 +75,7 @@ bash -x -e tests/test_gcloud/test_gcs.sh gcs-emulator
7575
bash -x -e tests/test_kafka/kafka_test.sh
7676
bash -x -e tests/test_pulsar/pulsar_test.sh
7777
bash -x -e tests/test_aws/aws_test.sh
78-
bash -x -e tests/test_pubsub/pubsub_test.sh pubsub
78+
bash -x -e tests/test_gcloud/test_pubsub_bigtable.sh
7979
bash -x -e tests/test_prometheus/prometheus_test.sh start
8080
bash -x -e tests/test_azure/start_azure.sh
8181
bash -x -e tests/test_sql/sql_test.sh sql

tests/test_bigtable_eager.py

Lines changed: 119 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,119 @@
1+
# Copyright 2018 The TensorFlow Authors. All Rights Reserved.
2+
#
3+
# Licensed under the Apache License, Version 2.0 (the "License"); you may not
4+
# use this file except in compliance with the License. You may obtain a copy of
5+
# 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, WITHOUT
11+
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
12+
# License for the specific language governing permissions and limitations under
13+
# the License.
14+
# ==============================================================================
15+
"""Stub Test"""
16+
17+
import os
18+
import sys
19+
import time
20+
import shutil
21+
import datetime
22+
import tempfile
23+
import numpy as np
24+
import pytest
25+
26+
import tensorflow as tf
27+
import tensorflow_io as tfio
28+
29+
30+
def bigtable_func(project_id, instance_id, table_id):
31+
from google.cloud import bigtable
32+
from google.cloud.bigtable import column_family
33+
from google.cloud.bigtable import row_filters
34+
from google.auth.credentials import AnonymousCredentials
35+
36+
os.environ["BIGTABLE_EMULATOR_HOST"] = "localhost:8086"
37+
38+
# [START bigtable_hw_connect]
39+
# The client must be created with admin=True because it will create a
40+
# table.
41+
client = bigtable.Client(
42+
project=project_id, admin=True, credentials=AnonymousCredentials()
43+
)
44+
instance = client.instance(instance_id)
45+
# [END bigtable_hw_connect]
46+
47+
# [START bigtable_hw_create_table]
48+
print("Creating the {} table.".format(table_id))
49+
table = instance.table(table_id)
50+
51+
print("Creating column family cf1 with Max Version GC rule...")
52+
# Create a column family with GC policy : most recent N versions
53+
# Define the GC policy to retain only the most recent 2 versions
54+
max_versions_rule = column_family.MaxVersionsGCRule(2)
55+
column_family_id = "cf1"
56+
column_families = {column_family_id: max_versions_rule}
57+
if not table.exists():
58+
table.create(column_families=column_families)
59+
else:
60+
print("Table {} already exists.".format(table_id))
61+
# [END bigtable_hw_create_table]
62+
63+
# [START bigtable_hw_write_rows]
64+
print("Writing some greetings to the table.")
65+
greetings = ["Hello World!", "Hello Cloud Bigtable!", "Hello Python!"]
66+
rows = []
67+
column = b"greeting"
68+
for i, value in enumerate(greetings):
69+
# Note: This example uses sequential numeric IDs for simplicity,
70+
# but this can result in poor performance in a production
71+
# application. Since rows are stored in sorted order by key,
72+
# sequential keys can result in poor distribution of operations
73+
# across nodes.
74+
#
75+
# For more information about how to design a Bigtable schema for
76+
# the best performance, see the documentation:
77+
#
78+
# https://cloud.google.com/bigtable/docs/schema-design
79+
row_key = "greeting{}".format(i).encode()
80+
row = table.direct_row(row_key)
81+
row.set_cell(
82+
column_family_id, column, value, timestamp=datetime.datetime.utcnow()
83+
)
84+
rows.append(row)
85+
table.mutate_rows(rows)
86+
# [END bigtable_hw_write_rows]
87+
88+
# [START bigtable_hw_create_filter]
89+
# Create a filter to only retrieve the most recent version of the cell
90+
# for each column accross entire row.
91+
row_filter = row_filters.CellsColumnLimitFilter(1)
92+
# [END bigtable_hw_create_filter]
93+
94+
# [START bigtable_hw_get_with_filter]
95+
print("Getting a single greeting by row key.")
96+
key = b"greeting0"
97+
98+
row = table.read_row(key, row_filter)
99+
cell = row.cells[column_family_id][column][0]
100+
print(cell.value.decode("utf-8"))
101+
# [END bigtable_hw_get_with_filter]
102+
103+
# [START bigtable_hw_scan_with_filter]
104+
print("Scanning for all greetings:")
105+
partial_rows = table.read_rows(filter_=row_filter)
106+
107+
for row in partial_rows:
108+
cell = row.cells[column_family_id][column][0]
109+
print(cell.value.decode("utf-8"))
110+
# [END bigtable_hw_scan_with_filter]
111+
112+
# [START bigtable_hw_delete_table]
113+
print("Deleting the {} table.".format(table_id))
114+
table.delete()
115+
# [END bigtable_hw_delete_table]
116+
117+
118+
def test_bigtable():
119+
bigtable_func("bigtable_project", "bigtable_instance", "bigtable_table")

tests/test_pubsub/pubsub_test.sh renamed to tests/test_gcloud/test_pubsub_bigtable.sh

Lines changed: 5 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,8 +23,10 @@ if [ "$#" -eq 1 ]; then
2323
echo pull google/cloud-sdk
2424
docker pull google/cloud-sdk:236.0.0
2525
echo pull google/cloud-sdk successfully
26-
docker run -d --rm --net=host --name=$container -v $base:/v -w /v google/cloud-sdk:236.0.0 bash -x -c 'gcloud beta emulators pubsub start'
26+
docker run -d --rm --net=host --name=$container-pubsub -v $base:/v -w /v google/cloud-sdk:236.0.0 bash -x -c 'gcloud beta emulators pubsub start'
2727
echo wait 10 secs until pubsub is up and running
28+
docker run -d --rm --net=host --name=$container-bigtable -v $base:/v -w /v google/cloud-sdk:236.0.0 bash -x -c 'gcloud beta emulators bigtable start'
29+
echo wait 10 secs until bigtable is up and running
2830
sleep 10
2931
exit 0
3032
fi
@@ -34,7 +36,9 @@ tar -xzf google-cloud-sdk-236.0.0-darwin-x86_64.tar.gz
3436
google-cloud-sdk/install.sh -q
3537
google-cloud-sdk/bin/gcloud -q components install beta
3638
google-cloud-sdk/bin/gcloud -q components install pubsub-emulator
39+
google-cloud-sdk/bin/gcloud -q components update beta
3740
google-cloud-sdk/bin/gcloud -q beta emulators pubsub start &
41+
google-cloud-sdk/bin/gcloud -q beta emulators bigtable start &
3842
exit 0
3943

4044

0 commit comments

Comments
 (0)