feat: Support seek subscription in AdminClient (#176)

Seek subscription performs an out-of-band seek for a subscription to a specified target, which may be a backlog location, publish timestamp or event timestamp.

google/cloud/pubsublite/__init__.py

@@ -62,6 +62,9 @@
 from google.cloud.pubsublite_v1.types.admin import ListTopicSubscriptionsResponse
 from google.cloud.pubsublite_v1.types.admin import ListTopicsRequest
 from google.cloud.pubsublite_v1.types.admin import ListTopicsResponse
+from google.cloud.pubsublite_v1.types.admin import OperationMetadata
+from google.cloud.pubsublite_v1.types.admin import SeekSubscriptionRequest
+from google.cloud.pubsublite_v1.types.admin import SeekSubscriptionResponse
 from google.cloud.pubsublite_v1.types.admin import TopicPartitions
 from google.cloud.pubsublite_v1.types.admin import UpdateSubscriptionRequest
 from google.cloud.pubsublite_v1.types.admin import UpdateTopicRequest
@@ -70,6 +73,7 @@
 from google.cloud.pubsublite_v1.types.common import PubSubMessage
 from google.cloud.pubsublite_v1.types.common import SequencedMessage
 from google.cloud.pubsublite_v1.types.common import Subscription
+from google.cloud.pubsublite_v1.types.common import TimeTarget
 from google.cloud.pubsublite_v1.types.common import Topic
 from google.cloud.pubsublite_v1.types.cursor import CommitCursorRequest
 from google.cloud.pubsublite_v1.types.cursor import CommitCursorResponse
@@ -148,6 +152,7 @@
     "MessagePublishRequest",
     "MessagePublishResponse",
     "MessageResponse",
+    "OperationMetadata",
     "PartitionAssignment",
     "PartitionAssignmentAck",
     "PartitionAssignmentRequest",
@@ -159,6 +164,8 @@
     "PublishResponse",
     "PublisherServiceAsyncClient",
     "PublisherServiceClient",
+    "SeekSubscriptionRequest",
+    "SeekSubscriptionResponse",
     "SeekRequest",
     "SeekResponse",
     "SequencedCommitCursorRequest",
@@ -171,6 +178,7 @@
     "SubscriberServiceAsyncClient",
     "SubscriberServiceClient",
     "Subscription",
+    "TimeTarget",
     "Topic",
     "TopicPartitions",
     "TopicStatsServiceAsyncClient",

google/cloud/pubsublite/admin_client.py

@@ -12,10 +12,11 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import Optional, List
+from typing import Optional, List, Union
 
 from overrides import overrides
 from google.api_core.client_options import ClientOptions
+from google.api_core.operation import Operation
 from google.auth.credentials import Credentials
 from google.protobuf.field_mask_pb2 import FieldMask
 
@@ -31,6 +32,8 @@
     LocationPath,
     TopicPath,
     BacklogLocation,
+    PublishTime,
+    EventTime,
 )
 from google.cloud.pubsublite.types.paths import ReservationPath
 from google.cloud.pubsublite_v1 import (
@@ -129,6 +132,14 @@ def update_subscription(
     ) -> Subscription:
         return self._impl.update_subscription(subscription, update_mask)
 
+    @overrides
+    def seek_subscription(
+        self,
+        subscription_path: SubscriptionPath,
+        target: Union[BacklogLocation, PublishTime, EventTime],
+    ) -> Operation:
+        return self._impl.seek_subscription(subscription_path, target)
+
     @overrides
     def delete_subscription(self, subscription_path: SubscriptionPath):
         return self._impl.delete_subscription(subscription_path)

google/cloud/pubsublite/admin_client_interface.py

@@ -13,14 +13,17 @@
 # limitations under the License.
 
 from abc import ABC, abstractmethod
-from typing import List
+from typing import List, Union
 
+from google.api_core.operation import Operation
 from google.cloud.pubsublite.types import (
     CloudRegion,
     TopicPath,
     LocationPath,
     SubscriptionPath,
     BacklogLocation,
+    PublishTime,
+    EventTime,
 )
 from google.cloud.pubsublite.types.paths import ReservationPath
 from google.cloud.pubsublite_v1 import Topic, Subscription, Reservation
@@ -88,6 +91,24 @@ def update_subscription(
     ) -> Subscription:
         """Update the masked fields of the provided subscription."""
 
+    @abstractmethod
+    def seek_subscription(
+        self,
+        subscription_path: SubscriptionPath,
+        target: Union[BacklogLocation, PublishTime, EventTime],
+    ) -> Operation:
+        """Initiate an out-of-band seek for a subscription to a specified target.
+
+        The seek target may be timestamps or named positions within the message
+        backlog See https://cloud.google.com/pubsub/lite/docs/seek for more
+        information.
+
+        Returns:
+            google.api_core.operation.Operation with:
+              result type: google.cloud.pubsublite.SeekSubscriptionResponse
+              metadata type: google.cloud.pubsublite.OperationMetadata
+        """
+
     @abstractmethod
     def delete_subscription(self, subscription_path: SubscriptionPath):
         """Delete a subscription and all associated messages."""

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

@@ -12,8 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from typing import List
+from typing import List, Union
 
+from google.api_core.exceptions import InvalidArgument
+from google.api_core.operation import Operation
 from google.protobuf.field_mask_pb2 import FieldMask
 
 from google.cloud.pubsublite.admin_client_interface import AdminClientInterface
@@ -23,6 +25,8 @@
     LocationPath,
     TopicPath,
     BacklogLocation,
+    PublishTime,
+    EventTime,
 )
 from google.cloud.pubsublite.types.paths import ReservationPath
 from google.cloud.pubsublite_v1 import (
@@ -31,6 +35,8 @@
     AdminServiceClient,
     TopicPartitions,
     Reservation,
+    TimeTarget,
+    SeekSubscriptionRequest,
 )
 
 
