Handshake protocols and pooled connections

# -*- coding: utf-8 -*-

# stdlib
import socket

# Zato
from zato.common.sdk import ConnectionLost, Field, PooledConnector

class MainframeConnection:
    """ One connection to a mainframe gateway. The gateway requires a logon handshake first and
    then serves one call at a time per session, which is why the platform pools these connections.
    """
    def __init__(self, host:'str', port:'int', logon_token:'str') -> 'None':

        # Connect and keep the socket open for as long as the connection exists ..
        self.socket = socket.create_connection((host, port))
        self.reader = self.socket.makefile('r', encoding='utf8')

        # .. log on first - the gateway answers with the session it assigned ..
        self._send_line(f'logon {logon_token}')
        response = self.reader.readline().strip()

        # .. anything other than an ok reply means the logon was rejected.
        if not response.startswith('ok '):
            raise Exception(f'Logon failed -> {response}')

        self.session_id = response[len('ok '):]

    def _send_line(self, data:'str') -> 'None':
        self.socket.sendall(f'{data}\n'.encode('utf8'))

    def send(self, data:'str') -> 'str':
        self._send_line(data)
        return self.reader.readline().strip()

    def close(self) -> 'None':
        self.reader.close()
        self.socket.close()

class MainframeConnector(PooledConnector):
    """ Wraps the mainframe gateway as a connection type that services access through self.out.mainframe.
    """
    type = 'mainframe'

    # Configuration schema
    host = Field.Text()
    port = Field.Int(default=9960)
    logon_token = Field.Secret()

    def create_client(self) -> 'MainframeConnection':
        conn = MainframeConnection(self.config.host, self.config.port, self.config.logon_token)
        self.logger.info('Mainframe session `%s` logged on for `%s`', conn.session_id, self.name)
        return conn

    def ping(self, conn:'MainframeConnection') -> 'None':

        # A closed socket produces an empty line - raising ConnectionLost makes the platform
        # discard the connection.
        response = conn.send('ping')
        if not response.endswith('ping'):
            raise ConnectionLost(f'The gateway did not answer a ping -> {response!r}')

    def on_stop(self, conn:'MainframeConnection') -> 'None':
        conn.close()
        self.logger.info('Mainframe session `%s` closed for `%s`', conn.session_id, self.name)

    def send_command(self, command:'str') -> 'str':

        # Each call borrows one connection for its whole duration - the gateway serves
        # one call at a time per session.
        with self.get_connection() as conn:
            return conn.send(command)

    def on_get_from_pool(self, conn:'MainframeConnection') -> 'None':
        # Runs each time a call borrows this connection - reset conversational state here
        ...

    def on_return_to_pool(self, conn:'MainframeConnection') -> 'None':
        # Runs each time a with block ends - clean up before the next caller borrows the connection
        ...

# -*- coding: utf-8 -*-

# Zato
from zato.server.service import Service

class SendMainframeCommand(Service):
    """ Sends one command to the mainframe gateway over a pooled connection.
    """
    name = 'demo.mainframe.send-command'

    input = 'command'

    def handle(self) -> 'None':
        conn = self.out.mainframe['My Mainframe']
        response = conn.send_command(self.request.input.command)
        self.response.payload = response