Auto partitioning settings to public api

vgvoleg · vgvoleg · commit 7c607f5a1aee · 2025-01-31T10:49:43.000+03:00
diff --git a/ydb/_grpc/grpcwrapper/ydb_topic.py b/ydb/_grpc/grpcwrapper/ydb_topic.py
@@ -967,30 +967,184 @@ def from_public(alter_consumer: ydb_topic_public_types.PublicAlterConsumer) -> A
 class PartitioningSettings(IToProto, IFromProto):
     min_active_partitions: int
     partition_count_limit: int
+    max_active_partitions: int
+    auto_partitioning_settings: AutoPartitioningSettings
 
     @staticmethod
     def from_proto(msg: ydb_topic_pb2.PartitioningSettings) -> "PartitioningSettings":
         return PartitioningSettings(
             min_active_partitions=msg.min_active_partitions,
             partition_count_limit=msg.partition_count_limit,
+            max_active_partitions=msg.max_active_partitions,
+            auto_partitioning_settings=AutoPartitioningSettings.from_proto(
+                msg.auto_partitioning_settings
+            ),
         )
 
     def to_proto(self) -> ydb_topic_pb2.PartitioningSettings:
         return ydb_topic_pb2.PartitioningSettings(
             min_active_partitions=self.min_active_partitions,
             partition_count_limit=self.partition_count_limit,
+            max_active_partitions=self.max_active_partitions,
+            auto_partitioning_settings=self.auto_partitioning_settings.to_proto()
+        )
+
+
+class AutoPartitioningStrategy(int, IFromProto, IFromPublic, IToPublic):
+    UNSPECIFIED = 0
+    DISABLED = 1
+    SCALE_UP = 2
+    SCALE_UP_AND_DOWN = 3
+    PAUSED = 4
+
+    @staticmethod
+    def from_public(
+        strategy: Optional[ydb_topic_public_types.PublicAutoPartitioningStrategy],
+    ) -> Optional["AutoPartitioningStrategy"]:
+        if strategy is None:
+            return None
+
+        return AutoPartitioningStrategy(strategy)
+
+    @staticmethod
+    def from_proto(code: Optional[int]) -> Optional["AutoPartitioningStrategy"]:
+        if code is None:
+            return None
+
+        return AutoPartitioningStrategy(code)
+
+    def to_public(self) -> ydb_topic_public_types.PublicAutoPartitioningStrategy:
+        try:
+            ydb_topic_public_types.PublicAutoPartitioningStrategy(int(self))
+        except KeyError:
+            return ydb_topic_public_types.PublicAutoPartitioningStrategy.UNSPECIFIED
+
+
+@dataclass
+class AutoPartitioningSettings(IToProto, IFromProto, IFromPublic, IToPublic):
+    strategy: AutoPartitioningStrategy
+    partition_write_speed: AutoPartitioningWriteSpeedStrategy
+
+    @staticmethod
+    def from_public(
+        settings: Optional[ydb_topic_public_types.PublicAutoPartitioningSettings]
+    ) -> Optional[AutoPartitioningSettings]:
+        if not settings:
+            return None
+
+        return AutoPartitioningSettings(
+            strategy=settings.strategy,
+            partition_write_speed=AutoPartitioningWriteSpeedStrategy(
+                stabilization_window=settings.stabilization_window,
+                up_utilization_percent=settings.up_utilization_percent,
+                down_utilization_percent=settings.down_utilization_percent,
+            )
+        )
+
+    @staticmethod
+    def from_proto(msg: ydb_topic_pb2.AutoPartitioningSettings) -> AutoPartitioningSettings:
+        return AutoPartitioningSettings(
+            strategy=AutoPartitioningStrategy.from_proto(msg.strategy),
+            partition_write_speed=AutoPartitioningWriteSpeedStrategy.from_proto(
+                msg.partition_write_speed
+            ),
+        )
+
+    def to_proto(self) -> ydb_topic_pb2.AutoPartitioningSettings:
+        return ydb_topic_pb2.AutoPartitioningSettings(
+            strategy=self.strategy,
+            partition_write_speed=self.partition_write_speed.to_proto()
+        )
+
+    def to_public(self) -> ydb_topic_public_types.PublicAutoPartitioningSettings:
+        return ydb_topic_public_types.PublicAutoPartitioningSettings(
+            strategy=self.strategy.to_public(),
+            stabilization_window=self.partition_write_speed.stabilization_window,
+            up_utilization_percent=self.partition_write_speed.up_utilization_percent,
+            down_utilization_percent=self.partition_write_speed.down_utilization_percent,
+        )
+
+
+@dataclass
+class AutoPartitioningWriteSpeedStrategy(IToProto, IFromProto):
+    stabilization_window: Optional[datetime.timedelta]
+    up_utilization_percent: Optional[int]
+    down_utilization_percent: Optional[int]
+
+    def to_proto(self):
+        return ydb_topic_pb2.AutoPartitioningWriteSpeedStrategy(
+            stabilization_window=proto_duration_from_timedelta(self.stabilization_window),
+            up_utilization_percent=self.up_utilization_percent,
+            down_utilization_percent=self.down_utilization_percent,
+        )
+
+    @staticmethod
+    def from_proto(msg: ydb_topic_pb2.AutoPartitioningWriteSpeedStrategy) -> AutoPartitioningWriteSpeedStrategy:
+        return AutoPartitioningWriteSpeedStrategy(
+            stabilization_window=timedelta_from_proto_duration(msg.stabilization_window),
+            up_utilization_percent=msg.up_utilization_percent,
+            down_utilization_percent=msg.down_utilization_percent,
         )
 
 
 @dataclass
 class AlterPartitioningSettings(IToProto):
     set_min_active_partitions: Optional[int]
     set_partition_count_limit: Optional[int]
