protocol.py

import struct
from gearman.constants import PRIORITY_NONE, PRIORITY_LOW, PRIORITY_HIGH
from gearman.errors import ProtocolError
from gearman import compat
# Protocol specific constants
NULL_CHAR = '\x00'
MAGIC_RES_STRING = '%sRES' % NULL_CHAR
MAGIC_REQ_STRING = '%sREQ' % NULL_CHAR

COMMAND_HEADER_SIZE = 12

# Gearman commands 1-9
GEARMAN_COMMAND_CAN_DO = 1
GEARMAN_COMMAND_CANT_DO = 2
GEARMAN_COMMAND_RESET_ABILITIES = 3
GEARMAN_COMMAND_PRE_SLEEP = 4
GEARMAN_COMMAND_NOOP = 6
GEARMAN_COMMAND_SUBMIT_JOB = 7
GEARMAN_COMMAND_JOB_CREATED = 8
GEARMAN_COMMAND_GRAB_JOB = 9

# Gearman commands 10-19
GEARMAN_COMMAND_NO_JOB = 10
GEARMAN_COMMAND_JOB_ASSIGN = 11
GEARMAN_COMMAND_WORK_STATUS = 12
GEARMAN_COMMAND_WORK_COMPLETE = 13
GEARMAN_COMMAND_WORK_FAIL = 14
GEARMAN_COMMAND_GET_STATUS = 15
GEARMAN_COMMAND_ECHO_REQ = 16
GEARMAN_COMMAND_ECHO_RES = 17
GEARMAN_COMMAND_SUBMIT_JOB_BG = 18
GEARMAN_COMMAND_ERROR = 19

# Gearman commands 20-29
GEARMAN_COMMAND_STATUS_RES = 20
GEARMAN_COMMAND_SUBMIT_JOB_HIGH = 21
GEARMAN_COMMAND_SET_CLIENT_ID = 22
GEARMAN_COMMAND_CAN_DO_TIMEOUT = 23
GEARMAN_COMMAND_ALL_YOURS = 24
GEARMAN_COMMAND_WORK_EXCEPTION = 25
GEARMAN_COMMAND_OPTION_REQ = 26
GEARMAN_COMMAND_OPTION_RES = 27
GEARMAN_COMMAND_WORK_DATA = 28
GEARMAN_COMMAND_WORK_WARNING = 29

# Gearman commands 30-39
GEARMAN_COMMAND_GRAB_JOB_UNIQ = 30
GEARMAN_COMMAND_JOB_ASSIGN_UNIQ = 31
GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG = 32
GEARMAN_COMMAND_SUBMIT_JOB_LOW = 33
GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG = 34

# Fake command code 
GEARMAN_COMMAND_TEXT_COMMAND = 9999

