def example_with_statement():
# DataAPI supports with statement for handling transaction
with DataAPI(database=database, resource_arn=resource_arn, secret_arn=secret_arn) as data_api:
# start transaction
insert: Insert = Insert(Pets, {'name': 'dog'})
# INSERT INTO pets (name) VALUES ('dog')
# `execute` accepts SQL statement as str or SQL Alchemy SQL objects
result: Result = data_api.execute(insert)
print(result.number_of_records_updated)
# 1
query = Query(Pets).filter(Pets.id == 1)
result: Result = data_api.execute(query) # or data_api.execute('select id, name from pets')
# SELECT pets.id, pets.name FROM pets WHERE pets.id = 1
# `Result` like a Result object in SQL Alchemy
print(result.scalar())
# 1
print(result.one())
# [Record<id=1, name='dog'>]
# `Result` is Sequence[Record]
records: List[Record] = list(result)
print(records)
# [Record<id=1, name='dog'>]
# Record is Sequence and Iterator
record = records[0]
print(record[0])
# 1
print(record[1])
# dog
for column in record:
print(column)
# 1 ...
# show record as dict()
print(record.dict())
# {'id': 1, 'name': 'dog'}
# batch insert
insert: Insert = Insert(Pets)
data_api.batch_execute(insert, [
{'id': 2, 'name': 'cat'},
{'id': 3, 'name': 'snake'},
{'id': 4, 'name': 'rabbit'},
])
result = data_api.execute('select * from pets')
print(list(result))
# [Record<id=1, name='dog'>, Record<id=2, name='cat'>, Record<id=3, name='snake'>, Record<id=4, name='rabbit'>]
# result is a sequence object
for record in result:
print(record)
def patch_sqlite_on_conflict_do_nothing():
# H/t https://stackoverflow.com/a/64902371/51685
from sqlalchemy.ext.compiler import compiles
from sqlalchemy.sql import Insert
@compiles(Insert, "sqlite")
def suffix_insert(insert, compiler, **kwargs):
stmt = compiler.visit_insert(insert, **kwargs)
if insert.dialect_kwargs.get("sqlite_on_conflict_do_nothing"):
stmt += " ON CONFLICT DO NOTHING"
return stmt
Insert.argument_for("sqlite", "on_conflict_do_nothing", False)
def test_with_statement(rds_data_client, db_connection):
with DataAPI(
database=database,
resource_arn=resource_arn,
secret_arn=secret_arn,
client=rds_data_client,
) as data_api:
insert: Insert = Insert(Pets, {'name': 'dog'})
result = data_api.execute(insert)
assert result.number_of_records_updated == 1
query = Query(Pets).filter(Pets.id == 1)
result = data_api.execute(query)
assert list(result) == [Record([1, 'dog'], [])]
result = data_api.execute('select * from pets')
assert result.one().dict() == {'id': 1, 'name': 'dog'}
insert: Insert = Insert(Pets)
data_api.batch_execute(
insert,
[
{
'id': 2,
'name': 'cat'
},
{
'id': 3,
'name': 'snake'
},
{
'id': 4,
'name': 'rabbit'
},
],
)
result = data_api.execute('select * from pets')
expected = [
Record([1, 'dog'], ['id', 'name']),
Record([2, 'cat'], ['id', 'name']),
Record([3, 'snake'], ['id', 'name']),
Record([4, 'rabbit'], ['id', 'name']),
]
assert list(result) == expected
for row, expected_row in zip(result, expected):
assert row == expected_row
def get_on_conflict_stmt(
stmt: Insert,
index: Any,
args: Any,
where: Any = None,
action: ConflictAction = ConflictAction.DO_UPDATE,
) -> Insert:
values = {attr: getattr(stmt.excluded, attr) for attr in args}
if hasattr(stmt.table.c, "last_seen_at") and "last_seen_at" not in values:
values["last_seen_at"] = datetime.utcnow()
if not args or action == ConflictAction.DO_NOTHING:
return stmt.on_conflict_do_nothing(index_elements=index)
return stmt.on_conflict_do_update(index_elements=index, set_=values, where=where)
def test_rollback_with_custom_exception(db_connection):
rds_data_client = boto3.client(
'rds-data',
endpoint_url='http://127.0.0.1:8080',
aws_access_key_id='aaa',
aws_secret_access_key='bbb',
)
class OriginalError(Exception):
pass
class OtherError(Exception):
pass
try:
with DataAPI(
resource_arn=resource_arn,
secret_arn=secret_arn,
rollback_exception=OriginalError,
database=database,
client=rds_data_client,
) as data_api:
data_api.execute(Insert(Pets, {'name': 'dog'}))
raise OriginalError # rollback
except:
pass
result = list(db_connection.execute('select * from pets'))
assert result == []
try:
with DataAPI(
resource_arn=resource_arn,
secret_arn=secret_arn,
rollback_exception=OriginalError,
database=database,
client=rds_data_client,
) as data_api:
data_api.execute(Insert(Pets, {'name': 'dog'}))
raise OtherError
except:
pass
result = list(get_connection().execute('select * from pets'))
assert result == [(2, 'dog')]
def example_rollback_with_custom_exception():
class OriginalError(Exception):
pass
with DataAPI(resource_arn=resource_arn, secret_arn=secret_arn, rollback_exception=rollback_exception=OriginalError) as data_api:
data_api.execute(Insert(Pets, {'name': 'dog'}))
# some logic ...
# rollback when happen `rollback_exception`
raise OriginalError # rollback
def test_rollback(rds_data_client, db_connection):
try:
with DataAPI(resource_arn, secret_arn) as data_api:
data_api.execute(Insert(Pets, {'name': 'dog'}))
# you can rollback by Exception
raise Exception
except:
pass
result = list(db_connection.execute('select * from pets'))
assert result == []
def test_generate_sql() -> None:
class Users(declarative_base()):
__tablename__ = 'users'
id = Column(Integer, primary_key=True, autoincrement=True)
name = Column(String(255, collation='utf8_unicode_ci'), default=None)
insert: Insert = Insert(Users, {'name': 'ken'})
assert generate_sql(insert) == "INSERT INTO users (name) VALUES ('ken')"
assert (generate_sql(
Query(Users).filter(Users.id == 1)) == "SELECT users.id, users.name \n"
"FROM users \n"
"WHERE users.id = 1")
def add_pet(data_api: DataAPI, pet_names: List[str]) -> None:
response = data_api.execute(Insert(Pets, {'name': pet_names[0]}))
assert response.generated_fields_first == 1
response = data_api.execute(Insert(Pets, {'name': pet_names[1]}))
assert response.generated_fields_first == 2
def add_pets(data_api: DataAPI, pet_names: List[str]) -> None:
# start transaction
for pet_name in pet_names:
data_api.execute(Insert(Pets, {'name': pet_name}))
def example_rollback():
with DataAPI(resource_arn=resource_arn, secret_arn=secret_arn) as data_api:
data_api.execute(Insert(Pets, {'name': 'dog'}))
# you can rollback by Exception
raise Exception