Refactor connection infrastructure with async io

Use selectors to manage sockets in parallel, and enable multiple connections
Simplifies code by a lot

Author: Alexandru Pisarenco

client.py

@@ -1,65 +1,62 @@
 import logging
 
-TYPE_SERVER=0
-TYPE_CLIENT=1
-
 class Connection:
-    def __init__(self, sock, target, name):
+    def __init__(self, sock, address, name, events, context):
         self.sock = sock
-        self.target = target
+        self.address = address
         self.name = name
-
-    def send(self, message):
-        logging.debug("sending message to {}:\n{}".format(self.name, message))
-        total = len(message)
-        total_sent = 0
-        remaining = message
-        while total_sent < total:
-            sent = self.sock.send(remaining)
-            total_sent += sent
-            remaining = remaining[sent:]
-
-    def __receive_raw(self, length):
-        total_received = 0
-        chunks = []
-        while total_received < length:
-            chunk = self.sock.recv(min([length - total_received]), 4096)
-            if chunk==b'':
-                raise RuntimeError("socket connection broken")
-            chunks.append(chunk)
-            total_received += len(chunk)
-        return b''.join(chunks)
-
-
-    def receive_packet(self):
-        pack_type = self.__receive_raw(1)
-        if pack_type == b'N':
-            # Null message? This message has no length. Just a single byte. Weird.
-            return pack_type, pack_type
-        if pack_type == b'\x00':
-            # Initialization packet. No type. This, and the next 3 bytes are the length
-            pack_length = self.__receive_raw(3)
-            pack_length = b''.join([pack_type, pack_length])
-            pack_header = pack_length
-        else:
-            pack_length = self.__receive_raw(4)
-            pack_header = b''.join([pack_type, pack_length])
-        pack_length = int.from_bytes(pack_length, 'big')
-        pack_body = self.__receive_raw(pack_length - 4)
-        pack = b''.join([pack_header, pack_body])
-        return pack, pack_type
-
-
-    def receive(self):
-        packets = []
-        logging.debug("receive message from {}:".format(self.name))
+        self.events = events
+        self.context = context
+        self.is_reading = False
+        self.is_writing = False
+        self.interceptor = None
+        self.redirect_conn = None
+        self.out_bytes = b''
+        self.in_bytes = b''
+
+    def parse_length(self, length_bytes):
+        return int.from_bytes(length_bytes, 'big')
+
+    def encode_length(self, length):
+        return length.to_bytes(4, byteorder='big')
+
+    def received(self, in_bytes):
+        self.in_bytes += in_bytes
+        # Read packet from byte array while there are enough bytes to make up a packet.
+        # Otherwise wait for more bytes to be received (break and exit)
         while True:
-            logging.debug("receive packet from {}:".format(self.name))
-            packet, pack_type = self.receive_packet()
-            packets.append(packet)
-            logging.debug("received packet from {}:\n{}".format(self.name, packet))
-            if not pack_type in (b'B', b'D', b'P', b'E'):
+            ptype = self.in_bytes[0:1]
+            if ptype == b'\x00':
+                if len(self.in_bytes) < 4:
+                    break
+                header_length = 4
+                body_length = self.parse_length(self.in_bytes[0:4]) - 4
+            elif ptype == b'N':
+                header_length = 1
+                body_length = 0
+            else:
+                if len(self.in_bytes) < 5:
+                    break
+                header_length = 5
+                body_length = self.parse_length(self.in_bytes[1:5]) - 4
+
+            length = header_length + body_length
+            if len(self.in_bytes) < length:
                 break
-        message = b''.join(packets)
-        logging.debug("received message from {}:\n{}".format(self.name, message))
-        return message
+            header = self.in_bytes[0:header_length]
+            body = self.in_bytes[header_length:length]
+            self.process_inbound_packet(header, body)
+            self.in_bytes = self.in_bytes[length:]
+
+    def process_inbound_packet(self, header, body):
+        if header != b'N':
+            packet_type = header[0:-4]
+            logging.info("intercepting packet of type '%s' from %s", packet_type, self.name)
+            body = self.interceptor.intercept(packet_type, body)
+            header = packet_type + self.encode_length(len(body) + 4)
+        message = header + body
+        logging.debug("Received message. Relaying. Speaker: %s, message:\n%s", self.name, message)
+        self.redirect_conn.out_bytes += message
+
+    def sent(self, num_bytes):
+        self.out_bytes = self.out_bytes[num_bytes:]

connection.py