GEARMAN_PARAMS_FOR_COMMAND = {
    # Gearman commands 1-9
    GEARMAN_COMMAND_CAN_DO: ['task'],
    GEARMAN_COMMAND_CANT_DO: ['task'],
    GEARMAN_COMMAND_RESET_ABILITIES: [],
    GEARMAN_COMMAND_PRE_SLEEP: [],
    GEARMAN_COMMAND_NOOP: [],
    GEARMAN_COMMAND_SUBMIT_JOB: ['task', 'unique', 'data'],
    GEARMAN_COMMAND_JOB_CREATED: ['job_handle'],
    GEARMAN_COMMAND_GRAB_JOB: [],

    # Gearman commands 10-19
    GEARMAN_COMMAND_NO_JOB: [],
    GEARMAN_COMMAND_JOB_ASSIGN: ['job_handle', 'task', 'data'],
    GEARMAN_COMMAND_WORK_STATUS: ['job_handle', 'numerator', 'denominator'],
    GEARMAN_COMMAND_WORK_COMPLETE: ['job_handle', 'data'],
    GEARMAN_COMMAND_WORK_FAIL: ['job_handle'],
    GEARMAN_COMMAND_GET_STATUS: ['job_handle'],
    GEARMAN_COMMAND_ECHO_REQ: ['data'],
    GEARMAN_COMMAND_ECHO_RES: ['data'],
    GEARMAN_COMMAND_SUBMIT_JOB_BG: ['task', 'unique', 'data'],
    GEARMAN_COMMAND_ERROR: ['error_code', 'error_text'],

    # Gearman commands 20-29
    GEARMAN_COMMAND_STATUS_RES: ['job_handle', 'known', 'running', 'numerator', 'denominator'],
    GEARMAN_COMMAND_SUBMIT_JOB_HIGH: ['task', 'unique', 'data'],
    GEARMAN_COMMAND_SET_CLIENT_ID: ['client_id'],
    GEARMAN_COMMAND_CAN_DO_TIMEOUT: ['task', 'timeout'],
    GEARMAN_COMMAND_ALL_YOURS: [],
    GEARMAN_COMMAND_WORK_EXCEPTION: ['job_handle', 'data'],
    GEARMAN_COMMAND_OPTION_REQ: ['option_name'],
    GEARMAN_COMMAND_OPTION_RES: ['option_name'],
    GEARMAN_COMMAND_WORK_DATA: ['job_handle', 'data'],
    GEARMAN_COMMAND_WORK_WARNING: ['job_handle', 'data'],

    # Gearman commands 30-39
    GEARMAN_COMMAND_GRAB_JOB_UNIQ: [],
    GEARMAN_COMMAND_JOB_ASSIGN_UNIQ: ['job_handle', 'task', 'unique', 'data'],
    GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG: ['task', 'unique', 'data'],
    GEARMAN_COMMAND_SUBMIT_JOB_LOW: ['task', 'unique', 'data'],
    GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG: ['task', 'unique', 'data'],

    # Fake gearman command
    GEARMAN_COMMAND_TEXT_COMMAND: ['raw_text']
}

GEARMAN_COMMAND_TO_NAME = {
    GEARMAN_COMMAND_CAN_DO: 'GEARMAN_COMMAND_CAN_DO',
    GEARMAN_COMMAND_CANT_DO: 'GEARMAN_COMMAND_CANT_DO',
    GEARMAN_COMMAND_RESET_ABILITIES: 'GEARMAN_COMMAND_RESET_ABILITIES',
    GEARMAN_COMMAND_PRE_SLEEP: 'GEARMAN_COMMAND_PRE_SLEEP',
    GEARMAN_COMMAND_NOOP: 'GEARMAN_COMMAND_NOOP',
    GEARMAN_COMMAND_SUBMIT_JOB: 'GEARMAN_COMMAND_SUBMIT_JOB',
    GEARMAN_COMMAND_JOB_CREATED: 'GEARMAN_COMMAND_JOB_CREATED',
    GEARMAN_COMMAND_GRAB_JOB: 'GEARMAN_COMMAND_GRAB_JOB',

    # Gearman commands 10-19
    GEARMAN_COMMAND_NO_JOB: 'GEARMAN_COMMAND_NO_JOB',
    GEARMAN_COMMAND_JOB_ASSIGN: 'GEARMAN_COMMAND_JOB_ASSIGN',
    GEARMAN_COMMAND_WORK_STATUS: 'GEARMAN_COMMAND_WORK_STATUS',
    GEARMAN_COMMAND_WORK_COMPLETE: 'GEARMAN_COMMAND_WORK_COMPLETE',
    GEARMAN_COMMAND_WORK_FAIL: 'GEARMAN_COMMAND_WORK_FAIL',
    GEARMAN_COMMAND_GET_STATUS: 'GEARMAN_COMMAND_GET_STATUS',
    GEARMAN_COMMAND_ECHO_REQ: 'GEARMAN_COMMAND_ECHO_REQ',
    GEARMAN_COMMAND_ECHO_RES: 'GEARMAN_COMMAND_ECHO_RES',
    GEARMAN_COMMAND_SUBMIT_JOB_BG: 'GEARMAN_COMMAND_SUBMIT_JOB_BG',
    GEARMAN_COMMAND_ERROR: 'GEARMAN_COMMAND_ERROR',

    # Gearman commands 20-29
    GEARMAN_COMMAND_STATUS_RES: 'GEARMAN_COMMAND_STATUS_RES',
    GEARMAN_COMMAND_SUBMIT_JOB_HIGH: 'GEARMAN_COMMAND_SUBMIT_JOB_HIGH',
    GEARMAN_COMMAND_SET_CLIENT_ID: 'GEARMAN_COMMAND_SET_CLIENT_ID',
    GEARMAN_COMMAND_CAN_DO_TIMEOUT: 'GEARMAN_COMMAND_CAN_DO_TIMEOUT',
    GEARMAN_COMMAND_ALL_YOURS: 'GEARMAN_COMMAND_ALL_YOURS',
    GEARMAN_COMMAND_WORK_EXCEPTION: 'GEARMAN_COMMAND_WORK_EXCEPTION',
    GEARMAN_COMMAND_OPTION_REQ: 'GEARMAN_COMMAND_OPTION_REQ',
    GEARMAN_COMMAND_OPTION_RES: 'GEARMAN_COMMAND_OPTION_RES',
    GEARMAN_COMMAND_WORK_DATA: 'GEARMAN_COMMAND_WORK_DATA',
    GEARMAN_COMMAND_WORK_WARNING: 'GEARMAN_COMMAND_WORK_WARNING',

    # Gearman commands 30-39
    GEARMAN_COMMAND_GRAB_JOB_UNIQ: 'GEARMAN_COMMAND_GRAB_JOB_UNIQ',
    GEARMAN_COMMAND_JOB_ASSIGN_UNIQ: 'GEARMAN_COMMAND_JOB_ASSIGN_UNIQ',
    GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG: 'GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG',
    GEARMAN_COMMAND_SUBMIT_JOB_LOW: 'GEARMAN_COMMAND_SUBMIT_JOB_LOW',
    GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG: 'GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG',

    GEARMAN_COMMAND_TEXT_COMMAND: 'GEARMAN_COMMAND_TEXT_COMMAND'
}

