feat: Create subscriptions at a seek target (#383)

Support creating subscriptions at a nominated target location within the message backlog. A seek is performed for publish and event timestamps. Export subscriptions (pre-release feature) are also supported.

Author: tmdiep

google/cloud/pubsublite/__init__.py

@@ -70,6 +70,7 @@
 from google.cloud.pubsublite_v1.types.admin import UpdateTopicRequest
 from google.cloud.pubsublite_v1.types.common import AttributeValues
 from google.cloud.pubsublite_v1.types.common import Cursor
+from google.cloud.pubsublite_v1.types.common import ExportConfig
 from google.cloud.pubsublite_v1.types.common import PubSubMessage
 from google.cloud.pubsublite_v1.types.common import Reservation
 from google.cloud.pubsublite_v1.types.common import SequencedMessage
@@ -131,6 +132,7 @@
  "CursorServiceClient",
  "DeleteSubscriptionRequest",
  "DeleteTopicRequest",
+ "ExportConfig",
  "FlowControlRequest",
  "GetSubscriptionRequest",
  "GetTopicPartitionsRequest",

google/cloud/pubsublite/admin_client.py

@@ -114,9 +114,10 @@ def list_topic_subscriptions(self, topic_path: TopicPath) -> List[SubscriptionPa
  def create_subscription(
  self,
  subscription: Subscription,
- starting_offset: BacklogLocation = BacklogLocation.END,
+ target: Union[BacklogLocation, PublishTime, EventTime] = BacklogLocation.END,
+ starting_offset: Optional[BacklogLocation] = None,
  ) -> Subscription:
- return self._impl.create_subscription(subscription, starting_offset)
+ return self._impl.create_subscription(subscription, target, starting_offset)
 
  @overrides
  def get_subscription(self, subscription_path: SubscriptionPath) -> Subscription:

google/cloud/pubsublite/admin_client_interface.py

@@ -13,7 +13,7 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import List, Union
+from typing import List, Optional, Union
 
 from google.api_core.operation import Operation
 from google.cloud.pubsublite.types import (
@@ -71,11 +71,20 @@ def list_topic_subscriptions(self, topic_path: TopicPath) -> List[SubscriptionPa
  def create_subscription(
  self,
  subscription: Subscription,
- starting_offset: BacklogLocation = BacklogLocation.END,
+ target: Union[BacklogLocation, PublishTime, EventTime] = BacklogLocation.END,
+ starting_offset: Optional[BacklogLocation] = None,
  ) -> Subscription:
  """Create a subscription, returns the created subscription. By default
  a subscription will only receive messages published after the
- subscription was created."""
+ subscription was created.
+
+ `starting_offset` is deprecated. Use `target` to initialize the
+ subscription to a target location within the message backlog instead.
+ `starting_offset` has higher precedence if `target` is also set.
+
+ A seek is initiated if the target location is a publish or event time.
+ If the seek fails, the created subscription is not deleted.
+ """
 
  @abstractmethod
  def get_subscription(self, subscription_path: SubscriptionPath) -> Subscription:

google/cloud/pubsublite/internal/wire/admin_client_impl.py

@@ -12,7 +12,8 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List, Union
+import logging
+from typing import List, Optional, Union
 
 from google.api_core.exceptions import InvalidArgument
 from google.api_core.operation import Operation
@@ -38,8 +39,11 @@
  TimeTarget,
  SeekSubscriptionRequest,
  CreateSubscriptionRequest,
+ ExportConfig,
 )
 
+log = logging.getLogger(__name__)
+
 
 class AdminClientImpl(AdminClientInterface):
  _underlying: AdminServiceClient
@@ -85,17 +89,51 @@ def list_topic_subscriptions(self, topic_path: TopicPath) -> List[SubscriptionPa
  def create_subscription(
  self,
  subscription: Subscription,
- starting_offset: BacklogLocation = BacklogLocation.END,
+ target: Union[BacklogLocation, PublishTime, EventTime] = BacklogLocation.END,
+ starting_offset: Optional[BacklogLocation] = None,
  ) -> Subscription:
+ if starting_offset:
+ log.warning("starting_offset is deprecated. Use target instead.")
+ target = starting_offset
  path = SubscriptionPath.parse(subscription.name)
- return self._underlying.create_subscription(
+ requires_seek = isinstance(target, PublishTime) or isinstance(target, EventTime)
+ requires_update = (
+ requires_seek
+ and subscription.export_config
+ and subscription.export_config.desired_state == ExportConfig.State.ACTIVE
+ )
+ if requires_update:
+ # Export subscriptions must be paused while seeking. The state is
+ # later updated to active.
+ subscription.export_config.desired_state = ExportConfig.State.PAUSED
+
+ # Request 1 - Create the subscription.
+ skip_backlog = False
+ if isinstance(target, BacklogLocation):
+ skip_backlog = target == BacklogLocation.END
+ response = self._underlying.create_subscription(
  request=CreateSubscriptionRequest(
  parent=str(path.to_location_path()),
  subscription=subscription,
  subscription_id=path.name,
- skip_backlog=(starting_offset == BacklogLocation.END),
+ skip_backlog=skip_backlog,
  )
  )
+ # Request 2 (optional) - seek the subscription.
+ if requires_seek:
+ self.seek_subscription(subscription_path=path, target=target)
+ # Request 3 (optional) - make the export subscription active.
+ if requires_update:
+ response = self.update_subscription(
+ subscription=Subscription(
+ name=response.name,
+ export_config=ExportConfig(
+ desired_state=ExportConfig.State.ACTIVE,
+ ),
+ ),
+ update_mask=FieldMask(paths=["export_config.desired_state"]),
+ )
+ return response
 
  def get_subscription(self, subscription_path: SubscriptionPath) -> Subscription:
  return self._underlying.get_subscription(name=str(subscription_path))