feat: Add snippets for Spanner DML with returning clause (#811)

Samples are provided for INSERT, DELETE, and UPDATE in both GoogleSQL and PostgreSQL dialects. To provide a more compelling example for the INSERT case, a generated column has been added in the "create_database" example so that the generated value can be returned in the INSERT examples.

Author: c2nes

samples/samples/pg_snippets.py

@@ -95,6 +95,8 @@ def create_table_using_ddl(database_name):
  FirstName character varying(1024),
  LastName character varying(1024),
  SingerInfo bytea,
+ FullName character varying(2048)
+ GENERATED ALWAYS AS (FirstName || ' ' || LastName) STORED,
  PRIMARY KEY (SingerId)
  )""",
  """CREATE TABLE Albums (
@@ -539,6 +541,38 @@ def insert_singers(transaction):
  # [END spanner_postgresql_dml_getting_started_insert]
 
 
+def insert_with_dml_returning(instance_id, database_id):
+ """Inserts sample data into the given database using a DML statement having a RETURNING clause. """
+ # [START spanner_postgresql_dml_insert_returning]
+ # instance_id = "your-spanner-instance"
+ # database_id = "your-spanner-db-id"
+
+ spanner_client = spanner.Client()
+ instance = spanner_client.instance(instance_id)
+ database = instance.database(database_id)
+
+ # Insert records into the SINGERS table and returns the
+ # generated column FullName of the inserted records using
+ # 'RETURNING FullName'.
+ # It is also possible to return all columns of all the
+ # inserted records by using 'RETURNING *'.
+ def insert_singers(transaction):
+ results = transaction.execute_sql(
+ "INSERT INTO Singers (SingerId, FirstName, LastName) VALUES "
+ "(21, 'Luann', 'Chizoba'), "
+ "(22, 'Denis', 'Patricio'), "
+ "(23, 'Felxi', 'Ronan'), "
+ "(24, 'Dominik', 'Martyna') "
+ "RETURNING FullName"
+ )
+ for result in results:
+ print("FullName: {}".format(*result))
+ print("{} record(s) inserted.".format(results.stats.row_count_exact))
+
+ database.run_in_transaction(insert_singers)
+ # [END spanner_postgresql_dml_insert_returning]
+
+
 def query_data_with_parameter(instance_id, database_id):
  """Queries sample data from the database using SQL with a parameter."""
  # [START spanner_postgresql_query_with_parameter]
@@ -852,6 +886,37 @@ def update_albums(transaction):
  # [END spanner_postgresql_dml_standard_update]
 
 
+def update_data_with_dml_returning(instance_id, database_id):
+ """Updates sample data from the database using a DML statement having a RETURNING clause."""
+ # [START spanner_postgresql_dml_update_returning]
+ # instance_id = "your-spanner-instance"
+ # database_id = "your-spanner-db-id"
+
+ spanner_client = spanner.Client()
+ instance = spanner_client.instance(instance_id)
+ database = instance.database(database_id)
+
+ # Update MarketingBudget column for records satisfying
+ # a particular condition and returns the modified
+ # MarketingBudget column of the updated records using
+ # 'RETURNING MarketingBudget'.
+ # It is also possible to return all columns of all the
+ # updated records by using 'RETURNING *'.
+ def update_albums(transaction):
+ results = transaction.execute_sql(
+ "UPDATE Albums "
+ "SET MarketingBudget = MarketingBudget * 2 "
+ "WHERE SingerId = 1 and AlbumId = 1 "
+ "RETURNING MarketingBudget"
+ )
+ for result in results:
+ print("MarketingBudget: {}".format(*result))
+ print("{} record(s) updated.".format(results.stats.row_count_exact))
+
+ database.run_in_transaction(update_albums)
+ # [END spanner_postgresql_dml_update_returning]
+
+
 def delete_data_with_dml(instance_id, database_id):
  """Deletes sample data from the database using a DML statement."""
  # [START spanner_postgresql_dml_standard_delete]
@@ -873,6 +938,35 @@ def delete_singers(transaction):
  # [END spanner_postgresql_dml_standard_delete]
 
 
+def delete_data_with_dml_returning(instance_id, database_id):
+ """Deletes sample data from the database using a DML statement having a RETURNING clause. """
+ # [START spanner_postgresql_dml_delete_returning]
+ # instance_id = "your-spanner-instance"
+ # database_id = "your-spanner-db-id"
+
+ spanner_client = spanner.Client()
+ instance = spanner_client.instance(instance_id)
+ database = instance.database(database_id)
+
+ # Delete records from SINGERS table satisfying a
+ # particular condition and returns the SingerId
+ # and FullName column of the deleted records using
+ # 'RETURNING SingerId, FullName'.
+ # It is also possible to return all columns of all the
+ # deleted records by using 'RETURNING *'.
+ def delete_singers(transaction):
+ results = transaction.execute_sql(
+ "DELETE FROM Singers WHERE FirstName = 'David' "
+ "RETURNING SingerId, FullName"
+ )
+ for result in results:
+ print("SingerId: {}, FullName: {}".format(*result))
+ print("{} record(s) deleted.".format(results.stats.row_count_exact))
+
+ database.run_in_transaction(delete_singers)
+ # [END spanner_postgresql_dml_delete_returning]
+
+
 def dml_write_read_transaction(instance_id, database_id):
  """First inserts data then reads it from within a transaction using DML."""
  # [START spanner_postgresql_dml_write_then_read]
@@ -1522,12 +1616,17 @@ def query_data_with_jsonb_parameter(instance_id, database_id):
  help=insert_data_with_dml.__doc__)
  subparsers.add_parser("update_data_with_dml",
  help=update_data_with_dml.__doc__)