GEARMAN_SERVER_COMMAND_STATUS = 'status'
GEARMAN_SERVER_COMMAND_VERSION = 'version'
GEARMAN_SERVER_COMMAND_WORKERS = 'workers'
GEARMAN_SERVER_COMMAND_MAXQUEUE = 'maxqueue'
GEARMAN_SERVER_COMMAND_SHUTDOWN = 'shutdown'

def get_command_name(cmd_type):
    return GEARMAN_COMMAND_TO_NAME.get(cmd_type, cmd_type)

def submit_cmd_for_background_priority(background, priority):
    cmd_type_lookup = {
        (True, PRIORITY_NONE): GEARMAN_COMMAND_SUBMIT_JOB_BG,
        (True, PRIORITY_LOW): GEARMAN_COMMAND_SUBMIT_JOB_LOW_BG,
        (True, PRIORITY_HIGH): GEARMAN_COMMAND_SUBMIT_JOB_HIGH_BG,
        (False, PRIORITY_NONE): GEARMAN_COMMAND_SUBMIT_JOB,
        (False, PRIORITY_LOW): GEARMAN_COMMAND_SUBMIT_JOB_LOW,
        (False, PRIORITY_HIGH): GEARMAN_COMMAND_SUBMIT_JOB_HIGH
    }
    lookup_tuple = (background, priority)
    cmd_type = cmd_type_lookup[lookup_tuple]
    return cmd_type

def parse_binary_command(in_buffer, is_response=True):
    """Parse data and return (command type, command arguments dict, command size)
    or (None, None, data) if there's not enough data for a complete command.
    """
    in_buffer_size = len(in_buffer)
    magic = None
    cmd_type = None
    cmd_args = None
    cmd_len = 0
    expected_packet_size = None

    # If we don't have enough data to parse, error early
    if in_buffer_size < COMMAND_HEADER_SIZE:
        return cmd_type, cmd_args, cmd_len

    # By default, we'll assume we're dealing with a gearman command
    magic, cmd_type, cmd_len = struct.unpack('!4sII', in_buffer[:COMMAND_HEADER_SIZE])
    magic = magic.decode('utf-8')

    received_bad_response = is_response and bool(magic != MAGIC_RES_STRING)
    received_bad_request = not is_response and bool(magic != MAGIC_REQ_STRING)
    if received_bad_response or received_bad_request:
        raise ProtocolError('Malformed Magic')

    expected_cmd_params = GEARMAN_PARAMS_FOR_COMMAND.get(cmd_type, None)

    # GEARMAN_COMMAND_TEXT_COMMAND is a faked command that we use to support server text-based commands
    if expected_cmd_params is None or cmd_type == GEARMAN_COMMAND_TEXT_COMMAND:
        raise ProtocolError('Received unknown binary command: %s' % cmd_type)

    # If everything indicates this is a valid command, we should check to see if we have enough stuff to read in our buffer
    expected_packet_size = COMMAND_HEADER_SIZE + cmd_len
    if in_buffer_size < expected_packet_size:
        return None, None, 0

    binary_payload = in_buffer[COMMAND_HEADER_SIZE:expected_packet_size]
    binary_payload = binary_payload.decode('utf-8')
    split_arguments = []

    if len(expected_cmd_params) > 0:
        split_arguments = binary_payload.split(NULL_CHAR, len(expected_cmd_params) - 1)
    elif binary_payload:
        raise ProtocolError('Expected no binary payload: %s' % get_command_name(cmd_type))

    # This is a sanity check on the binary_payload.split() phase
    # We should never be able to get here with any VALID gearman data
    if len(split_arguments) != len(expected_cmd_params):
        raise ProtocolError('Received %d argument(s), expecting %d argument(s): %s' % (len(split_arguments), len(expected_cmd_params), get_command_name(cmd_type)))

    # Iterate through the split arguments and assign them labels based on their order
    cmd_args = dict((param_label, param_value) for param_label, param_value in zip(expected_cmd_params, split_arguments))
    return cmd_type, cmd_args, expected_packet_size


