diff --git a/samples/python/sqlalchemy2-sample/model.py b/samples/python/sqlalchemy2-sample/model.py index f49c490f3..100decd40 100644 --- a/samples/python/sqlalchemy2-sample/model.py +++ b/samples/python/sqlalchemy2-sample/model.py @@ -31,7 +31,7 @@ def format_timestamp(timestamp: datetime) -> str: BaseMixin contains properties that are common to all models in this sample. The created_at and updated_at properties are automatically filled with the - current client system time when a model is created or updated. + current client system time when a model is created or updated. """ class BaseMixin(object): __prepare_statements__ = None diff --git a/samples/python/sqlalchemy2-sample/sample.py b/samples/python/sqlalchemy2-sample/sample.py index 1685eb408..ed2ed2b4a 100644 --- a/samples/python/sqlalchemy2-sample/sample.py +++ b/samples/python/sqlalchemy2-sample/sample.py @@ -12,6 +12,7 @@ See the License for the specific language governing permissions and limitations under the License. """ +import uuid from random import random from connect import create_test_engine @@ -20,9 +21,10 @@ from model import Singer, Album, Track, Venue, Concert, TicketSale from util_random_names import random_first_name, random_last_name, \ random_album_title, random_release_date, random_marketing_budget, \ - random_cover_picture + random_cover_picture, random_boolean from uuid import uuid4 from datetime import datetime, date, timezone +from sqlalchemy.dialects.postgresql import insert # This is the default engine that is connected to PostgreSQL (PGAdapter). # This engine will by default use read/write transactions. @@ -83,6 +85,14 @@ def run_sample(): print("Album was found using a stale read, even though it has already " "been deleted.") + # Randomly either insert or update a new Singer record using an + # INSERT .. ON CONFLICT UPDATE .. statement. + if random_boolean(): + singer_id = session.query(Singer).first().id + else: + singer_id = uuid.uuid4() + insert_or_update_singer(singer_id, random_first_name(), random_last_name()) + print() print("Finished running sample") @@ -327,6 +337,39 @@ def update_singer(singer_id, first_name, last_name): session.commit() +# Inserts-or-updates a singer record in the database. This generates an +# INSERT ... ON CONFLICT (id) UPDATE ... statement. +def insert_or_update_singer(singer_id, first_name, last_name): + with Session(engine) as session: + singer = Singer() + singer.id = singer_id + singer.first_name = first_name + singer.last_name = last_name + insert_stmt = insert(Singer).values(id=singer.id, + first_name=singer.first_name, + last_name=singer.last_name, + created_at=datetime.now(), + updated_at=datetime.now(), + version_id=1) + # Spanner requires that the ON CONFLICT clause specifies the conflict + # column(s), and that all columns that are included in the INSERT statement + # are also included in the UPDATE statement. The UPDATE statement may only + # contain 'my_col=excluded.my_col' assignments. + do_update_stmt = insert_stmt.on_conflict_do_update( + index_elements=['id'], + set_=dict(id=insert_stmt.excluded.id, + first_name=insert_stmt.excluded.first_name, + last_name=insert_stmt.excluded.last_name, + created_at=insert_stmt.excluded.created_at, + updated_at=insert_stmt.excluded.updated_at, + version_id=insert_stmt.excluded.version_id), + ) + session.execute(do_update_stmt) + print("Inserted or updated singer {} with last name {}" + .format(singer.id, singer.last_name)) + session.commit() + + # Loads the given album from the database and prints its properties, including # the tracks related to this album. The table `tracks` is interleaved in # `albums`. This is handled as a normal relationship in SQLAlchemy. diff --git a/samples/python/sqlalchemy2-sample/util_random_names.py b/samples/python/sqlalchemy2-sample/util_random_names.py index f8c65d393..f0a9b4bc1 100644 --- a/samples/python/sqlalchemy2-sample/util_random_names.py +++ b/samples/python/sqlalchemy2-sample/util_random_names.py @@ -14,7 +14,7 @@ """ import secrets -from random import seed, randrange, random +from random import choice, seed, randrange, random from datetime import date seed() @@ -24,6 +24,11 @@ Helper functions for generating random names and titles. """ + +def random_boolean(): + return choice([True, False]) + + def random_first_name(): return first_names[randrange(len(first_names))]