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))