| """Provides the 1.0, 1.1 and 1.2 protocol classes. |
| """ |
| |
| import uuid |
| |
| from ambari_stomp.backward import encode |
| from ambari_stomp.constants import * |
| from ambari_stomp.exception import ConnectFailedException |
| from ambari_stomp.listener import * |
| import ambari_stomp.utils as utils |
| |
| |
| log = logging.getLogger('stomp.py') |
| |
| |
| class Protocol10(ConnectionListener): |
| """ |
| Represents version 1.0 of the protocol (see https://stomp.github.io/stomp-specification-1.0.html). |
| |
| Most users should not instantiate the protocol directly. See :py:mod:`stomp.connect` for connection classes. |
| |
| :param transport: |
| :param bool auto_content_length: Whether to calculate and send the content-length header automatically if it has not been set |
| """ |
| def __init__(self, transport, auto_content_length=True): |
| self.transport = transport |
| self.auto_content_length = auto_content_length |
| transport.set_listener('protocol-listener', self) |
| self.version = '1.0' |
| |
| def send_frame(self, cmd, headers=None, body=''): |
| """ |
| Encode and send a stomp frame |
| through the underlying transport. |
| |
| :param str cmd: the protocol command |
| :param dict headers: a map of headers to include in the frame |
| :param body: the content of the message |
| """ |
| frame = utils.Frame(cmd, headers, body) |
| self.transport.transmit(frame) |
| |
| def abort(self, transaction, headers=None, **keyword_headers): |
| """ |
| Abort a transaction. |
| |
| :param str transaction: the identifier of the transaction |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert transaction is not None, "'transaction' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_ABORT, headers) |
| |
| def ack(self, id, transaction=None): |
| """ |
| Acknowledge 'consumption' of a message by id. |
| |
| :param str id: identifier of the message |
| :param str transaction: include the acknowledgement in the specified transaction |
| """ |
| assert id is not None, "'id' is required" |
| headers = {HDR_MESSAGE_ID: id} |
| if transaction: |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_ACK, headers) |
| |
| def begin(self, transaction=None, headers=None, **keyword_headers): |
| """ |
| Begin a transaction. |
| |
| :param str transaction: the identifier for the transaction (optional - if not specified |
| a unique transaction id will be generated) |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| |
| :return: the transaction id |
| :rtype: str |
| """ |
| headers = utils.merge_headers([headers, keyword_headers]) |
| if not transaction: |
| transaction = str(uuid.uuid4()) |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_BEGIN, headers) |
| return transaction |
| |
| def commit(self, transaction=None, headers=None, **keyword_headers): |
| """ |
| Commit a transaction. |
| |
| :param str transaction: the identifier for the transaction |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert transaction is not None, "'transaction' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_COMMIT, headers) |
| |
| def connect(self, username=None, passcode=None, wait=False, headers=None, **keyword_headers): |
| """ |
| Start a connection. |
| |
| :param str username: the username to connect with |
| :param str passcode: the password used to authenticate with |
| :param bool wait: if True, wait for the connection to be established/acknowledged |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| cmd = CMD_CONNECT |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_ACCEPT_VERSION] = self.version |
| |
| if username is not None: |
| headers[HDR_LOGIN] = username |
| |
| if passcode is not None: |
| headers[HDR_PASSCODE] = passcode |
| |
| self.send_frame(cmd, headers) |
| |
| if wait: |
| self.transport.wait_for_connection() |
| if self.transport.connection_error: |
| raise ConnectFailedException() |
| |
| def disconnect(self, receipt=None, headers=None, **keyword_headers): |
| """ |
| Disconnect from the server. |
| |
| :param str receipt: the receipt to use (once the server acknowledges that receipt, we're |
| officially disconnected; optional - if not specified a unique receipt id will |
| be generated) |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| if not self.transport.is_connected(): |
| log.debug('Not sending disconnect, already disconnected') |
| return |
| headers = utils.merge_headers([headers, keyword_headers]) |
| rec = receipt or str(uuid.uuid4()) |
| headers[HDR_RECEIPT] = rec |
| self.set_receipt(rec, CMD_DISCONNECT) |
| self.send_frame(CMD_DISCONNECT, headers) |
| |
| def send(self, destination, body, content_type=None, headers=None, **keyword_headers): |
| """ |
| Send a message to a destination. |
| |
| :param str destination: the destination of the message (e.g. queue or topic name) |
| :param body: the content of the message |
| :param str content_type: the content type of the message |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert destination is not None, "'destination' is required" |
| assert body is not None, "'body' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_DESTINATION] = destination |
| if content_type: |
| headers[HDR_CONTENT_TYPE] = content_type |
| body = encode(body) |
| if self.auto_content_length and body and HDR_CONTENT_LENGTH not in headers: |
| headers[HDR_CONTENT_LENGTH] = len(body) |
| self.send_frame(CMD_SEND, headers, body) |
| |
| def subscribe(self, destination, id=None, ack='auto', headers=None, **keyword_headers): |
| """ |
| Subscribe to a destination. |
| |
| :param str destination: the topic or queue to subscribe to |
| :param str id: a unique id to represent the subscription |
| :param str ack: acknowledgement mode, either auto, client, or client-individual |
| (see http://stomp.github.io/stomp-specification-1.2.html#SUBSCRIBE_ack_Header) |
| for more information |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert destination is not None, "'destination' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_DESTINATION] = destination |
| if id: |
| headers[HDR_ID] = id |
| headers[HDR_ACK] = ack |
| self.send_frame(CMD_SUBSCRIBE, headers) |
| |
| def unsubscribe(self, destination=None, id=None, headers=None, **keyword_headers): |
| """ |
| Unsubscribe from a destination by either id or the destination name. |
| |
| :param str destination: the name of the topic or queue to unsubscribe from |
| :param str id: the unique identifier of the topic or queue to unsubscribe from |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert id is not None or destination is not None, "'id' or 'destination' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| if id: |
| headers[HDR_ID] = id |
| if destination: |
| headers[HDR_DESTINATION] = destination |
| self.send_frame(CMD_UNSUBSCRIBE, headers) |
| |
| |
| class Protocol11(HeartbeatListener, ConnectionListener): |
| """ |
| Represents version 1.1 of the protocol (see https://stomp.github.io/stomp-specification-1.1.html). |
| |
| Most users should not instantiate the protocol directly. See :py:mod:`stomp.connect` for connection classes. |
| |
| :param transport: |
| :param (int,int) heartbeats: |
| :param bool auto_content_length: Whether to calculate and send the content-length header automatically if it has not been set |
| """ |
| def __init__(self, transport, heartbeats=(0, 0), auto_content_length=True): |
| HeartbeatListener.__init__(self, heartbeats) |
| self.transport = transport |
| self.auto_content_length = auto_content_length |
| transport.set_listener('protocol-listener', self) |
| self.version = '1.1' |
| |
| def _escape_headers(self, headers): |
| """ |
| :param dict(str,str) headers: |
| """ |
| for key, val in headers.items(): |
| try: |
| val = val.replace('\\', '\\\\').replace('\n', '\\n').replace(':', '\\c') |
| except: |
| pass |
| headers[key] = val |
| |
| def send_frame(self, cmd, headers=None, body=''): |
| """ |
| Encode and send a stomp frame |
| through the underlying transport: |
| |
| :param str cmd: the protocol command |
| :param dict headers: a map of headers to include in the frame |
| :param body: the content of the message |
| """ |
| if cmd != CMD_CONNECT: |
| if headers is None: |
| headers = {} |
| self._escape_headers(headers) |
| frame = utils.Frame(cmd, headers, body) |
| self.transport.transmit(frame) |
| |
| def abort(self, transaction, headers=None, **keyword_headers): |
| """ |
| Abort a transaction. |
| |
| :param str transaction: the identifier of the transaction |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert transaction is not None, "'transaction' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_ABORT, headers) |
| |
| def ack(self, id, subscription, transaction=None): |
| """ |
| Acknowledge 'consumption' of a message by id. |
| |
| :param str id: identifier of the message |
| :param str subscription: the subscription this message is associated with |
| :param str transaction: include the acknowledgement in the specified transaction |
| """ |
| assert id is not None, "'id' is required" |
| assert subscription is not None, "'subscription' is required" |
| headers = {HDR_MESSAGE_ID: id, HDR_SUBSCRIPTION: subscription} |
| if transaction: |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_ACK, headers) |
| |
| def begin(self, transaction=None, headers=None, **keyword_headers): |
| """ |
| Begin a transaction. |
| |
| :param str transaction: the identifier for the transaction (optional - if not specified |
| a unique transaction id will be generated) |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| |
| :return: the transaction id |
| :rtype: str |
| """ |
| headers = utils.merge_headers([headers, keyword_headers]) |
| if not transaction: |
| transaction = str(uuid.uuid4()) |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_BEGIN, headers) |
| return transaction |
| |
| def commit(self, transaction=None, headers=None, **keyword_headers): |
| """ |
| Commit a transaction. |
| |
| :param str transaction: the identifier for the transaction |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert transaction is not None, "'transaction' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_COMMIT, headers) |
| |
| def connect(self, username=None, passcode=None, wait=False, headers=None, **keyword_headers): |
| """ |
| Start a connection. |
| |
| :param str username: the username to connect with |
| :param str passcode: the password used to authenticate with |
| :param bool wait: if True, wait for the connection to be established/acknowledged |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| cmd = CMD_STOMP |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_ACCEPT_VERSION] = self.version |
| |
| if self.transport.vhost: |
| headers[HDR_HOST] = self.transport.vhost |
| |
| if username is not None: |
| headers[HDR_LOGIN] = username |
| |
| if passcode is not None: |
| headers[HDR_PASSCODE] = passcode |
| |
| self.send_frame(cmd, headers) |
| |
| if wait: |
| self.transport.wait_for_connection() |
| if self.transport.connection_error: |
| raise ConnectFailedException() |
| |
| def disconnect(self, receipt=None, headers=None, **keyword_headers): |
| """ |
| Disconnect from the server. |
| |
| :param str receipt: the receipt to use (once the server acknowledges that receipt, we're |
| officially disconnected; optional - if not specified a unique receipt id will |
| be generated) |
| :param dict headers: a map of any additional headers the broker requires |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| if not self.transport.is_connected(): |
| log.debug('Not sending disconnect, already disconnected') |
| return |
| headers = utils.merge_headers([headers, keyword_headers]) |
| rec = receipt or str(uuid.uuid4()) |
| headers[HDR_RECEIPT] = rec |
| self.set_receipt(rec, CMD_DISCONNECT) |
| self.send_frame(CMD_DISCONNECT, headers) |
| |
| def nack(self, id, subscription, transaction=None): |
| """ |
| Let the server know that a message was not consumed. |
| |
| :param str id: the unique id of the message to nack |
| :param str subscription: the subscription this message is associated with |
| :param str transaction: include this nack in a named transaction |
| """ |
| assert id is not None, "'id' is required" |
| assert subscription is not None, "'subscription' is required" |
| headers = {HDR_MESSAGE_ID: id, HDR_SUBSCRIPTION: subscription} |
| if transaction: |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_NACK, headers) |
| |
| def send(self, destination, body, content_type=None, headers=None, **keyword_headers): |
| """ |
| Send a message to a destination in the messaging system (as per https://stomp.github.io/stomp-specification-1.2.html#SEND) |
| |
| :param str destination: the destination (such as a message queue - for example '/queue/test' - or a message topic) |
| :param body: the content of the message |
| :param str content_type: the MIME type of message |
| :param dict headers: additional headers to send in the message frame |
| :param keyword_headers: any additional headers the broker requires |
| """ |
| assert destination is not None, "'destination' is required" |
| assert body is not None, "'body' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_DESTINATION] = destination |
| if content_type: |
| headers[HDR_CONTENT_TYPE] = content_type |
| body = encode(body) |
| if self.auto_content_length and body and HDR_CONTENT_LENGTH not in headers: |
| headers[HDR_CONTENT_LENGTH] = len(body) |
| self.send_frame(CMD_SEND, headers, body) |
| |
| def subscribe(self, destination, id, ack='auto', headers=None, **keyword_headers): |
| """ |
| Subscribe to a destination |
| |
| :param str destination: the topic or queue to subscribe to |
| :param str id: the identifier to uniquely identify the subscription |
| :param str ack: either auto, client or client-individual (see https://stomp.github.io/stomp-specification-1.2.html#SUBSCRIBE for more info) |
| :param dict headers: a map of any additional headers to send with the subscription |
| :param keyword_headers: any additional headers to send with the subscription |
| """ |
| assert destination is not None, "'destination' is required" |
| assert id is not None, "'id' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_DESTINATION] = destination |
| headers[HDR_ID] = id |
| headers[HDR_ACK] = ack |
| self.send_frame(CMD_SUBSCRIBE, headers) |
| |
| def unsubscribe(self, id, headers=None, **keyword_headers): |
| """ |
| Unsubscribe from a destination by its unique identifier |
| |
| :param str id: the unique identifier to unsubscribe from |
| :param dict headers: additional headers to send with the unsubscribe |
| :param keyword_headers: any additional headers to send with the subscription |
| """ |
| assert id is not None, "'id' is required" |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_ID] = id |
| self.send_frame(CMD_UNSUBSCRIBE, headers) |
| |
| |
| class Protocol12(Protocol11): |
| """ |
| Represents version 1.2 of the protocol (see https://stomp.github.io/stomp-specification-1.2.html). |
| |
| Most users should not instantiate the protocol directly. See :py:mod:`stomp.connect` for connection classes. |
| |
| :param transport: |
| :param (int,int) heartbeats: |
| :param bool auto_content_length: Whether to calculate and send the content-length header automatically if it has not been set |
| """ |
| def __init__(self, transport, heartbeats=(0, 0), auto_content_length=True): |
| Protocol11.__init__(self, transport, heartbeats, auto_content_length) |
| self.version = '1.2' |
| |
| def _escape_headers(self, headers): |
| """ |
| :param dict(str,str) headers: |
| """ |
| for key, val in headers.items(): |
| try: |
| val = val.replace('\\', '\\\\').replace('\n', '\\n').replace(':', '\\c').replace('\r', '\\r') |
| except: |
| pass |
| headers[key] = val |
| |
| def ack(self, id, transaction=None): |
| """ |
| Acknowledge 'consumption' of a message by id. |
| |
| :param str id: identifier of the message |
| :param str transaction: include the acknowledgement in the specified transaction |
| """ |
| assert id is not None, "'id' is required" |
| headers = {HDR_ID: id} |
| if transaction: |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_ACK, headers) |
| |
| def nack(self, id, transaction=None): |
| """ |
| Let the server know that a message was not consumed. |
| |
| :param str id: the unique id of the message to nack |
| :param str transaction: include this nack in a named transaction |
| """ |
| assert id is not None, "'id' is required" |
| headers = {HDR_ID: id} |
| if transaction: |
| headers[HDR_TRANSACTION] = transaction |
| self.send_frame(CMD_NACK, headers) |
| |
| def connect(self, username=None, passcode=None, wait=False, headers=None, **keyword_headers): |
| """ |
| Send a STOMP CONNECT frame. Differs from 1.0 and 1.1 versions in that the HOST header is enforced. |
| |
| :param str username: optionally specify the login user |
| :param str passcode: optionally specify the user password |
| :param bool wait: wait for the connection to complete before returning |
| :param dict headers: a map of any additional headers to send with the subscription |
| :param keyword_headers: any additional headers to send with the subscription |
| """ |
| cmd = CMD_STOMP |
| headers = utils.merge_headers([headers, keyword_headers]) |
| headers[HDR_ACCEPT_VERSION] = self.version |
| headers[HDR_HOST] = self.transport.current_host_and_port[0] |
| |
| if self.transport.vhost: |
| headers[HDR_HOST] = self.transport.vhost |
| |
| if username is not None: |
| headers[HDR_LOGIN] = username |
| |
| if passcode is not None: |
| headers[HDR_PASSCODE] = passcode |
| |
| self.send_frame(cmd, headers) |
| |
| if wait: |
| self.transport.wait_for_connection() |
| if self.transport.connection_error: |
| raise ConnectFailedException() |