+ subparsers.add_parser("update_data_with_dml",
+ help=update_data_with_dml_returning.__doc__)
  subparsers.add_parser("delete_data_with_dml",
  help=delete_data_with_dml.__doc__)
+ subparsers.add_parser("delete_data_with_dml_returning",
+ help=delete_data_with_dml_returning.__doc__)
  subparsers.add_parser(
  "dml_write_read_transaction", help=dml_write_read_transaction.__doc__
  )
  subparsers.add_parser("insert_with_dml", help=insert_with_dml.__doc__)
+ subparsers.add_parser("insert_with_dml_returning", help=insert_with_dml_returning.__doc__)
  subparsers.add_parser(
  "query_data_with_parameter", help=query_data_with_parameter.__doc__
  )
@@ -1628,12 +1727,18 @@ def query_data_with_jsonb_parameter(instance_id, database_id):
  insert_data_with_dml(args.instance_id, args.database_id)
  elif args.command == "update_data_with_dml":
  update_data_with_dml(args.instance_id, args.database_id)
+ elif args.command == "update_data_with_dml_returning":
+ update_data_with_dml_returning(args.instance_id, args.database_id)
  elif args.command == "delete_data_with_dml":
  delete_data_with_dml(args.instance_id, args.database_id)
+ elif args.command == "delete_data_with_dml_returning":
+ delete_data_with_dml_returning(args.instance_id, args.database_id)
  elif args.command == "dml_write_read_transaction":
  dml_write_read_transaction(args.instance_id, args.database_id)
  elif args.command == "insert_with_dml":
  insert_with_dml(args.instance_id, args.database_id)
+ elif args.command == "insert_with_dml_returning":
+ insert_with_dml_returning(args.instance_id, args.database_id)
  elif args.command == "query_data_with_parameter":
  query_data_with_parameter(args.instance_id, args.database_id)
  elif args.command == "write_with_dml_transaction":

samples/samples/pg_snippets_test.py

@@ -28,6 +28,8 @@
  FirstName CHARACTER VARYING(1024),
  LastName CHARACTER VARYING(1024),
  SingerInfo BYTEA,
