Bigtable Row Set snippets (googleapis#7016)

* add conditional row snippets

* add snippets

* Add row set snippets

* remove unwanted changes in snippets.py

* Rectify snippets file name and tags

* blacken

Author: sangramql

bigtable/docs/snippets_table.py

@@ -42,7 +42,6 @@
 INSTANCE_ID = "snippet-" + unique_resource_id("-")
 CLUSTER_ID = "clus-1-" + unique_resource_id("-")
 TABLE_ID = "tabl-1-" + unique_resource_id("-")
-COLUMN_FAMILY_ID = "col_fam_id-" + unique_resource_id("-")
 LOCATION_ID = "us-central1-f"
 ALT_LOCATION_ID = "us-central1-a"
 PRODUCTION = enums.Instance.Type.PRODUCTION
@@ -55,8 +54,14 @@
     .strftime("%Y-%m-%dt%H-%M-%S")
 )
 LABELS = {LABEL_KEY: str(LABEL_STAMP)}
+COLUMN_FAMILY_ID = "col_fam_id1"
 COL_NAME1 = b"col-name1"
 CELL_VAL1 = b"cell-val"
+ROW_KEY1 = b"row_key_id1"
+COLUMN_FAMILY_ID2 = "col_fam_id2"
+COL_NAME2 = b"col-name2"
+CELL_VAL2 = b"cell-val2"
+ROW_KEY2 = b"row_key_id2"
 
 
 class Config(object):
@@ -90,6 +95,9 @@ def setup_module():
     gc_rule = column_family.MaxVersionsGCRule(2)
     column_family1 = Config.TABLE.column_family(COLUMN_FAMILY_ID, gc_rule=gc_rule)
     column_family1.create()
+    gc_rule2 = column_family.MaxVersionsGCRule(4)
+    column_family2 = Config.TABLE.column_family(COLUMN_FAMILY_ID2, gc_rule=gc_rule2)
+    column_family2.create()
 
 
 def teardown_module():
@@ -403,5 +411,80 @@ def test_bigtable_table_row():
     table.truncate(timeout=300)
 
 
+def test_bigtable_add_row_add_row_range_add_row_range_from_keys():
+    row_keys = [
+        b"row_key_1",
+        b"row_key_2",
+        b"row_key_3",
+        b"row_key_4",
+        b"row_key_5",
+        b"row_key_6",
+        b"row_key_7",
+        b"row_key_8",
+        b"row_key_9",
+    ]
+
+    rows = []
+    for row_key in row_keys:
+        row = Config.TABLE.row(row_key)
+        row.set_cell(COLUMN_FAMILY_ID, COL_NAME1, CELL_VAL1)
+        rows.append(row)
+    Config.TABLE.mutate_rows(rows)
+
+    # [START bigtable_add_row_key]
+    from google.cloud.bigtable import Client
+    from google.cloud.bigtable.row_set import RowSet
+
+    client = Client(admin=True)
+    instance = client.instance(INSTANCE_ID)
+    table = instance.table(TABLE_ID)
+
+    row_set = RowSet()
+    row_set.add_row_key(b"row_key_5")
+    # [END bigtable_add_row_key]
+
+    read_rows = table.read_rows(row_set=row_set)
+    expected_row_keys = [b"row_key_5"]
+    found_row_keys = [row.row_key for row in read_rows]
+    assert found_row_keys == expected_row_keys
+
+    # [START bigtable_add_row_range]
+    from google.cloud.bigtable import Client
+    from google.cloud.bigtable.row_set import RowSet
+    from google.cloud.bigtable.row_set import RowRange
+
+    client = Client(admin=True)
+    instance = client.instance(INSTANCE_ID)
+    table = instance.table(TABLE_ID)
+
+    row_set = RowSet()
+    row_set.add_row_range(RowRange(start_key=b"row_key_3", end_key=b"row_key_7"))
+    # [END bigtable_add_row_range]
+
+    read_rows = table.read_rows(row_set=row_set)
+    expected_row_keys = [b"row_key_3", b"row_key_4", b"row_key_5", b"row_key_6"]
+    found_row_keys = [row.row_key for row in read_rows]
+    assert found_row_keys == expected_row_keys
+
+    # [START bigtable_row_range_from_keys]
+    from google.cloud.bigtable import Client
+    from google.cloud.bigtable.row_set import RowSet
+
+    client = Client(admin=True)
+    instance = client.instance(INSTANCE_ID)
+    table = instance.table(TABLE_ID)
+
+    row_set = RowSet()
+    row_set.add_row_range_from_keys(start_key=b"row_key_3", end_key=b"row_key_7")
+    # [END bigtable_row_range_from_keys]
+
+    read_rows = table.read_rows(row_set=row_set)
+    expected_row_keys = [b"row_key_3", b"row_key_4", b"row_key_5", b"row_key_6"]
+    found_row_keys = [row.row_key for row in read_rows]
+    assert found_row_keys == expected_row_keys
+
+    table.truncate(timeout=200)
+
+
 if __name__ == "__main__":
     pytest.main()

bigtable/google/cloud/bigtable/row_set.py

@@ -53,6 +53,12 @@ def __ne__(self, other):
     def add_row_key(self, row_key):
         """Add row key to row_keys list.
 
+        For example:
+
+        .. literalinclude:: snippets_table.py
+            :start-after: [START bigtable_add_row_key]
+            :end-before: [END bigtable_add_row_key]
+
         :type row_key: bytes
         :param row_key: The key of a row to read
         """
@@ -61,6 +67,12 @@ def add_row_key(self, row_key):
     def add_row_range(self, row_range):
         """Add row_range to row_ranges list.
 
+        For example:
+
+        .. literalinclude:: snippets_table.py
+            :start-after: [START bigtable_add_row_range]
+            :end-before: [END bigtable_add_row_range]
+
         :type row_range: class:`RowRange`
         :param row_range: The row range object having start and end key
         """
@@ -71,6 +83,12 @@ def add_row_range_from_keys(
     ):
         """Add row range to row_ranges list from the row keys
 
+        For example:
+
+        .. literalinclude:: snippets_table.py
+            :start-after: [START bigtable_row_range_from_keys]
+            :end-before: [END bigtable_row_range_from_keys]
+
         :type start_key: bytes
         :param start_key: (Optional) Start key of the row range. If left empty,
                           will be interpreted as the empty string.