def pack_binary_command(cmd_type, cmd_args, is_response=False):
    """Packs the given command using the parameter ordering specified in GEARMAN_PARAMS_FOR_COMMAND.
    *NOTE* Expects that all arguments in cmd_args are already str's.
    """
    expected_cmd_params = GEARMAN_PARAMS_FOR_COMMAND.get(cmd_type, None)
    if expected_cmd_params is None or cmd_type == GEARMAN_COMMAND_TEXT_COMMAND:
        raise ProtocolError('Received unknown binary command: %s' % get_command_name(cmd_type))

    expected_parameter_set = set(expected_cmd_params)
    received_parameter_set = set(cmd_args.keys())
    if expected_parameter_set != received_parameter_set:
        raise ProtocolError('Received arguments did not match expected arguments: %r != %r' % (expected_parameter_set, received_parameter_set))

    # Select the right expected magic
    if is_response:
        magic = MAGIC_RES_STRING
    else:
        magic = MAGIC_REQ_STRING

    # !NOTE! str should be replaced with bytes in Python 3.x
    # We will iterate in ORDER and str all our command arguments
    if compat.any(type(param_value) != str for param_value in cmd_args.values()):
        raise ProtocolError('Received non-binary arguments: %r' % cmd_args)

    data_items = [cmd_args[param] for param in expected_cmd_params]
    binary_payload = NULL_CHAR.join(data_items)

    # Pack the header in the !4sII format then append the binary payload
    payload_size = len(binary_payload)
    packing_format = '!4sII%ds' % payload_size
    return struct.pack(packing_format, magic.encode(), cmd_type, payload_size, binary_payload.encode())

def parse_text_command(in_buffer):
    """Parse a text command and return a single line at a time"""
    cmd_type = None
    cmd_args = None
    cmd_len = 0
    in_buffer = in_buffer.decode()
    if '\n' not in in_buffer:
        return cmd_type, cmd_args, cmd_len

    text_command, in_buffer = in_buffer.split('\n', 1)
    if NULL_CHAR in text_command:
        raise ProtocolError('Received unexpected character: %s' % text_command)

    # Fake gearman command "TEXT_COMMAND" used to process server admin client responses
    cmd_type = GEARMAN_COMMAND_TEXT_COMMAND
    cmd_args = dict(raw_text=text_command)
    cmd_len = len(text_command) + 1

    return cmd_type, cmd_args, cmd_len

def pack_text_command(cmd_type, cmd_args):
    """Parse a text command and return a single line at a time"""
    if cmd_type != GEARMAN_COMMAND_TEXT_COMMAND:
        raise ProtocolError('Unknown cmd_type: Received %s, expecting %s' % (get_command_name(cmd_type), get_command_name(GEARMAN_COMMAND_TEXT_COMMAND)))

    cmd_line = cmd_args.get('raw_text')
    if cmd_line is None:
        raise ProtocolError('Did not receive arguments any valid arguments: %s' % cmd_args)

    return str(cmd_line)