Adding Spanner STRUCT param samples

spanner/cloud-client/snippets.py

@@ -23,6 +23,7 @@
 import argparse
 
 from google.cloud import spanner
+from google.cloud.spanner_v1 import param_types
 
 
 # [START spanner_create_database]
@@ -598,6 +599,164 @@ def query_data_with_timestamp(instance_id, database_id):
 # [END spanner_query_data_with_timestamp_column]
 
 
+# [START spanner_write_data_for_struct_queries]
+def write_struct_data(instance_id, database_id):
+    """Inserts sample data that can be used to test STRUCT parameters in queries.
+    """
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    with database.batch() as batch:
+        batch.insert(
+            table='Singers',
+            columns=('SingerId', 'FirstName', 'LastName',),
+            values=[
+                (6, u'Elena', u'Campbell'),
+                (7, u'Gabriel', u'Wright'),
+                (8, u'Benjamin', u'Martinez'),
+                (9, u'Hannah', u'Harris')])
+
+    print('Inserted sample data for STRUCT queries')
+# [END spanner_write_data_for_struct_queries]
+
+
+def query_with_struct(instance_id, database_id):
+    """
+       Query a table using STRUCT parameters.
+    """
+    # [START spanner_create_struct_with_data]
+    record_type = param_types.Struct([
+            param_types.StructField('FirstName', param_types.STRING),
+            param_types.StructField('LastName', param_types.STRING)
+    ])
+    record_value = ('Elena', 'Campbell')
+    # [END spanner_create_struct_with_data]
+
+    # [START spanner_query_data_with_struct]
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+
+    database = instance.database(database_id)
+
+    with database.snapshot() as snapshot:
+        results = snapshot.execute_sql(
+            "SELECT SingerId FROM Singers WHERE "
+            "(FirstName, LastName) = @name",
+            params={'name': record_value},
+            param_types={'name': record_type})
+
+    for row in results:
+            print(u'SingerId: {}'.format(*row))
+    # [END spanner_query_data_with_struct]
+
+
+def query_with_array_of_struct(instance_id, database_id):
+    """
+       Query a table using an array of STRUCT parameters.
+    """
+    # [START spanner_create_user_defined_struct]
+    name_type = param_types.Struct([
+        param_types.StructField('FirstName', param_types.STRING),
+        param_types.StructField('LastName', param_types.STRING)
+        ])
+    # [END spanner_create_user_defined_struct]
+
+    # [START spanner_create_array_of_struct_with_data]
+    band_members = [("Elena", "Campbell"),
+                    ("Gabriel", "Wright"),
+                    ("Benjamin", "Martinez")]
+    # [END spanner_create_array_of_struct_with_data]
+
+    # [START spanner_query_data_with_array_of_struct]
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    with database.snapshot() as snapshot:
+        results = snapshot.execute_sql(
+            "SELECT SingerId FROM Singers WHERE "
+            "STRUCT<FirstName STRING, LastName STRING>"
+            "(FirstName, LastName) IN UNNEST(@names)",
+            params={'names': band_members},
+            param_types={'names': param_types.Array(name_type)})
+
+    for row in results:
+            print(u'SingerId: {}'.format(*row))
+    # [END spanner_query_data_with_array_of_struct]
+
+
+# [START spanner_field_access_on_struct_parameters]
+def query_struct_field(instance_id, database_id):
+    """
+       Query a table using field access on a STRUCT parameter.
+    """
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    name_type = param_types.Struct([
+        param_types.StructField('FirstName', param_types.STRING),
+        param_types.StructField('LastName', param_types.STRING)
+        ])
+
+    with database.snapshot() as snapshot:
+        results = snapshot.execute_sql(
+            "SELECT SingerId FROM Singers "
+            "WHERE FirstName = @name.FirstName",
+            params={'name': ("Elena", "Campbell")},
+            param_types={'name': name_type})
+
+    for row in results:
+            print(u'SingerId: {}'.format(*row))
+# [START spanner_field_access_on_struct_parameters]
+
+
+# [START spanner_field_access_on_nested_struct_parameters]
+def query_nested_struct_field(instance_id, database_id):
+    """
+       Query a table using nested field access on a STRUCT parameter.
+    """
+    spanner_client = spanner.Client()
+    instance = spanner_client.instance(instance_id)
+    database = instance.database(database_id)
+
+    song_info_type = param_types.Struct([
+        param_types.StructField('SongName', param_types.STRING),
+        param_types.StructField(
+            'ArtistNames', param_types.Array(
+                param_types.Struct([
+                     param_types.StructField(
+                         'FirstName', param_types.STRING),
+                     param_types.StructField(
+                         'LastName', param_types.STRING)
+                ])
+            )
+        )
+    ])
+
+    song_info = ('Imagination', [('Elena', 'Campbell'), ('Hannah', 'Harris')])
+
+    with database.snapshot() as snapshot:
+        results = snapshot.execute_sql(
+            "SELECT SingerId, @song_info.SongName "
+            "FROM Singers WHERE "
+            "STRUCT<FirstName STRING, LastName STRING>"
+            "(FirstName, LastName) "
+            "IN UNNEST(@song_info.ArtistNames)",
+            params={
+                'song_info': song_info
+                },
+            param_types={
+                'song_info': song_info_type
+                }
+        )
+
+    for row in results:
+            print(u'SingerId: {} SongName: {}'.format(*row))
+# [END spanner_field_access_on_nested_struct_parameters]
+
+
 if __name__ == '__main__':  # noqa: C901
     parser = argparse.ArgumentParser(
         description=__doc__,
@@ -644,6 +803,14 @@ def query_data_with_timestamp(instance_id, database_id):
         'update_data_with_timestamp', help=update_data_with_timestamp.__doc__)
     subparsers.add_parser(
         'query_data_with_timestamp', help=query_data_with_timestamp.__doc__)
