def add_rides(engine, num_rides, city):
chunk_size = 800
datagen = Faker()
def add_rides_helper(sess, chunk, n):
users = sess.query(User).filter_by(city=city).all()
vehicles = sess.query(Vehicle).filter_by(city=city).all()
rides = []
for i in range(chunk, min(chunk + chunk_size, num_rides)):
start_time = datetime.datetime.now() - datetime.timedelta(days=random.randint(0, 30))
rides.append(Ride(id=MovRGenerator.generate_uuid(),
city=city,
rider_id=random.choice(users).id,
vehicle_id=random.choice(vehicles).id,
start_time=start_time,
start_address=datagen.address(),
end_address=datagen.address(),
revenue=MovRGenerator.generate_revenue(),
end_time=start_time + datetime.timedelta(minutes=random.randint(0, 60))))
sess.bulk_save_objects(rides)
for chunk in range(0, num_rides, chunk_size):
run_transaction(sessionmaker(bind=engine),
lambda s: add_rides_helper(s, chunk, min(chunk + chunk_size, num_rides)))
def callback(barrier):
def txn_body(conn):
balances = self.get_balances(conn)
barrier()
self.perform_transfer(conn, balances)
with testing.db.connect() as conn:
run_transaction(conn, txn_body)
def update_ride_location(self, city, ride_id, lat, long):
def update_ride_location_helper(session, city, ride_id, lat, long):
h = VehicleLocationHistory(city=city,
ride_id=ride_id,
lat=lat,
long=long)
session.add(h)
run_transaction(
sessionmaker(bind=self.engine),
lambda session: update_ride_location_helper(
session, city, ride_id, lat, long))
def run_queries_in_separate_transactions(self, queries):
for query in queries:
try:
run_transaction(sessionmaker(bind=self.engine),
lambda session: session.execute(query))
except ProgrammingError as err:
if 'Duplicate' in str(err):
logging.info(
"The following query will be skipped, due to a duplicate object error:"
)
logging.info(str(query))
continue
else:
raise err
def callback(barrier):
Session = sessionmaker(testing.db)
def txn_body(session):
accounts = list(
session.query(Account).filter(Account.acct.in_(
(1, 2))).order_by(Account.acct))
barrier()
if accounts[0].balance > accounts[1].balance:
accounts[0].balance -= 100
accounts[1].balance += 100
else:
accounts[0].balance += 100
accounts[1].balance -= 100
run_transaction(Session, txn_body)
def apply_promo_code(self, user_city, user_id, promo_code):
def apply_promo_code_helper(session, user_city, user_id, code):
pc = session.query(PromoCode).filter_by(code=code).one_or_none()
if pc:
# see if it has already been applied
upc = session.query(UserPromoCode).filter_by(
user_id=user_id, code=code).one_or_none()
if not upc:
upc = UserPromoCode(city=user_city,
user_id=user_id,
code=code)
session.add(upc)
run_transaction(
sessionmaker(bind=self.engine), lambda session:
apply_promo_code_helper(session, user_city, user_id, promo_code))
def start_ride(self, city, rider_id, vehicle_id):
def start_ride_helper(session, city, rider_id, vehicle_id):
vehicle = session.query(Vehicle).filter_by(id=vehicle_id)
vehicle.update({'status': 'in_use'})
v = vehicle.first()
# get promo codes associated with this user's account
upcs = session.query(UserPromoCode).filter_by(
user_id=rider_id).all()
# determine which codes are valid
for upc in upcs:
promo_code = session.query(PromoCode).filter_by(
code=upc.code).first()
if promo_code and promo_code.expiration_time > datetime.datetime.now(
):
code_to_update = session.query(UserPromoCode).filter_by(
user_id=rider_id, code=upc.code)
code_to_update.update({'usage_count': upc.usage_count + 1})
r = Ride(city=city,
id=MovRGenerator.generate_uuid(),
rider_id=rider_id,
vehicle_id=vehicle_id,
start_address=v.current_location)
session.add(r)
return {'city': r.city, 'id': r.id}
return run_transaction(
sessionmaker(bind=self.engine), lambda session: start_ride_helper(
session, city, rider_id, vehicle_id))
def end_ride(self, city, ride_id):
def end_ride_helper(session, city, ride_id):
ride = session.query(Ride).filter_by(id=ride_id)
r = ride.first()
vehicle = session.query(Vehicle).filter_by(id=r.vehicle_id)
vehicle.update({'status': 'available'})
v = vehicle.first()
ride.update({
'end_address': v.current_location,
'revenue': MovRGenerator.generate_revenue(),
'end_time': datetime.datetime.now()
})
run_transaction(
sessionmaker(bind=self.engine),
lambda session: end_ride_helper(session, city, ride_id))
def callback(barrier):
Session = sessionmaker(testing.db)
def txn_body(session):
accounts = list(
session.query(Account).filter(Account.acct.in_((1, 2))).order_by(Account.acct)
)
barrier()
if accounts[0].balance > accounts[1].balance:
accounts[0].balance -= 100
accounts[1].balance += 100
else:
accounts[0].balance += 100
accounts[1].balance -= 100
with testing.expect_deprecated_20("The Session.autocommit parameter is deprecated"):
run_transaction(Session, txn_body)
def add_users(engine, num_users, city):
chunk_size = 1000
datagen = Faker()
def add_users_helper(sess, chunk, n):
users = []
for i in range(chunk, n):
users.append(User(id=MovRGenerator.generate_uuid(),
city=city,
name=datagen.name(),
address=datagen.address(),
credit_card=datagen.credit_card_number()))
sess.bulk_save_objects(users)
for chunk in range(0, num_users, chunk_size):
run_transaction(sessionmaker(bind=engine),
lambda s: add_users_helper(s, chunk, min(chunk + chunk_size, num_users)))
def add_promo_codes(engine, num_codes):
chunk_size = 800
datagen = Faker()
def add_codes_helper(sess, chunk, n):
codes = []
for i in range(chunk, min(chunk + chunk_size, num_codes)):
code = "_".join(datagen.words(nb=3)) + "_" + str(time.time())
codes.append(PromoCode(code=code,
description=datagen.paragraph(),
expiration_time=datetime.datetime.now(
) + datetime.timedelta(days=random.randint(0, 30)),
rules={"type": "percent_discount", "value": "10%"}))
sess.bulk_save_objects(codes)
for chunk in range(0, num_codes, chunk_size):
run_transaction(sessionmaker(bind=engine),
lambda s: add_codes_helper(s, chunk, min(chunk + chunk_size, num_codes)))
def update_region(self, table, region, cities):
region_list = self.get_regions()
region_enum = Enum(*region_list,
name='crdb_internal_region',
create_type=False,
native_enum=False)
def update_region_helper(session, table, region, cities):
crdb_region = Column('crdb_region', region_enum)
table.append_column(crdb_region)
query = table.update().where(column('city').in_(cities)).values(
{crdb_region: region})
session.execute(query)
run_transaction(
sessionmaker(bind=self.engine), lambda session:
update_region_helper(session, table, region, cities))
def test_run_transaction_retry(self):
def txn_body(sess):
rs = sess.execute(text("select acct, balance from account where acct = 1"))
sess.execute(text("select crdb_internal.force_retry('1s')"))
return [r for r in rs]
with testing.expect_deprecated_20("The Session.autocommit parameter is deprecated"):
Session = sessionmaker(testing.db)
rs = run_transaction(Session, txn_body)
assert rs[0] == (1, 100)
def test_run_transaction_retry(self):
def txn_body(conn):
rs = conn.execute(
text("select acct, balance from account where acct = 1"))
conn.execute(text("select crdb_internal.force_retry('1s')"))
return [r for r in rs]
with testing.db.connect() as conn:
rs = run_transaction(conn, txn_body)
assert rs[0] == (1, 100)
def test_run_transaction_retry(self):
def txn_body(sess):
rs = sess.execute(
text("select acct, balance from account where acct = 1"))
sess.execute(text("select crdb_internal.force_retry('1s')"))
return [r for r in rs]
Session = sessionmaker(testing.db)
rs = run_transaction(Session, txn_body)
assert rs[0] == (1, 100)
def add_vehicle_location_histories(engine, num_histories, city):
chunk_size = 5000
def add_vehicle_location_histories_helper(sess, chunk, n):
rides = sess.query(Ride).filter_by(city=city).all()
histories = []
for i in range(chunk, min(chunk + chunk_size, num_histories)):
latlong = MovRGenerator.generate_random_latlong()
histories.append(VehicleLocationHistory(
city=city,
ride_id=random.choice(rides).id,
lat=latlong["lat"],
long=latlong["long"]))
sess.bulk_save_objects(histories)
for chunk in range(0, num_histories, chunk_size):
run_transaction(sessionmaker(bind=engine),
lambda s: add_vehicle_location_histories_helper(s, chunk, min(chunk + chunk_size, num_histories)))
def add_vehicles(engine, num_vehicles, city):
chunk_size = 1000
datagen = Faker()
def add_vehicles_helper(sess, chunk, n):
owners = sess.query(User).filter_by(city=city).all()
vehicles = []
for i in range(chunk, n):
vehicle_type = MovRGenerator.generate_random_vehicle()
vehicles.append(Vehicle(id=MovRGenerator.generate_uuid(),
type=vehicle_type,
city=city,
current_location=datagen.address(),
owner_id=(random.choice(owners)).id,
status=MovRGenerator.get_vehicle_availability(),
ext=MovRGenerator.generate_vehicle_metadata(vehicle_type)))
sess.bulk_save_objects(vehicles)
for chunk in range(0, num_vehicles, chunk_size):
run_transaction(sessionmaker(bind=engine),
lambda s: add_vehicles_helper(s, chunk, min(chunk + chunk_size, num_vehicles)))
def add_user(self, city, name, address, credit_card_number):
def add_user_helper(session, city, name, address, credit_card_number):
u = User(city=city,
id=MovRGenerator.generate_uuid(),
name=name,
address=address,
credit_card=credit_card_number)
session.add(u)
return {'city': u.city, 'id': u.id}
return run_transaction(
sessionmaker(bind=self.engine), lambda session: add_user_helper(
session, city, name, address, credit_card_number))
def get_promo_codes(self, follower_reads=False, limit=None):
def get_promo_codes_helper(session, follower_reads, limit=None):
if follower_reads:
session.execute(
text(
'SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()'
))
pcs = session.query(PromoCode).limit(limit).all()
return list(map(lambda pc: pc.code, pcs))
return run_transaction(
sessionmaker(bind=self.engine), lambda session:
get_promo_codes_helper(session, follower_reads, limit))
def get_cities(self, follower_reads=False):
def get_cities_helper(session, follower_reads):
if follower_reads:
session.execute(
text(
'SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()'
))
users = self.session.query(User).distinct(User.city).all()
return tuple(user.city for user in users)
return run_transaction(
sessionmaker(bind=self.engine),
lambda session: get_cities_helper(session, follower_reads))
def get_active_rides(self, city, follower_reads=False, limit=None):
def get_active_rides_helper(session, city, follower_reads, limit=None):
if follower_reads:
session.execute(
text(
'SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()'
))
rides = session.query(Ride).filter_by(
city=city, end_time=None).limit(limit).all()
return list(map(lambda ride: {'city': city, 'id': ride.id}, rides))
return run_transaction(
sessionmaker(bind=self.engine), lambda session:
get_active_rides_helper(session, city, follower_reads, limit))
def create_promo_code(self, code, description, expiration_time, rules):
def add_promo_code_helper(session, code, description, expiration_time,
rules):
pc = PromoCode(code=code,
description=description,
expiration_time=expiration_time,
rules=rules)
session.add(pc)
return pc.code
return run_transaction(
sessionmaker(bind=self.engine),
lambda session: add_promo_code_helper(session, code, description,
expiration_time, rules))
def get_users(self, city, follower_reads=False, limit=None):
def get_users_helper(session, city, follower_reads, limit=None):
if follower_reads:
session.execute(
text(
'SET TRANSACTION AS OF SYSTEM TIME follower_read_timestamp()'
))
users = session.query(User).filter_by(city=city).limit(limit).all()
return list(
map(lambda user: {
'city': user.city,
'id': user.id
}, users))
return run_transaction(
sessionmaker(bind=self.engine), lambda session: get_users_helper(
session, city, follower_reads, limit))
def add_vehicle(self, city, owner_id, current_location, type,
vehicle_metadata, status):
def add_vehicle_helper(session, city, owner_id, current_location, type,
vehicle_metadata, status):
vehicle_type = type
vehicle = Vehicle(id=MovRGenerator.generate_uuid(),
type=vehicle_type,
city=city,
owner_id=owner_id,
current_location=current_location,
status=status,
ext=vehicle_metadata)
session.add(vehicle)
return {'city': vehicle.city, 'id': vehicle.id}
return run_transaction(
sessionmaker(bind=self.engine), lambda session: add_vehicle_helper(
session, city, owner_id, current_location, type,
vehicle_metadata, status))
