Add server-qualified job references

README.md

@@ -63,6 +63,18 @@ with client.using("some_tube") as inserter:
 client.put_job_into("some_tube", "some message")
 ```
 
+When a caller needs to identify the inserted job later, use the additive reference-returning variants. A
+`JobReference` includes a server ID because integer job IDs are only unique within one beanstalkd instance:
+
+```python
+reference = client.put_job_into_with_reference("some_tube", "some message")
+serialized_reference = reference.to_dict()
+```
+
+The server ID defaults to `host:port`. Pass a stable `server_id` to `BeanstalkClient` or `from_uri` when producers and
+consumers use different aliases for the same server. Do not use a load-balanced endpoint for references unless it maps
+one-to-one to a beanstalkd server.
+
 ### Consuming All Available Jobs
 
 The following script will walk through all currently-READY jobs and then exit:

docs/source/CHANGES.rst

@@ -2,6 +2,11 @@
 pystalk ChangeLog
 #################
 
+======
+0.10.0
+======
+* Add server-qualified job references without changing existing insertion APIs.
+
 =====
 0.9.1
 =====

pystalk/__init__.py

@@ -1,7 +1,7 @@
-from .client import BeanstalkClient, BeanstalkError, BeanstalkConnectionError
+from .client import BeanstalkClient, BeanstalkError, BeanstalkConnectionError, JobReference
 from .pool import ProductionPool
 
-__version__ = '0.9.1'
+__version__ = '0.10.0'
 
 __author__ = 'EasyPost <oss@easypost.com>'
 
@@ -10,5 +10,6 @@
     'BeanstalkClient',
     'BeanstalkConnectionError',
     'BeanstalkError',
+    'JobReference',
     'ProductionPool',
 ]

pystalk/client.py

@@ -5,7 +5,7 @@
 import socket
 import yaml
 import re
-from typing import Optional, Union
+from typing import Any, Dict, Optional, Union
 
 
 @attr.s(frozen=True)
@@ -24,6 +24,32 @@ class Job(object):
     job_data = attr.ib()
 
 
+@attr.s(frozen=True, slots=True)
+class JobReference(object):
+    """Serializable identity for a job on a specific beanstalkd server."""
+
+    server_id: str = attr.ib(validator=attr.validators.instance_of(str))
+    job_id: int = attr.ib(converter=int)
+
+    def to_dict(self) -> Dict[str, Union[str, int]]:
+        return {
+            'server_id': self.server_id,
+            'job_id': self.job_id,
+        }
+
+    @classmethod
+    def from_dict(cls, value: Dict[str, Any]) -> 'JobReference':
+        if not isinstance(value, dict):
+            raise TypeError('Job reference must be a dictionary')
+        try:
+            return cls(
+                server_id=value['server_id'],
+                job_id=value['job_id'],
+            )
+        except KeyError as exc:
+            raise ValueError('Job reference is missing {0}'.format(exc.args[0])) from exc
+
+
 @attr.s(hash=True, eq=True, init=False)
 class BeanstalkError(Exception):
 
@@ -105,6 +131,13 @@ def put_job(self, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr
         self.beanstalk_client.use(self.tube)
         return self.beanstalk_client.put_job(data=data, pri=pri, delay=delay, ttr=ttr)
 
+    def put_job_with_reference(
+        self, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr: int = 120
+    ) -> JobReference:
+        """Insert a job and return its server-qualified reference."""
+        self.beanstalk_client.use(self.tube)
+        return self.beanstalk_client.put_job_with_reference(data=data, pri=pri, delay=delay, ttr=ttr)
+
 
 class BeanstalkClient(object):
     """
@@ -128,20 +161,21 @@ class BeanstalkClient(object):
          :func:`reserve_job()` will cause errors!
     """
     def __init__(self, host: str, port: int = 11300, socket_timeout: Optional[float] = None,
-                 auto_decode: bool = False):
+                 auto_decode: bool = False, server_id: Optional[str] = None):
         """
         Construct a synchronous Beanstalk Client. Does not connect!
         """
         self.host = host
         self.port = port
+        self.server_id = server_id or '{0}:{1}'.format(host, port)
         self.socket_timeout = socket_timeout
         self._reset_state()
         self.desired_tube = 'default'
         self.desired_watchlist = set(['default'])
         self.auto_decode = auto_decode
 
     @classmethod
-    def from_uri(cls, uri, socket_timeout=None, auto_decode=False):
+    def from_uri(cls, uri, socket_timeout=None, auto_decode=False, server_id=None):
         """
         Construct a synchronous Beanstalk Client from a URI.
 
@@ -163,7 +197,7 @@ def from_uri(cls, uri, socket_timeout=None, auto_decode=False):
             host = parts.netloc
             port = 11300
         port = int(port)
-        return cls(host, port, socket_timeout=socket_timeout, auto_decode=auto_decode)
+        return cls(host, port, socket_timeout=socket_timeout, auto_decode=auto_decode, server_id=server_id)
 
     def _reset_state(self):
         self._watchlist = set(['default'])
@@ -363,6 +397,13 @@ def put_job(self, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr
             self._send_message(message, socket)
             return self._receive_id(socket)
 
+    def put_job_with_reference(
+        self, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr: int = 120
+    ) -> JobReference:
+        """Insert a job and return an identity qualified by this client's server."""
+        _, job_id = self.put_job(data=data, pri=pri, delay=delay, ttr=ttr)
+        return JobReference(server_id=self.server_id, job_id=job_id)
+
     def put_job_into(self, tube_name: str, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr: int = 120):
         """Insert a new job into a specific queue. Wrapper around :func:`put_job`.
 
@@ -389,6 +430,13 @@ def put_job_into(self, tube_name: str, data: Union[str, bytes], pri: int = 65536
         with self.using(tube_name) as inserter:
             return inserter.put_job(data=data, pri=pri, delay=delay, ttr=ttr)
 
+    def put_job_into_with_reference(
+        self, tube_name: str, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr: int = 120
+    ) -> JobReference:
+        """Insert a job into a tube and return its server-qualified reference."""
+        with self.using(tube_name) as inserter:
+            return inserter.put_job_with_reference(data=data, pri=pri, delay=delay, ttr=ttr)
+
     @property
     def watchlist(self):
         return self._watchlist

pystalk/pool.py

@@ -7,7 +7,7 @@
 
 import attr
 
-from .client import BeanstalkClient, BeanstalkConnectionError, BeanstalkError
+from .client import BeanstalkClient, BeanstalkConnectionError, BeanstalkError, JobReference
 
 
 RETRIABLE_ERRORS = ('INTERNAL_ERROR', 'OUT_OF_MEMORY')
@@ -172,6 +172,14 @@ def put_job(self, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr
             lambda client: client.put_job(data=data, pri=pri, delay=delay, ttr=120)
         )
 
+    def put_job_with_reference(
+        self, data: Union[str, bytes], pri: int = 65536, delay: int = 0, ttr: int = 120
+    ) -> JobReference:
+        """Insert a job and return a reference identifying the selected server."""
+        return self._attempt_on_all_clients(
+            lambda client: client.put_job_with_reference(data=data, pri=pri, delay=delay, ttr=ttr)
+        )
+
     def put_job_into(self, tube_name: str, data: Union[str, bytes], pri: int = 65536,
                      delay: int = 0, ttr: int = 120):
         """Insert a new job into a specific queue. Wrapper around :func:`put_job`.
@@ -199,3 +207,14 @@ def put_job_into(self, tube_name: str, data: Union[str, bytes], pri: int = 65536
         return self._attempt_on_all_clients(
             lambda client: client.put_job_into(tube_name=tube_name, data=data, pri=pri, delay=delay, ttr=120)
         )
+
+    def put_job_into_with_reference(
+        self, tube_name: str, data: Union[str, bytes], pri: int = 65536,
+        delay: int = 0, ttr: int = 120
+    ) -> JobReference:
+        """Insert into a tube and return a reference identifying the selected server."""
+        return self._attempt_on_all_clients(
+            lambda client: client.put_job_into_with_reference(
+                tube_name=tube_name, data=data, pri=pri, delay=delay, ttr=ttr
+            )
+        )

tests/unit/test_pool.py

@@ -3,7 +3,7 @@
 
 import pytest
 
-from pystalk import BeanstalkConnectionError
+from pystalk import BeanstalkConnectionError, JobReference
 from pystalk import pool
 from pystalk.pool import ClientRecord, NoMoreClients, ProductionPool
 
@@ -42,6 +42,38 @@ def test_connection_error_fails_over_to_next_client():
     bad_client.close.assert_called_once_with()
 
 
+def test_job_reference_identifies_client_selected_after_failover():
+    calls = []
+    connection_error = BeanstalkConnectionError('bad', 11300, ConnectionRefusedError('refused'))
+    bad_client = make_client('bad', calls)
+    bad_client.put_job_with_reference.side_effect = connection_error
+    good_client = make_client('good', calls)
+    expected_reference = JobReference('good-beanstalk', 42)
+    good_client.put_job_with_reference.return_value = expected_reference
+    production_pool = ProductionPool([bad_client, good_client], initial_shuffle=False, round_robin=False)
+
+    reference = production_pool.put_job_with_reference(b'job', pri=10, delay=20, ttr=30)
+
+    assert reference == expected_reference
+    bad_client.put_job_with_reference.assert_called_once_with(data=b'job', pri=10, delay=20, ttr=30)
+    good_client.put_job_with_reference.assert_called_once_with(data=b'job', pri=10, delay=20, ttr=30)
+
+
+def test_put_job_into_with_reference_delegates_to_selected_client():
+    calls = []
+    client = make_client('client', calls)
+    expected_reference = JobReference('client-beanstalk', 42)
+    client.put_job_into_with_reference.return_value = expected_reference
+    production_pool = ProductionPool([client], initial_shuffle=False)
+
+    reference = production_pool.put_job_into_with_reference('jobs', b'job', pri=10, delay=20, ttr=30)
+
+    assert reference == expected_reference
+    client.put_job_into_with_reference.assert_called_once_with(
+        tube_name='jobs', data=b'job', pri=10, delay=20, ttr=30
+    )
+
+
 def test_all_clients_are_attempted_once_with_zero_backoff():
     calls = []
     first_client = make_client('first', calls, error=socket.error('disconnected'))

tests/unit/test_pystalk.py

@@ -62,6 +62,44 @@ def test_put_job_uses_utf8_byte_length_for_non_ascii(client, server):
     assert server.received == [b'put 65536 0 120 14\r\ncaf\xc3\xa9 (coffee)\r\n']
 
 
+def test_put_job_with_reference_qualifies_job_id_with_server(client, server):
+    server.responses.append(b'INSERTED 42\r\n')
+
+    reference = client.put_job_with_reference('job')
+
+    assert reference == pystalk.JobReference('pystalk.example.com:0', 42)
+    assert server.received == [b'put 65536 0 120 3\r\njob\r\n']
+
+
+def test_put_job_into_with_reference_uses_selected_tube(client, monkeypatch):
+    client.current_tube = 'jobs'
+    put_job = Mock(return_value=(b'INSERTED', 42))
+    monkeypatch.setattr(client, 'put_job', put_job)
+
+    reference = client.put_job_into_with_reference('jobs', 'job', pri=10, delay=20, ttr=30)
+
+    assert reference == pystalk.JobReference('pystalk.example.com:0', 42)
+    put_job.assert_called_once_with(data='job', pri=10, delay=20, ttr=30)
+
+
+def test_job_reference_round_trips_through_dictionary():
+    reference = pystalk.JobReference('primary-beanstalk', 42)
+
+    assert pystalk.JobReference.from_dict(reference.to_dict()) == reference
+
+
+@pytest.mark.parametrize(
+    'value, exception_type, message',
+    [
+        (None, TypeError, 'Job reference must be a dictionary'),
+        ({'server_id': 'primary-beanstalk'}, ValueError, 'Job reference is missing job_id'),
+    ]
+)
+def test_job_reference_rejects_invalid_dictionary(value, exception_type, message):
+    with pytest.raises(exception_type, match=message):
+        pystalk.JobReference.from_dict(value)
+
+
 @pytest.mark.parametrize('uri,expected_host,expected_port', [
     ('beanstalkd://foo', 'foo', 11300),
     ('beanstalk://foo', 'foo', 11300),
@@ -77,6 +115,15 @@ def test_from_uri(uri, expected_host, expected_port):
     assert client.port == expected_port
 
 
+def test_client_accepts_stable_server_id_for_job_references():
+    client = pystalk.BeanstalkClient.from_uri(
+        'beanstalk://alias.example.com:11300',
+        server_id='primary-beanstalk',
+    )
+
+    assert client.server_id == 'primary-beanstalk'
+
+
 @pytest.mark.parametrize('uri', [
     'branstalk://foo:12345',
     'beanstalk://foo:bar',