def add_subscriber(self, subscriber_data: SubscriberData): """ Add the subscriber to store. """ sid = SIDUtils.to_str(subscriber_data.sid) data_str = subscriber_data.SerializeToString() db_location = self._db_locations[self._sid2bucket(sid)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (sid, ), ) if res.fetchone(): raise DuplicateSubscriberError(sid) conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str), ) finally: conn.close() self._on_ready.add_subscriber(subscriber_data)
def resync(self, subscribers): """ Method that should resync the store with the mentioned list of subscribers. The resync leaves the current state of subscribers intact. Args: subscribers - list of subscribers to be in the store. """ with self.conn: # Capture the current state of the subscribers res = self.conn.execute( "SELECT subscriber_id, data FROM subscriberdb") current_state = {} for row in res: sub = SubscriberData() sub.ParseFromString(row[1]) current_state[row[0]] = sub.state # Clear all subscribers self.conn.execute("DELETE FROM subscriberdb") # Add the subscribers with the current state for sub in subscribers: sid = SIDUtils.to_str(sub.sid) if sid in current_state: sub.state.CopyFrom(current_state[sid]) data_str = sub.SerializeToString() self.conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str)) self._on_ready.resync(subscribers)
def edit_subscriber(self, subscriber_id): """ Context manager to modify the subscriber data. """ db_location = self._db_locations[self._sid2bucket(subscriber_id)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, ), ) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) yield subscriber_data data_str = subscriber_data.SerializeToString() conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, subscriber_id), ) finally: conn.close()
def upsert_subscriber(self, subscriber_data: SubscriberData) -> None: """ Check if the given subscriber exists in store. If so, update subscriber data; otherwise, add subscriber. """ sid = SIDUtils.to_str(subscriber_data.sid) data_str = subscriber_data.SerializeToString() db_location = self._db_locations[self._sid2bucket(sid)] conn = sqlite3.connect(db_location, uri=True) try: with conn: res = conn.execute( "SELECT subscriber_id FROM subscriberdb WHERE " "subscriber_id = ?", (sid, ), ) row = res.fetchone() if row is None: conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str), ) else: conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, sid), ) finally: conn.close() self._on_ready.upsert_subscriber(subscriber_data)
def edit_subscriber(self, subscriber_id): """ Context manager to modify the subscriber data. """ start = time.clock_gettime(time.CLOCK_THREAD_CPUTIME_ID) start2 = time.clock_gettime(time.CLOCK_MONOTONIC) with self.conn: res = self.conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, ), ) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) yield subscriber_data data_str = subscriber_data.SerializeToString() self.conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, subscriber_id), ) end = time.clock_gettime(time.CLOCK_THREAD_CPUTIME_ID) end2 = time.clock_gettime(time.CLOCK_MONOTONIC) logging.warning('LTE edit sub spends: {} ms'.format( (end - start) * 1e3)) logging.warning('LTE edit sub takes: {} ms'.format( (end2 - start2) * 1e3))
def edit_subscriber(self, subscriber_id): """ Context manager to modify the subscriber data. """ with self.conn: res = self.conn.execute( "SELECT data FROM subscriberdb WHERE " "subscriber_id = ?", (subscriber_id, )) row = res.fetchone() if not row: raise SubscriberNotFoundError(subscriber_id) subscriber_data = SubscriberData() subscriber_data.ParseFromString(row[0]) yield subscriber_data data_str = subscriber_data.SerializeToString() self.conn.execute( "UPDATE subscriberdb SET data = ? " "WHERE subscriber_id = ?", (data_str, subscriber_id))
def resync(self, subscribers): """ Method that should resync the store with the mentioned list of subscribers. The resync leaves the current state of subscribers intact. Args: subscribers - list of subscribers to be in the store. """ bucket_subs = defaultdict(list) for sub in subscribers: sid = SIDUtils.to_str(sub.sid) bucket_subs[self._sid2bucket(sid)].append(sub) for i, db_location in enumerate(self._db_locations): conn = sqlite3.connect(db_location, uri=True) try: with conn: # Capture the current state of the subscribers res = conn.execute( "SELECT subscriber_id, data FROM subscriberdb", ) current_state = {} for row in res: sub = SubscriberData() sub.ParseFromString(row[1]) current_state[row[0]] = sub.state # Clear all subscribers conn.execute("DELETE FROM subscriberdb") # Add the subscribers with the current state for sub in bucket_subs[i]: sid = SIDUtils.to_str(sub.sid) if sid in current_state: sub.state.CopyFrom(current_state[sid]) data_str = sub.SerializeToString() conn.execute( "INSERT INTO subscriberdb(subscriber_id, data) " "VALUES (?, ?)", (sid, data_str), ) finally: conn.close() self._on_ready.resync(subscribers)