@@ -105,6 +111,25 @@ def update_subscription(
             subscription=subscription, update_mask=update_mask
         )
 
+    def seek_subscription(
+        self,
+        subscription_path: SubscriptionPath,
+        target: Union[BacklogLocation, PublishTime, EventTime],
+    ) -> Operation:
+        request = SeekSubscriptionRequest(name=str(subscription_path))
+        if isinstance(target, PublishTime):
+            request.time_target = TimeTarget(publish_time=target.value)
+        elif isinstance(target, EventTime):
+            request.time_target = TimeTarget(event_time=target.value)
+        elif isinstance(target, BacklogLocation):
+            if target == BacklogLocation.END:
+                request.named_target = SeekSubscriptionRequest.NamedTarget.HEAD
+            else:
+                request.named_target = SeekSubscriptionRequest.NamedTarget.TAIL
+        else:
+            raise InvalidArgument("A valid seek target must be specified.")
+        return self._underlying.seek_subscription(request=request)
+
     def delete_subscription(self, subscription_path: SubscriptionPath):
         self._underlying.delete_subscription(name=str(subscription_path))
 

google/cloud/pubsublite/types/__init__.py

@@ -18,7 +18,7 @@
 from .paths import LocationPath, TopicPath, SubscriptionPath
 from .message_metadata import MessageMetadata
 from .flow_control_settings import FlowControlSettings, DISABLED_FLOW_CONTROL
-from .backlog_location import BacklogLocation
+from .backlog_location import BacklogLocation, PublishTime, EventTime
 
 __all__ = (
     "CloudRegion",
@@ -30,4 +30,6 @@
     "SubscriptionPath",
     "TopicPath",
     "BacklogLocation",
+    "PublishTime",
+    "EventTime",
 )

google/cloud/pubsublite/types/backlog_location.py

@@ -13,6 +13,8 @@
 # limitations under the License.
 
 import enum
+from datetime import datetime
+from typing import NamedTuple
 
 
 class BacklogLocation(enum.Enum):
@@ -22,3 +24,15 @@ class BacklogLocation(enum.Enum):
 
     BEGINNING = 0
     END = 1
+
+
+class PublishTime(NamedTuple):
+    """The publish timestamp of a message."""
+
+    value: datetime
+
+
+class EventTime(NamedTuple):
+    """A user-defined event timestamp of a message."""
+
+    value: datetime

samples/snippets/quickstart_test.py