+    set_max_active_partitions: Optional[int]
+    alter_auto_partitioning_settings: Optional[AlterAutoPartitioningSettings]
 
     def to_proto(self) -> ydb_topic_pb2.AlterPartitioningSettings:
         return ydb_topic_pb2.AlterPartitioningSettings(
             set_min_active_partitions=self.set_min_active_partitions,
             set_partition_count_limit=self.set_partition_count_limit,
+            set_max_active_partitions=self.set_max_active_partitions,
+        )
+
+
+@dataclass
+class AlterAutoPartitioningSettings(IToProto, IFromPublic):
+    set_strategy: Optional[AutoPartitioningStrategy]
+    set_partition_write_speed: Optional[AlterAutoPartitioningWriteSpeedStrategy]
+
+    @staticmethod
+    def from_public(
+        settings: Optional[ydb_topic_public_types.PublicAlterAutoPartitioningSettings]
+    ) -> Optional[AlterAutoPartitioningSettings]:
+        if not settings:
+            return None
+
+        return AutoPartitioningSettings(
+            strategy=settings.set_strategy,
+            partition_write_speed=AlterAutoPartitioningWriteSpeedStrategy(
+                stabilization_window=settings.set_stabilization_window,
+                up_utilization_percent=settings.set_up_utilization_percent,
+                down_utilization_percent=settings.set_down_utilization_percent,
+            )
+        )
+
+    def to_proto(self) -> ydb_topic_pb2.AlterAutoPartitioningSettings:
+        set_partition_write_speed = None
+        if self.set_partition_write_speed:
+            set_partition_write_speed = self.set_partition_write_speed.to_proto()
+
+        return ydb_topic_pb2.AlterAutoPartitioningSettings(
+            set_strategy=self.set_strategy,
+            set_partition_write_speed=set_partition_write_speed,
+        )
+
+
+@dataclass
+class AlterAutoPartitioningWriteSpeedStrategy(IToProto):
+    set_stabilization_window: Optional[datetime.timedelta]
+    set_up_utilization_percent: Optional[int]
+    set_down_utilization_percent: Optional[int]
+
+    def to_proto(self) -> ydb_topic_pb2.AlterAutoPartitioningWriteSpeedStrategy:
+        return ydb_topic_pb2.AlterAutoPartitioningWriteSpeedStrategy(
+            set_stabilization_window=proto_duration_from_timedelta(self.set_stabilization_window),
+            set_up_utilization_percent=self.set_up_utilization_percent,
+            set_down_utilization_percent=self.set_down_utilization_percent,
         )
 
 