+    subparsers.add_parser('write_struct_data', help=write_struct_data.__doc__)
+    subparsers.add_parser('query_with_struct', help=query_with_struct.__doc__)
+    subparsers.add_parser(
+        'query_with_array_of_struct', help=query_with_array_of_struct.__doc__)
+    subparsers.add_parser(
+            'query_struct_field', help=query_struct_field.__doc__)
+    subparsers.add_parser(
+        'query_nested_struct_field', help=query_nested_struct_field.__doc__)
 
     args = parser.parse_args()
 
@@ -689,3 +856,13 @@ def query_data_with_timestamp(instance_id, database_id):
         update_data_with_timestamp(args.instance_id, args.database_id)
     elif args.command == 'query_data_with_timestamp':
         query_data_with_timestamp(args.instance_id, args.database_id)
+    elif args.command == 'write_struct_data':
+        write_struct_data(args.instance_id, args.database_id)
+    elif args.command == 'query_with_struct':
+        query_with_struct(args.instance_id, args.database_id)
+    elif args.command == 'query_with_array_of_struct':
+        query_with_array_of_struct(args.instance_id, args.database_id)
+    elif args.command == 'query_struct_field':
+        query_struct_field(args.instance_id, args.database_id)
+    elif args.command == 'query_nested_struct_field':
+        query_nested_struct_field(args.instance_id, args.database_id)

spanner/cloud-client/snippets_test.py

@@ -234,3 +234,64 @@ def _():
         out, _ = capsys.readouterr()
 
         assert 'Go, Go, Go' in out
+
+
+@pytest.mark.slow
+def test_query_data_with_struct(temporary_database, capsys):
+    @eventually_consistent.call
+    def _():
+        snippets.write_struct_data(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        snippets.query_with_struct(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        out, _ = capsys.readouterr()
+
+        assert 'SingerId: 6' in out
+
+
+@pytest.mark.slow
+def test_query_data_with_array_struct(temporary_database, capsys):
+    @eventually_consistent.call
+    def _():
+        snippets.write_struct_data(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        snippets.query_with_array_of_struct(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        out, _ = capsys.readouterr()
+
+        assert 'SingerId: 6\nSingerId: 7' in out
+
+
+@pytest.mark.slow
+def test_query_data_with_field_struct(temporary_database, capsys):
+    @eventually_consistent.call
+    def _():
+        snippets.write_struct_data(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        snippets.query_struct_field(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        out, _ = capsys.readouterr()
+
+        assert 'SingerId: 6' in out
+
+
+@pytest.mark.slow
+def test_query_data_with_nested_field_struct(temporary_database, capsys):
+    @eventually_consistent.call
+    def _():
+        snippets.write_struct_data(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        snippets.query_nested_struct_field(
+            SPANNER_INSTANCE,
+            temporary_database.database_id)
+        out, _ = capsys.readouterr()
+
+        assert 'SingerId: 6 SongName: Imagination' in out
+        assert 'SingerId: 9 SongName: Imagination' in out