@@ -22,6 +22,7 @@
 from google.api_core.exceptions import NotFound
 from google.cloud.pubsublite import AdminClient
 from google.cloud.pubsublite.types import (
+    BacklogLocation,
     CloudRegion,
     CloudZone,
     SubscriptionPath,
@@ -259,6 +260,16 @@ def test_subscriber_example(topic_path, subscription_path, capsys):
         assert f"Received {message}" in out
 
 
+def test_seek_lite_subscription_example(capsys):
+    import seek_lite_subscription_example
+
+    seek_lite_subscription_example.seek_lite_subscription(
+        PROJECT_NUMBER, CLOUD_REGION, ZONE_ID, SUBSCRIPTION_ID, BacklogLocation.BEGINNING, False
+    )
+    out, _ = capsys.readouterr()
+    assert "Seek operation" in out
+
+
 def test_delete_lite_subscription_example(subscription_path, capsys):
     import delete_lite_subscription_example
 

samples/snippets/seek_lite_subscription_example.py

@@ -0,0 +1,101 @@
+#!/usr/bin/env python
+
+# Copyright 2021 Google Inc. All Rights Reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""This application demonstrates how to initiate an out-of-band seek for a
+subscription with the Pub/Sub Lite API. For more information, see the
+documentation at https://cloud.google.com/pubsub/lite/docs/seek.
+"""
+
+import argparse
+
+
+def seek_lite_subscription(project_number, cloud_region, zone_id, subscription_id, seek_target, wait_for_operation):
+    # [START pubsublite_seek_subscription]
+    from google.api_core.exceptions import NotFound
+    from google.cloud.pubsublite import AdminClient
+    from google.cloud.pubsublite.types import CloudRegion, CloudZone, SubscriptionPath
+
+    # TODO(developer):
+    # project_number = 1122334455
+    # cloud_region = "us-central1"
+    # zone_id = "a"
+    # subscription_id = "your-subscription-id"
+    # seek_target = BacklogLocation.BEGINNING
+    # wait_for_operation = False
+
+    # Possible values for seek_target:
+    # - BacklogLocation.BEGINNING: replays from the beginning of all retained
+    #   messages.
+    # - BacklogLocation.END: skips past all current published messages.
+    # - PublishTime(<datetime>): delivers messages with publish time greater
+    #   than or equal to the specified timestamp.
+    # - EventTime(<datetime>): seeks to the first message with event time
+    #   greater than or equal to the specified timestamp.
+
+    # Waiting for the seek operation to complete is optional. It indicates when
+    # subscribers for all partitions are receiving messages from the seek
+    # target. If subscribers are offline, the operation will complete once they
+    # are online.
+
+    cloud_region = CloudRegion(cloud_region)
+    location = CloudZone(cloud_region, zone_id)
+    subscription_path = SubscriptionPath(project_number, location, subscription_id)
+
+    client = AdminClient(cloud_region)
+    try:
+        # Initiate an out-of-band seek for a subscription to the specified
+        # target. If an operation is returned, the seek has been successfully
+        # registered and will eventually propagate to subscribers.
+        seek_operation = client.seek_subscription(subscription_path, seek_target)
+        print(f"Seek operation: {seek_operation.operation.name}")
+    except NotFound:
+        print(f"{subscription_path} not found.")
+        return
+
+    if wait_for_operation:
+        print("Waiting for operation to complete...")
+        seek_operation.result()
+        print(f"Operation completed. Metadata:\n{seek_operation.metadata}")
+    # [END pubsublite_seek_subscription]
+
+
+if __name__ == "__main__":
+    from datetime import datetime
+    from google.cloud.pubsublite.types import BacklogLocation, PublishTime
+
+    parser = argparse.ArgumentParser(
+        description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter,
+    )
+    parser.add_argument("project_number", help="Your Google Cloud Project Number")
+    parser.add_argument("cloud_region", help="Your Cloud Region, e.g. 'us-central1'")
+    parser.add_argument("zone_id", help="Your Zone ID, e.g. 'a'")
+    parser.add_argument("subscription_id", help="Your subscription ID")
+    parser.add_argument("--target", default="BEGINNING", help="Seek target, e.g. 'BEGINNING, 'END' or a timestamp")
+    parser.add_argument("--wait_for_operation", help="Wait for the seek operation to complete")
+
+    args = parser.parse_args()
+
+    if args.target == "BEGINNING":
+        seek_target = BacklogLocation.BEGINNING
+    elif args.target == "END":
+        seek_target = BacklogLocation.END
+    else:
+        seek_target = PublishTime(datetime.strptime(args.target, "%Y-%m-%d %H:%M:%S"))
+
+    seek_lite_subscription(
+        args.project_number, args.cloud_region, args.zone_id,
+        args.subscription_id, seek_target, args.wait_for_operation
+    )