feat: add sample for pre-split feature (#1333)

* feat: add sample for pre-split feature

* build error fixes

* build failure fixes

* build fixes

* lint fixes

* fixes lint

* fixed the build error

* fixed the build error

* chore: fix positional argument issue

Signed-off-by: Sri Harsha CH <[email protected]>

* fixed the index test case

* added comment on the splits for idex keys

* fixed indent

* lint fixes

* lint fixes

* chore: tests fix

Signed-off-by: Sri Harsha CH <[email protected]>

* chore: update sample to not change editions due to failing test case

Signed-off-by: Sri Harsha CH <[email protected]>

---------

Signed-off-by: Sri Harsha CH <[email protected]>
Co-authored-by: Sri Harsha CH <[email protected]>
Co-authored-by: Sri Harsha CH <[email protected]>

Author: aksharauke

samples/samples/snippets.py

@@ -33,6 +33,7 @@
 from google.cloud.spanner_v1 import DirectedReadOptions, param_types
 from google.cloud.spanner_v1.data_types import JsonObject
 from google.protobuf import field_mask_pb2  # type: ignore
+from google.protobuf import struct_pb2  # type: ignore
 
 from testdata import singer_pb2
 
@@ -90,7 +91,7 @@ def update_instance(instance_id):
             labels={
                 "sample_name": "snippets-update_instance-explicit",
             },
-            edition=spanner_instance_admin.Instance.Edition.ENTERPRISE,  # Optional
+            edition=spanner_instance_admin.Instance.Edition.STANDARD,  # Optional
         ),
         field_mask=field_mask_pb2.FieldMask(paths=["labels", "edition"]),
     )
@@ -3204,6 +3205,7 @@ def create_instance_with_autoscaling_config(instance_id):
                 "sample_name": "snippets-create_instance_with_autoscaling_config",
                 "created": str(int(time.time())),
             },
+            edition=spanner_instance_admin.Instance.Edition.ENTERPRISE,  # Optional
         ),
     )
 
@@ -3509,6 +3511,90 @@ def query_data_with_proto_types_parameter(instance_id, database_id):
     # [END spanner_query_with_proto_types_parameter]
 
 
+# [START spanner_database_add_split_points]
+def add_split_points(instance_id, database_id):
+    """Adds split points to table and index."""
+
+    from google.cloud.spanner_admin_database_v1.types import spanner_database_admin
+
+    spanner_client = spanner.Client()
+    database_admin_api = spanner_client.database_admin_api
+
+    request = spanner_database_admin.UpdateDatabaseDdlRequest(
+        database=database_admin_api.database_path(
+            spanner_client.project, instance_id, database_id
+        ),
+        statements=["CREATE INDEX IF NOT EXISTS SingersByFirstLastName ON Singers(FirstName, LastName)"],
+    )
+
+    operation = database_admin_api.update_database_ddl(request)
+
+    print("Waiting for operation to complete...")
+    operation.result(OPERATION_TIMEOUT_SECONDS)
+
+    print("Added the SingersByFirstLastName index.")
+
+    addSplitPointRequest = spanner_database_admin.AddSplitPointsRequest(
+        database=database_admin_api.database_path(
+            spanner_client.project, instance_id, database_id
+        ),
+        # Table split
+        # Index split without table key part
+        # Index split with table key part: first key is the index key and second the table key
+        split_points=[
+            spanner_database_admin.SplitPoints(
+                table="Singers",
+                keys=[
+                    spanner_database_admin.SplitPoints.Key(
+                        key_parts=struct_pb2.ListValue(
+                            values=[struct_pb2.Value(string_value="42")]
+                        )
+                    )
+                ],
+            ),
+            spanner_database_admin.SplitPoints(
+                index="SingersByFirstLastName",
+                keys=[
+                    spanner_database_admin.SplitPoints.Key(
+                        key_parts=struct_pb2.ListValue(
+                            values=[
+                                struct_pb2.Value(string_value="John"),
+                                struct_pb2.Value(string_value="Doe"),
+                            ]
+                        )
+                    )
+                ],
+            ),
+            spanner_database_admin.SplitPoints(
+                index="SingersByFirstLastName",
+                keys=[
+                    spanner_database_admin.SplitPoints.Key(
+                        key_parts=struct_pb2.ListValue(
+                            values=[
+                                struct_pb2.Value(string_value="Jane"),
+                                struct_pb2.Value(string_value="Doe"),
+                            ]
+                        )
+                    ),
+                    spanner_database_admin.SplitPoints.Key(
+                        key_parts=struct_pb2.ListValue(
+                            values=[struct_pb2.Value(string_value="38")]
+                        )
+                    ),
+
+                ],
+            ),
+        ],
+    )
+
+    operation = database_admin_api.add_split_points(addSplitPointRequest)
+
+    print("Added split points.")
+
+
+# [END spanner_database_add_split_points]
+
+
 if __name__ == "__main__":  # noqa: C901
     parser = argparse.ArgumentParser(
         description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter
@@ -3666,6 +3752,10 @@ def query_data_with_proto_types_parameter(instance_id, database_id):
         "query_data_with_proto_types_parameter",
         help=query_data_with_proto_types_parameter.__doc__,
     )
+    subparsers.add_parser(
+        "add_split_points",
+        help=add_split_points.__doc__,
+    )
 
     args = parser.parse_args()
 
@@ -3815,3 +3905,5 @@ def query_data_with_proto_types_parameter(instance_id, database_id):
         update_data_with_proto_types_with_dml(args.instance_id, args.database_id)
     elif args.command == "query_data_with_proto_types_parameter":
         query_data_with_proto_types_parameter(args.instance_id, args.database_id)
+    elif args.command == "add_split_points":
+        add_split_points(args.instance_id, args.database_id)

samples/samples/snippets_test.py

@@ -1009,3 +1009,10 @@ def test_query_data_with_proto_types_parameter(
     )
     out, _ = capsys.readouterr()
     assert "SingerId: 2, SingerInfo: singer_id: 2" in out
+
+
+@pytest.mark.dependency(name="add_split_points", depends=["insert_data"])
+def test_add_split_points(capsys, instance_id, sample_database):
+    snippets.add_split_points(instance_id, sample_database.database_id)
+    out, _ = capsys.readouterr()
+    assert "Added split points." in out