fix table rebalance issue

faust/stores/base.py

@@ -68,10 +68,10 @@ async def need_active_standby_for(self, tp: TP) -> bool:
 
     async def on_rebalance(
         self,
-        table: CollectionT,
         assigned: Set[TP],
         revoked: Set[TP],
         newly_assigned: Set[TP],
+        generation_id: int = 0,
     ) -> None:
         """Handle rebalancing of the cluster."""
         ...

faust/stores/rocksdb.py

@@ -329,7 +329,6 @@ def _del(self, key: bytes) -> None:
 
     async def on_rebalance(
         self,
-        table: CollectionT,
         assigned: Set[TP],
         revoked: Set[TP],
         newly_assigned: Set[TP],
@@ -338,16 +337,16 @@ async def on_rebalance(
         """Rebalance occurred.
 
         Arguments:
-            table: The table that we store data for.
             assigned: Set of all assigned topic partitions.
             revoked: Set of newly revoked topic partitions.
             newly_assigned: Set of newly assigned topic partitions,
                 for which we were not assigned the last time.
+            generation_id: the metadata generation identifier for the re-balance
         """
         self.rebalance_ack = False
         async with self.db_lock:
-            self.revoke_partitions(table, revoked)
-            await self.assign_partitions(table, newly_assigned, generation_id)
+            self.revoke_partitions(self.table, revoked)
+            await self.assign_partitions(self.table, newly_assigned, generation_id)
 
     async def stop(self) -> None:
         self.logger.info("Closing rocksdb on stop")

faust/types/stores.py

@@ -89,7 +89,6 @@ def reset_state(self) -> None:
     @abc.abstractmethod
     async def on_rebalance(
         self,
-        table: _CollectionT,
         assigned: Set[TP],
         revoked: Set[TP],
         newly_assigned: Set[TP],

tests/unit/stores/test_rocksdb.py

@@ -323,7 +323,7 @@ async def test_on_rebalance(self, *, store, table):
         newly_assigned = {TP2}
         generation_id = 1
         await store.on_rebalance(
-            table, assigned, revoked, newly_assigned, generation_id=generation_id
+            assigned, revoked, newly_assigned, generation_id=generation_id
         )
 
         store.revoke_partitions.assert_called_once_with(table, revoked)