Rewrite protocol implementation

Author: 1st1

asyncpg/__init__.py

@@ -1,7 +1,6 @@
 import asyncio
 import getpass
 import os
-import socket
 import urllib.parse
 
 from .exceptions import *  # NOQA
@@ -20,18 +19,14 @@ async def connect(dsn=None, *,
                   loop=None,
                   timeout=60,
                   statement_cache_size=100,
-                  **kwargs):
+                  **opts):
 
     if loop is None:
         loop = asyncio.get_event_loop()
 
-    host, port, user, password, database, kwargs = _parse_connect_params(
+    host, port, opts = _parse_connect_params(
         dsn=dsn, host=host, port=port, user=user, password=password,
-        database=database, kwargs=kwargs)
-
-    if kwargs:
-        raise RuntimeError(
-            'arbitrary connection arguments are not yet supported')
+        database=database, opts=opts)
 
     last_ex = None
     for h in host:
@@ -42,25 +37,18 @@ async def connect(dsn=None, *,
             # UNIX socket name
             sname = os.path.join(h, '.s.PGSQL.{}'.format(port))
             conn = loop.create_unix_connection(
-                lambda: protocol.Protocol(sname, connected, user,
-                                          password, database, loop),
+                lambda: protocol.Protocol(sname, connected, opts, loop),
                 sname)
         else:
             conn = loop.create_connection(
-                lambda: protocol.Protocol((h, port), connected, user,
-                                          password, database, loop),
+                lambda: protocol.Protocol((h, port), connected, opts, loop),
                 h, port)
 
         try:
             tr, pr = await asyncio.wait_for(conn, timeout=timeout, loop=loop)
         except (OSError, asyncio.TimeoutError) as ex:
             last_ex = ex
         else:
-            if not unix:
-                sock = tr.get_extra_info('socket')
-                if hasattr(socket, 'TCP_NODELAY'):
-                    sock.setsockopt(socket.IPPROTO_TCP,
-                                    socket.TCP_NODELAY, 1)
             break
     else:
         raise last_ex
@@ -76,7 +64,7 @@ async def connect(dsn=None, *,
 
 
 def _parse_connect_params(*, dsn, host, port, user,
-                          password, database, kwargs):
+                          password, database, opts):
 
     if dsn:
         parsed = urllib.parse.urlparse(dsn)
@@ -140,7 +128,7 @@ def _parse_connect_params(*, dsn, host, port, user,
                     password = val
 
             if query:
-                kwargs = {**query, **kwargs}
+                opts = {**query, **opts}
 
     # On env-var -> connection parameter conversion read here:
     # https://www.postgresql.org/docs/current/static/libpq-envars.html
@@ -176,7 +164,24 @@ def _parse_connect_params(*, dsn, host, port, user,
     if database is None:
         database = os.getenv('PGDATABASE')
 
-    return host, port, user, password, database, kwargs
+    if user is not None:
+        opts['user'] = user
+    if password is not None:
+        opts['password'] = password
+    if database is not None:
+        opts['database'] = database
+
+    for param in opts:
+        if not isinstance(param, str):
+            raise ValueError(
+                'invalid connection parameter {!r} (str expected)'
+                .format(param))
+        if not isinstance(opts[param], str):
+            raise ValueError(
+                'invalid connection parameter {!r}: {!r} (str expected)'
+                .format(param, opts[param]))
+
+    return host, port, opts
 
 
 def _create_future(loop):

asyncpg/connection.py

@@ -169,13 +169,12 @@ async def close(self):
             return
         self._close_stmts()
         self._aborted = True
-        self._transport.abort()
         await self._protocol.close()
 
     def terminate(self):
         self._close_stmts()
         self._aborted = True
-        self._transport.abort()
+        self._protocol.abort()
 
     def _get_unique_id(self):
         self._uid += 1

asyncpg/protocol/buffer.pxd

@@ -96,7 +96,7 @@ cdef class ReadBuffer:
     cdef inline read_cstr(self)
     cdef int32_t has_message(self) except -1
     cdef inline char* try_consume_message(self, int32_t* len)
-    cdef consume_message(self)
+    cdef Memory consume_message(self)
     cdef discard_message(self)
     cdef inline _discard_message(self)
     cdef inline char get_message_type(self)

asyncpg/protocol/buffer.pyx

@@ -484,7 +484,7 @@ cdef class ReadBuffer:
             self._discard_message()
         return buf
 
-    cdef consume_message(self):
+    cdef Memory consume_message(self):
         if not self._current_message_ready:
             raise BufferError('no message to consume')
         mem = self.read(self._current_message_len_unread)
@@ -500,10 +500,14 @@ cdef class ReadBuffer:
             raise BufferError('no message to discard')
 
         if self._current_message_len_unread:
+            IF DEBUG:
+                mtype = chr(self._current_message_type)
+
             discarded = self.consume_message()
+
             IF DEBUG:
                 print('!!! discarding message {!r} unread data: {!r}'.format(
-                    chr(self._current_message_type),
+                    mtype,
                     (<Memory>discarded).as_bytes()))
 
         self._discard_message()

asyncpg/protocol/coreproto.pxd

@@ -1,139 +1,99 @@
 cdef enum ConnectionStatus:
-    CONNECTION_OK = 0
-    CONNECTION_BAD = 1
-    CONNECTION_STARTED = 2           # Waiting for connection to be made.
-    CONNECTION_MADE = 3              # Connection OK; waiting to send.
-    CONNECTION_AWAITING_RESPONSE = 4 # Waiting for a response from the
-                                     # postmaster.
-    CONNECTION_AUTH_OK = 5           # Received authentication; waiting for
-                                     # backend startup.
-    CONNECTION_SETENV = 6            # Negotiating environment.
-    CONNECTION_SSL_STARTUP = 7       # Negotiating SSL.
-    CONNECTION_NEEDED = 8            # Internal state: connect() needed
-
-cdef enum AsyncStatus:
-    # defines the state of the query-execution state machine
-    PGASYNC_IDLE = 0                 # nothing's happening, dude
-    PGASYNC_BUSY = 1                 # query in progress
-    PGASYNC_READY = 2                # result ready for PQgetResult
-    PGASYNC_COPY_IN = 3              # Copy In data transfer in progress
-    PGASYNC_COPY_OUT = 4             # Copy Out data transfer in progress
-    PGASYNC_COPY_BOTH = 5            # Copy In/Out data transfer in progress
-
-cdef enum QueryClass:
-    # tracks which query protocol we are now executing
-    PGQUERY_SIMPLE = 0               # simple Query protocol (PQexec)
-    PGQUERY_EXTENDED = 1             # full Extended protocol (PQexecParams)
-    PGQUERY_PREPARE = 2              # Parse only (PQprepare)
-    PGQUERY_DESCRIBE = 3             # Describe Statement or Portal
-    PGQUERY_CLOSE = 4
-
-cdef enum ExecStatusType:
-    PGRES_EMPTY_QUERY = 0            # empty query string was executed
-    PGRES_COMMAND_OK = 1             # a query command that doesn't return
-                                     # anything was executed properly by the
-                                     # backend
-    PGRES_TUPLES_OK = 2              # a query command that returns tuples was
-                                     # executed properly by the backend,
-                                     # PGresult contains the result tuples
-    PGRES_COPY_OUT = 3               # Copy Out data transfer in progress
-    PGRES_COPY_IN = 4                # Copy In data transfer in progress
-    PGRES_BAD_RESPONSE = 5           # an unexpected response was recv'd from
-                                     # the backend
-    PGRES_NONFATAL_ERROR = 6         # notice or warning message
-    PGRES_FATAL_ERROR = 7            # query failed
-    PGRES_COPY_BOTH = 8              # Copy In/Out data transfer in progress
-    PGRES_SINGLE_TUPLE = 9           # single tuple from larger resultset
+    CONNECTION_OK = 1
+    CONNECTION_BAD = 2
+    CONNECTION_STARTED = 3           # Waiting for connection to be made.
 
-cdef enum TransactionStatus:
-    PQTRANS_IDLE = 0                 # connection idle
-    PQTRANS_ACTIVE = 1               # command in progress
-    PQTRANS_INTRANS = 2              # idle, within transaction block
-    PQTRANS_INERROR = 3              # idle, within failed transaction
-    PQTRANS_UNKNOWN = 4              # cannot determine status
 
-cdef enum MessageDispatchLoop:
-    DISPATCH_CONTINUE = 0
-    DISPATCH_STOP = 1
-    DISPATCH_CONTINUE_NO_DISCARD = 2
+cdef enum ProtocolState:
+    PROTOCOL_IDLE = 0
 
+    PROTOCOL_FAILED = 1
 
-cdef class Result:
-    cdef:
-        ExecStatusType status
+    PROTOCOL_ERROR_CONSUME = 2
+
+    PROTOCOL_AUTH = 10
+    PROTOCOL_PREPARE = 11
+    PROTOCOL_BIND_EXECUTE = 12
+    PROTOCOL_CLOSE_STMT_PORTAL = 13
+    PROTOCOL_SIMPLE_QUERY = 14
 
-        # message broken into fields
-        dict err_fields
-        # text of triggering query, if available
-        str  err_query
 
-        # cmd status from the query
-        bytes cmd_status
+cdef enum ResultType:
+    RESULT_OK = 1
+    RESULT_FAILED = 2
 
-        object parameters_desc
-        object row_desc
-        list rows
 
-    @staticmethod
-    cdef Result new(ExecStatusType status)
+cdef enum TransactionStatus:
+    PQTRANS_IDLE = 0                 # connection idle
+    PQTRANS_ACTIVE = 1               # command in progress
+    PQTRANS_INTRANS = 2              # idle, within transaction block
+    PQTRANS_INERROR = 3              # idle, within failed transaction
+    PQTRANS_UNKNOWN = 4              # cannot determine status
 
 
 ctypedef object (*decode_row_method)(object, const char*, int32_t)
 
 
 cdef class CoreProtocol:
     cdef:
-        object transport
         ReadBuffer buffer
+        bint _skip_discard
 
-        ####### Options:
+        ConnectionStatus con_status
+        ProtocolState state
+        TransactionStatus xact_status
 
-        str _user
-        str _password
-        str _dbname
-        str _encoding
+        str encoding
 
-        ####### Connection State:
-        int _backend_pid
-        int _backend_secret
-
-        # result being constructed
-        Result            _result
+        object transport
 
-        ConnectionStatus  _status
-        AsyncStatus       _async_status
-        QueryClass        _query_class
-        TransactionStatus _xact_status
+        # Dict with all connection arguments
+        dict con_args
+
+        int32_t backend_pid
+        int32_t backend_secret
+
+        ## Result
+        ResultType result_type
+        object result
+        bytes result_param_desc
+        bytes result_row_desc
+        bytes result_status_msg
+
+    cdef _process__auth(self, char mtype)
+    cdef _process__prepare(self, char mtype)
+    cdef _process__bind_exec(self, char mtype)
+    cdef _process__close_stmt_portal(self, char mtype)
+    cdef _process__simple_query(self, char mtype)
+
+    cdef _parse_msg_authentication(self)
+    cdef _parse_msg_parameter_status(self)
+    cdef _parse_msg_backend_key_data(self)
+    cdef _parse_msg_ready_for_query(self)
+    cdef _parse_data_msgs(self)
+    cdef _parse_msg_error_response(self, is_error)
+    cdef _parse_msg_command_complete(self)
+
+    cdef _write(self, buf)
+    cdef inline _write_sync_message(self)
 
-        WriteBuffer _after_sync
+    cdef _read_server_messages(self)
 
-    cdef inline _write(self, WriteBuffer buf)
-    cdef inline _write_sync_message(self)
-    cdef inline _read_server_messages(self)
-    cdef inline MessageDispatchLoop _dispatch_server_message(self, char mtype)
-    cdef _parse_server_authentication(self)
-    cdef _parse_server_parameter_status(self)
-    cdef _parse_server_backend_key_data(self)
-    cdef _parse_server_parameter_description(self)
-    cdef _parse_server_ready_for_query(self)
-    cdef _parse_server_row_description(self)
-    cdef _parse_server_data_rows(self)
-    cdef _parse_server_error_response(self, is_error)
-    cdef _fatal_error(self, exc)
     cdef _push_result(self)
-    cdef _sync(self)
-    cdef _ensure_ready_state(self)
+    cdef _reset_result(self)
+    cdef _set_state(self, ProtocolState new_state)
+
+    cdef _ensure_connected(self)
 
     cdef _connect(self)
-    cdef _query(self, str query)
     cdef _prepare(self, str stmt_name, str query)
-    cdef _bind(self, str portal_name, str stmt_name,
-               WriteBuffer bind_data, int32_t limit)
+    cdef _bind_and_execute(self, str portal_name, str stmt_name,
+                           WriteBuffer bind_data, int32_t limit)
     cdef _close(self, str name, bint is_portal)
+    cdef _simple_query(self, str query)
 
     cdef _decode_row(self, const char* buf, int32_t buf_len)
 
-    cdef _on_result(self, Result result)
-    cdef _on_fatal_error(self, exc)
+    cdef _on_result(self)
+    cdef _set_server_parameter(self, name, val)
     cdef _on_connection_lost(self, exc)
-    cdef _set_server_parameter(self, key, val)