@@ -0,0 +1,82 @@
+import logging
+
+class Interceptor:
+    def __init__(self, interceptor_config, plugins, context):
+        self.interceptor_config = interceptor_config
+        self.plugins = plugins
+        self.context = context
+
+    def intercept(self, packet_type, data):
+        return data
+
+    def get_codec(self):
+        if self.context is not None and 'connect_params' in self.context:
+            if self.context['connect_params'] is not None and 'client_encoding' in self.context['connect_params']:
+                return self.context['connect_params']['client_encoding']
+        return 'utf-8'
+
+
+class CommandInterceptor(Interceptor):
+    def intercept(self, packet_type, data):
+        if self.interceptor_config.queries is not None:
+            ic_queries = self.interceptor_config.queries
+            if packet_type == b'Q':
+                # Query, ends with b'\x00'
+                data = self.__intercept_query(data, ic_queries)
+            elif packet_type == b'P':
+                # Statement that needs parsing.
+                # First byte of the body is some Statement flag. Ignore, don't lose
+                # Next is the query, same as above, ends with an b'\x00'
+                # Last 2 bytes are the number of parameters. Ignore, don't lose
+                statement = data[0:1]
+                query = self.__intercept_query(data[1:-2], ic_queries)
+                params = data[-2:]
+                data = statement + query + params
+            elif packet_type == b'':
+                self.__intercept_context_data(data)
+        return data
+
+
+    def __intercept_context_data(self, data):
+        # first 4 bytes and last zero byte are not interesting
+        relevant_data = data[4:-1]
+        # Each entry is terminated by b'\x00'
+        entries = relevant_data.split(b'\x00')[:-1]
+        entries = dict(zip(entries[0::2], entries[1::2]))
+        self.context['connect_params'] = {}
+        # Try to set codec, then transcode the dict
+        if b'client_encoding' in entries:
+            self.context['connect_params']['client_encoding'] = entries[b'client_encoding'].decode('ascii')
+        codec = self.get_codec()
+        for k, v in entries.items():
+            self.context['connect_params'][k.decode(codec)] = v.decode(codec)
+
+
+    def __intercept_query(self, query, interceptors):
+        logging.getLogger('intercept').debug("intercepting query\n%s", query)
+        # Remove zero byte at the end
+        query = query[:-1].decode('utf-8')
+        for interceptor in interceptors:
+            if interceptor.plugin in self.plugins:
+                plugin = self.plugins[interceptor.plugin]
+                if hasattr(plugin, interceptor.function):
+                    func = getattr(plugin, interceptor.function)
+                    query = func(query, self.context)
+                    logging.getLogger('intercept').debug(
+                        "modifying query using interceptor %s.%s\n%s",
+                        interceptor.plugin,
+                        interceptor.function,
+                        query)
+                else:
+                    raise Exception("Can't find function {} in plugin {}".format(
+                        interceptor.function,
+                        interceptor.plugin
+                    ))
+            else:
+                raise Exception("Plugin {} not loaded".format(interceptor.plugin))
+        # Append the zero byte at the end
+        return query.encode('utf-8') + b'\x00'
+
+
+class ResponseInterceptor(Interceptor):
+    pass

interceptors.py

@@ -27,7 +27,7 @@ def replace(match):
 
             # need to know which columns are hll
             if not hll_table.lower() in column_cache:
-                db_conn_info = context.instance_config.redirect
+                db_conn_info = context['instance_config'].redirect
                 conn = None
                 try:
                     conn = psycopg2.connect(
@@ -36,8 +36,8 @@ def replace(match):
                             db_conn_info.host,
                             db_conn_info.port,
                             # Get auth information from the proxied request
-                            context.connect_params['database'],
-                            context.connect_params['user']
+                            context['connect_params']['database'],
+                            context['connect_params']['user']
                         )
                     )
 
@@ -83,7 +83,6 @@ def replace(match):
         return match.group(0)
 
 
-    query = query.decode('utf-8')
     # Matches this string. The 2 groups are `schema.table` and `"alias"`
     # FROM schema.table) "alias"
     table_result = table_pattern.search(query)
@@ -94,4 +93,4 @@ def replace(match):
     # Replaces count(distinct ...) with hll_cardinality(hll_union_agg(...)) :: BIGINT
     # where and how it is appropriate
     # the inner function `replace` uses the variables `original_table` and `table_alias` from this scope (smelly code)
-    return field_pattern.sub(replace, query).encode('utf-8')
+    return field_pattern.sub(replace, query)

pg_connection.py

@@ -3,18 +3,19 @@
 import plugins.tableau_hll as hll
 import yaml
 
-class TestContext():
-    def __init__(self):
-        with open(os.path.dirname(__file__) + '/config.yml', 'r') as fp:
-            self.config = yaml.load(fp)
-        InstanceConfig = collections.namedtuple('InstanceConfig', 'redirect')
-        Redirect = collections.namedtuple('Redirect', 'name host port')
-        self.instance_config = InstanceConfig(redirect=Redirect(**self.config['redirect']))
-        self.connect_params = self.config['connect_params']
+def test_context():
+    with open(os.path.dirname(__file__) + '/config.yml', 'r') as fp:
+        config = yaml.load(fp)
+    InstanceConfig = collections.namedtuple('InstanceConfig', 'redirect')
+    Redirect = collections.namedtuple('Redirect', 'name host port')
+    return {
+        'instance_config': InstanceConfig(redirect=Redirect(**config['redirect']))
+        'connect_params': config['connect_params']
+    }
 
 
 def run():
-    query = b'SELECT COUNT(DISTINCT "crm_data_source"."Set of Customers") AS "ctd:Set of Customers:ok"\nFROM "crm_dim"."crm_data_source" "crm_data_source"\nHAVING (COUNT(1) > 0);'
-    out_query = b'SELECT  hll_cardinality(hll_union_agg("crm_data_source"."Set of Customers")) :: BIGINT  AS "ctd:Set of Customers:ok"\nFROM "crm_dim"."crm_data_source" "crm_data_source"\nHAVING (COUNT(1) > 0);'
-    context = TestContext()
+    query = 'SELECT COUNT(DISTINCT "crm_data_source"."Set of Customers") AS "ctd:Set of Customers:ok"\nFROM "crm_dim"."crm_data_source" "crm_data_source"\nHAVING (COUNT(1) > 0);'
+    out_query = 'SELECT  hll_cardinality(hll_union_agg("crm_data_source"."Set of Customers")) :: BIGINT  AS "ctd:Set of Customers:ok"\nFROM "crm_dim"."crm_data_source" "crm_data_source"\nHAVING (COUNT(1) > 0);'
+    context = test_context()
     assert hll.rewrite_query(query, context) == out_query