@@ -1063,11 +1217,17 @@ def from_public(req: ydb_topic_public_types.CreateTopicRequestParams):
                     consumer = ydb_topic_public_types.PublicConsumer(name=consumer)
                 consumers.append(Consumer.from_public(consumer))
 
+        auto_partitioning_settings = None
+        if req.auto_partitioning_settings is not None:
+            auto_partitioning_settings = AutoPartitioningSettings.from_public(req.auto_partitioning_settings)
+
         return CreateTopicRequest(
             path=req.path,
             partitioning_settings=PartitioningSettings(
                 min_active_partitions=req.min_active_partitions,
                 partition_count_limit=req.partition_count_limit,
+                max_active_partitions=req.max_active_partitions,
+                auto_partitioning_settings=auto_partitioning_settings
             ),
             retention_period=req.retention_period,
             retention_storage_mb=req.retention_storage_mb,
@@ -1138,13 +1298,22 @@ def from_public(req: ydb_topic_public_types.AlterTopicRequestParams) -> AlterTop
                     consumer = ydb_topic_public_types.PublicAlterConsumer(name=consumer)
                 alter_consumers.append(AlterConsumer.from_public(consumer))
 
+        alter_auto_partitioning_settings = None
+        if req.alter_auto_partitioning_settings is not None:
+            alter_auto_partitioning_settings = AutoPartitioningSettings.from_public(
+                req.alter_auto_partitioning_settings
+            )
+
+
         drop_consumers = req.drop_consumers if req.drop_consumers else []
 
         return AlterTopicRequest(
             path=req.path,
             alter_partitioning_settings=AlterPartitioningSettings(
                 set_min_active_partitions=req.set_min_active_partitions,
                 set_partition_count_limit=req.set_partition_count_limit,
+                set_max_active_partitions=req.set_max_active_partitions,
+                alter_auto_partitioning_settings=alter_auto_partitioning_settings,
             ),
             add_consumers=add_consumers,
             set_retention_period=req.set_retention_period,
@@ -1205,6 +1374,8 @@ def to_public(self) -> ydb_topic_public_types.PublicDescribeTopicResult:
         return ydb_topic_public_types.PublicDescribeTopicResult(
             self=scheme._wrap_scheme_entry(self.self_proto),
             min_active_partitions=self.partitioning_settings.min_active_partitions,
+            max_active_partitions=self.partitioning_settings.max_active_partitions,
+            auto_partitioning_settings=self.partitioning_settings.auto_partitioning_settings.to_public(),
             partition_count_limit=self.partitioning_settings.partition_count_limit,
             partitions=list(map(DescribeTopicResult.PartitionInfo.to_public, self.partitions)),
             retention_period=self.retention_period,
diff --git a/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py b/ydb/_grpc/grpcwrapper/ydb_topic_public_types.py
@@ -19,6 +19,7 @@
 class CreateTopicRequestParams:
     path: str
     min_active_partitions: Optional[int]
+    max_active_partitions: Optional[int]
     partition_count_limit: Optional[int]
     retention_period: Optional[datetime.timedelta]
     retention_storage_mb: Optional[int]
@@ -28,12 +29,14 @@ class CreateTopicRequestParams:
     attributes: Optional[Dict[str, str]]
     consumers: Optional[List[Union["PublicConsumer", str]]]
     metering_mode: Optional["PublicMeteringMode"]
+    auto_partitioning_settings: Optional["PublicAutoPartitioningSettings"]
 
 
 @dataclass
 class AlterTopicRequestParams:
     path: str
     set_min_active_partitions: Optional[int]
+    set_max_active_partitions: Optional[int]
     set_partition_count_limit: Optional[int]
     add_consumers: Optional[List[Union["PublicConsumer", str]]]
     alter_consumers: Optional[List[Union["PublicAlterConsumer", str]]]
@@ -45,6 +48,7 @@ class AlterTopicRequestParams:
     set_retention_period: Optional[datetime.timedelta]
     set_retention_storage_mb: Optional[int]
     set_supported_codecs: Optional[List[Union["PublicCodec", int]]]
+    alter_auto_partitioning_settings: Optional["PublicAlterAutoPartitioningSettings"]
 
 
 class PublicCodec(int):
@@ -68,6 +72,30 @@ class PublicMeteringMode(IntEnum):
     REQUEST_UNITS = 2
 
 
+class PublicAutoPartitioningStrategy(IntEnum):
+    UNSPECIFIED = 0
+    DISABLED = 1
+    SCALE_UP = 2
+    SCALE_UP_AND_DOWN = 3
+    PAUSED = 4
+
+
+@dataclass
+class PublicAutoPartitioningSettings:
+    strategy: Optional["PublicAutoPartitioningStrategy"]
+    stabilization_window: Optional[datetime.timedelta]
+    up_utilization_percent: Optional[int]
+    down_utilization_percent: Optional[int]
+
+
+@dataclass
+class PublicAlterAutoPartitioningSettings:
+    set_strategy: Optional["PublicAutoPartitioningStrategy"]
+    set_stabilization_window: Optional[datetime.timedelta]
+    set_up_utilization_percent: Optional[int]
+    set_down_utilization_percent: Optional[int]
+
+
 @dataclass
 class PublicConsumer:
     name: str
@@ -138,6 +166,9 @@ class PublicDescribeTopicResult:
     min_active_partitions: int
     "Minimum partition count auto merge would stop working at"
 
+    max_active_partitions: int
+    "Minimum partition count auto split would stop working at"
+
     partition_count_limit: int
     "Limit for total partition count, including active (open for write) and read-only partitions"
 
@@ -171,6 +202,8 @@ class PublicDescribeTopicResult:
     topic_stats: "PublicDescribeTopicResult.TopicStats"
     "Statistics of topic"
 
+    auto_partitioning_settings: "PublicAutoPartitioningSettings"
+
     @dataclass
     class PartitionInfo:
         partition_id: int
diff --git a/ydb/_topic_reader/topic_reader_asyncio.py b/ydb/_topic_reader/topic_reader_asyncio.py
@@ -582,7 +582,7 @@ def _on_partition_session_stop(self, message: StreamReadMessage.StopPartitionSes
             )
 
     def _on_end_partition_session(self, message: StreamReadMessage.EndPartitionSession):
-        logger.debug(f"End partition session with id: {message.partition_session_id}")
+        logger.info(f"End partition session with id: {message.partition_session_id}, child partitions: {message.child_partition_ids}")
 
     def _on_read_response(self, message: StreamReadMessage.ReadResponse):
         self._buffer_consume_bytes(message.bytes_size)
diff --git a/ydb/topic.py b/ydb/topic.py
@@ -80,6 +80,9 @@
     PublicConsumer as TopicConsumer,
     PublicAlterConsumer as TopicAlterConsumer,
     PublicMeteringMode as TopicMeteringMode,
+    PublicAutoPartitioningStrategy as TopicAutoPartitioningStrategy,
+    PublicAutoPartitioningSettings as TopicAutoPartitioningSettings,
+    PublicAlterAutoPartitioningSettings as TopicAlterAutoPartitioningSettings,
 )
 
 
@@ -117,6 +120,7 @@ async def create_topic(
         attributes: Optional[Dict[str, str]] = None,
         consumers: Optional[List[Union[TopicConsumer, str]]] = None,
         metering_mode: Optional[TopicMeteringMode] = None,
+        auto_partitioning_settings: Optional[TopicAutoPartitioningSettings] = None,
     ):
         """
         create topic command
@@ -151,6 +155,7 @@ async def alter_topic(
         self,
         path: str,
         set_min_active_partitions: Optional[int] = None,
+        set_max_active_partitions: Optional[int] = None,
         set_partition_count_limit: Optional[int] = None,
         add_consumers: Optional[List[Union[TopicConsumer, str]]] = None,
         alter_consumers: Optional[List[Union[TopicAlterConsumer, str]]] = None,
@@ -162,6 +167,7 @@ async def alter_topic(
         set_retention_period: Optional[datetime.timedelta] = None,
         set_retention_storage_mb: Optional[int] = None,
         set_supported_codecs: Optional[List[Union[TopicCodec, int]]] = None,
+        alter_auto_partitioning_settings: Optional[TopicAlterAutoPartitioningSettings] = None,
     ):
         """
         alter topic command
@@ -226,6 +232,7 @@ def reader(
         # custom decoder executor for call builtin and custom decoders. If None - use shared executor pool.
         # if max_worker in the executor is 1 - then decoders will be called from the thread without parallel
         decoder_executor: Optional[concurrent.futures.Executor] = None,
+        auto_partitioning_support: Optional[bool] = False,  # Auto partitioning feature flag. Default - False.
     ) -> TopicReaderAsyncIO:
 
         if not decoder_executor:
@@ -314,6 +321,7 @@ def create_topic(
         attributes: Optional[Dict[str, str]] = None,
         consumers: Optional[List[Union[TopicConsumer, str]]] = None,
         metering_mode: Optional[TopicMeteringMode] = None,
+        auto_partitioning_settings: Optional[TopicAutoPartitioningSettings] = None,
     ):
         """
         create topic command
@@ -350,6 +358,7 @@ def alter_topic(
         self,
         path: str,
         set_min_active_partitions: Optional[int] = None,
+        set_max_active_partitions: Optional[int] = None,
         set_partition_count_limit: Optional[int] = None,
         add_consumers: Optional[List[Union[TopicConsumer, str]]] = None,
         alter_consumers: Optional[List[Union[TopicAlterConsumer, str]]] = None,
@@ -361,6 +370,7 @@ def alter_topic(
         set_retention_period: Optional[datetime.timedelta] = None,
         set_retention_storage_mb: Optional[int] = None,
         set_supported_codecs: Optional[List[Union[TopicCodec, int]]] = None,
+        alter_auto_partitioning_settings: Optional[TopicAlterAutoPartitioningSettings] = None,
     ):
         """
         alter topic command
@@ -431,6 +441,7 @@ def reader(
         # custom decoder executor for call builtin and custom decoders. If None - use shared executor pool.
         # if max_worker in the executor is 1 - then decoders will be called from the thread without parallel
         decoder_executor: Optional[concurrent.futures.Executor] = None,  # default shared client executor pool
+        auto_partitioning_support: Optional[bool] = False,  # Auto partitioning feature flag. Default - False.
     ) -> TopicReader:
         if not decoder_executor:
             decoder_executor = self._executor

Original file line number	Diff line number	Diff line change
`@@ -582,7 +582,7 @@ def _on_partition_session_stop(self, message: StreamReadMessage.StopPartitionSes`
`582`	`582`	`)`
`583`	`583`
`584`	`584`	`def _on_end_partition_session(self, message: StreamReadMessage.EndPartitionSession):`
`585`		`- logger.debug(f"End partition session with id: {message.partition_session_id}")`
	`585`	`+ logger.info(f"End partition session with id: {message.partition_session_id}, child partitions: {message.child_partition_ids}")`
`586`	`586`
`587`	`587`	`def _on_read_response(self, message: StreamReadMessage.ReadResponse):`
`588`	`588`	`self._buffer_consume_bytes(message.bytes_size)`