+ FullName CHARACTER VARYING(2048)
+ GENERATED ALWAYS AS (FirstName || ' ' || LastName) STORED,
  PRIMARY KEY (SingerId)
 )
 """
@@ -287,13 +289,27 @@ def test_update_data_with_dml(capsys, instance_id, sample_database):
  assert "1 record(s) updated." in out
 
 
+@pytest.mark.dependency(depends=["add_column"])
+def test_update_data_with_dml_returning(capsys, instance_id, sample_database):
+ snippets.update_data_with_dml_returning(instance_id, sample_database.database_id)
+ out, _ = capsys.readouterr()
+ assert "1 record(s) updated." in out
+
+
 @pytest.mark.dependency(depends=["insert_data"])
 def test_delete_data_with_dml(capsys, instance_id, sample_database):
  snippets.delete_data_with_dml(instance_id, sample_database.database_id)
  out, _ = capsys.readouterr()
  assert "1 record(s) deleted." in out
 
 
+@pytest.mark.dependency(depends=["insert_data"])
+def test_delete_data_with_dml_returning(capsys, instance_id, sample_database):
+ snippets.delete_data_with_dml_returning(instance_id, sample_database.database_id)
+ out, _ = capsys.readouterr()
+ assert "1 record(s) deleted." in out
+
+
 @pytest.mark.dependency(name="dml_write_read_transaction")
 def test_dml_write_read_transaction(capsys, instance_id, sample_database):
  snippets.dml_write_read_transaction(instance_id,
@@ -310,6 +326,13 @@ def test_insert_with_dml(capsys, instance_id, sample_database):
  assert "4 record(s) inserted" in out
 
 
+@pytest.mark.dependency(name="insert_with_dml_returning")
+def test_insert_with_dml_returning(capsys, instance_id, sample_database):
+ snippets.insert_with_dml_returning(instance_id, sample_database.database_id)
+ out, _ = capsys.readouterr()
+ assert "4 record(s) inserted" in out
+
+
 @pytest.mark.dependency(depends=["insert_with_dml"])
 def test_query_data_with_parameter(capsys, instance_id, sample_database):
  snippets.query_data_with_parameter(instance_id, sample_database.database_id)
@@ -333,12 +356,12 @@ def update_data_with_partitioned_dml(capsys, instance_id, sample_database):
  assert "3 record(s) updated" in out
 
 
-@pytest.mark.dependency(depends=["insert_with_dml"])
+@pytest.mark.dependency(depends=["insert_with_dml", "insert_with_dml_returning"])
 def test_delete_data_with_partitioned_dml(capsys, instance_id, sample_database):
  snippets.delete_data_with_partitioned_dml(instance_id,
  sample_database.database_id)
  out, _ = capsys.readouterr()
- assert "5 record(s) deleted" in out
+ assert "9 record(s) deleted" in out
 
 
 @pytest.mark.dependency(depends=["add_column"])

samples/samples/snippets.py

@@ -169,7 +169,10 @@ def create_database(instance_id, database_id):
  SingerId INT64 NOT NULL,
  FirstName STRING(1024),
  LastName STRING(1024),
- SingerInfo BYTES(MAX)
+ SingerInfo BYTES(MAX),
+ FullName STRING(2048) AS (
+ ARRAY_TO_STRING([FirstName, LastName], " ")
+ ) STORED
  ) PRIMARY KEY (SingerId)""",
  """CREATE TABLE Albums (
  SingerId INT64 NOT NULL,
@@ -1344,6 +1347,37 @@ def update_albums(transaction):
  # [END spanner_dml_standard_update]
 
 
+def update_data_with_dml_returning(instance_id, database_id):
+ """Updates sample data from the database using a DML statement having a THEN RETURN clause."""
+ # [START spanner_dml_update_returning]
+ # instance_id = "your-spanner-instance"
+ # database_id = "your-spanner-db-id"
+
+ spanner_client = spanner.Client()
+ instance = spanner_client.instance(instance_id)
+ database = instance.database(database_id)
+
+ # Update MarketingBudget column for records satisfying
+ # a particular condition and returns the modified
+ # MarketingBudget column of the updated records using
+ # 'THEN RETURN MarketingBudget'.
+ # It is also possible to return all columns of all the
+ # updated records by using 'THEN RETURN *'.
+ def update_albums(transaction):
+ results = transaction.execute_sql(
+ "UPDATE Albums "
+ "SET MarketingBudget = MarketingBudget * 2 "
+ "WHERE SingerId = 1 and AlbumId = 1 "
+ "THEN RETURN MarketingBudget"
+ )
+ for result in results:
+ print("MarketingBudget: {}".format(*result))
+ print("{} record(s) updated.".format(results.stats.row_count_exact))
+
+ database.run_in_transaction(update_albums)
+ # [END spanner_dml_update_returning]
+
+
 def delete_data_with_dml(instance_id, database_id):
  """Deletes sample data from the database using a DML statement."""
  # [START spanner_dml_standard_delete]
@@ -1365,6 +1399,35 @@ def delete_singers(transaction):
  # [END spanner_dml_standard_delete]
 
 
+def delete_data_with_dml_returning(instance_id, database_id):
+ """Deletes sample data from the database using a DML statement having a THEN RETURN clause. """
+ # [START spanner_dml_delete_returning]
+ # instance_id = "your-spanner-instance"
+ # database_id = "your-spanner-db-id"
+
+ spanner_client = spanner.Client()
+ instance = spanner_client.instance(instance_id)
+ database = instance.database(database_id)
+
+ # Delete records from SINGERS table satisfying a
+ # particular condition and returns the SingerId
+ # and FullName column of the deleted records using
+ # 'THEN RETURN SingerId, FullName'.
+ # It is also possible to return all columns of all the
+ # deleted records by using 'THEN RETURN *'.
+ def delete_singers(transaction):
+ results = transaction.execute_sql(
+ "DELETE FROM Singers WHERE FirstName = 'David' "
+ "THEN RETURN SingerId, FullName"
+ )
+ for result in results:
+ print("SingerId: {}, FullName: {}".format(*result))
+ print("{} record(s) deleted.".format(results.stats.row_count_exact))
+
+ database.run_in_transaction(delete_singers)
+ # [END spanner_dml_delete_returning]
+
+
 def update_data_with_dml_timestamp(instance_id, database_id):
  """Updates data with Timestamp from the database using a DML statement."""
  # [START spanner_dml_standard_update_with_timestamp]
@@ -1472,6 +1535,38 @@ def insert_singers(transaction):
  # [END spanner_dml_getting_started_insert]
 
 
+def insert_with_dml_returning(instance_id, database_id):
+ """Inserts sample data into the given database using a DML statement having a THEN RETURN clause. """
+ # [START spanner_dml_insert_returning]
+ # instance_id = "your-spanner-instance"
+ # database_id = "your-spanner-db-id"
+
+ spanner_client = spanner.Client()
+ instance = spanner_client.instance(instance_id)
+ database = instance.database(database_id)
+
+ # Insert records into the SINGERS table and returns the
+ # generated column FullName of the inserted records using
+ # 'THEN RETURN FullName'.
+ # It is also possible to return all columns of all the
+ # inserted records by using 'THEN RETURN *'.
+ def insert_singers(transaction):
+ results = transaction.execute_sql(
+ "INSERT INTO Singers (SingerId, FirstName, LastName) VALUES "
+ "(21, 'Luann', 'Chizoba'), "
+ "(22, 'Denis', 'Patricio'), "
+ "(23, 'Felxi', 'Ronan'), "
+ "(24, 'Dominik', 'Martyna') "
+ "THEN RETURN FullName"
+ )
+ for result in results:
+ print("FullName: {}".format(*result))
+ print("{} record(s) inserted.".format(results.stats.row_count_exact))
+
+ database.run_in_transaction(insert_singers)
+ # [END spanner_dml_insert_returning]
+
+
 def query_data_with_parameter(instance_id, database_id):
  """Queries sample data from the database using SQL with a parameter."""
  # [START spanner_query_with_parameter]
@@ -2273,7 +2368,9 @@ def list_instance_config_operations():
  subparsers.add_parser("insert_data_with_dml", help=insert_data_with_dml.__doc__)
  subparsers.add_parser("log_commit_stats", help=log_commit_stats.__doc__)
  subparsers.add_parser("update_data_with_dml", help=update_data_with_dml.__doc__)
+ subparsers.add_parser("update_data_with_dml_returning", help=update_data_with_dml_returning.__doc__)
  subparsers.add_parser("delete_data_with_dml", help=delete_data_with_dml.__doc__)
+ subparsers.add_parser("delete_data_with_dml_returning", help=delete_data_with_dml_returning.__doc__)
  subparsers.add_parser(
  "update_data_with_dml_timestamp", help=update_data_with_dml_timestamp.__doc__
  )
@@ -2284,6 +2381,7 @@ def list_instance_config_operations():
  "update_data_with_dml_struct", help=update_data_with_dml_struct.__doc__
  )
  subparsers.add_parser("insert_with_dml", help=insert_with_dml.__doc__)
+ subparsers.add_parser("insert_with_dml_returning", help=insert_with_dml_returning.__doc__)
  subparsers.add_parser(
  "query_data_with_parameter", help=query_data_with_parameter.__doc__
  )
@@ -2386,8 +2484,12 @@ def list_instance_config_operations():
  log_commit_stats(args.instance_id, args.database_id)
  elif args.command == "update_data_with_dml":
  update_data_with_dml(args.instance_id, args.database_id)
+ elif args.command == "update_data_with_dml_returning":
+ update_data_with_dml_returning(args.instance_id, args.database_id)
  elif args.command == "delete_data_with_dml":
  delete_data_with_dml(args.instance_id, args.database_id)
+ elif args.command == "delete_data_with_dml_returning":
+ delete_data_with_dml_returning(args.instance_id, args.database_id)
  elif args.command == "update_data_with_dml_timestamp":
  update_data_with_dml_timestamp(args.instance_id, args.database_id)
  elif args.command == "dml_write_read_transaction":
@@ -2396,6 +2498,8 @@ def list_instance_config_operations():
  update_data_with_dml_struct(args.instance_id, args.database_id)
  elif args.command == "insert_with_dml":
  insert_with_dml(args.instance_id, args.database_id)
+ elif args.command == "insert_with_dml_returning":
+ insert_with_dml_returning(args.instance_id, args.database_id)
  elif args.command == "query_data_with_parameter":
  query_data_with_parameter(args.instance_id, args.database_id)
  elif args.command == "write_with_dml_transaction":