- enterprise it!

Author: zzzeek

sqlalchemy_collectd/client/__init__.py

@@ -1,10 +1,9 @@
-import os
 import socket
 import sys
 
 from sqlalchemy.engine import CreateEnginePlugin
 
-from . import collectd
+from .. import protocol
 from . import sender
 from . import worker
 from . import collector
@@ -40,7 +39,7 @@ def engine_created(self, engine):
             progname)
         collector.EngineCollector(collection_target, engine)
 
-        connection = collectd.ClientConnection.for_host_port(
+        connection = protocol.ClientConnection.for_host_port(
             collectd_hostname, collectd_port)
 
         worker.add_target(

sqlalchemy_collectd/collector.py renamed to sqlalchemy_collectd/client/collector.py

@@ -1,13 +1,13 @@
-from . import collectd
-from . import types
+from .. import protocol
+from .. import types
 
 
 senders = []
 
 
-def sends(collectd_type):
+def sends(protocol_type):
     def decorate(fn):
-        senders.append((collectd_type, fn))
+        senders.append((protocol_type, fn))
         return fn
 
     return decorate
@@ -19,9 +19,9 @@ def __init__(self, hostname, stats_name):
         self.stats_name = stats_name
 
     def send(self, connection, collection_target, timestamp, interval, pid):
-        for collectd_type, sender in senders:
-            message_sender = collectd.MessageSender(
-                collectd_type, self.hostname, "sqlalchemy",
+        for protocol_type, sender in senders:
+            message_sender = protocol.MessageSender(
+                protocol_type, self.hostname, "sqlalchemy",
                 plugin_instance=self.stats_name, type_instance=str(pid),
                 interval=interval
             )

sqlalchemy_collectd/plugin.py renamed to sqlalchemy_collectd/client/plugin.py

@@ -1,3 +1,7 @@
+"""Using the collectd protocol because we originally were going to
+connect straight to the network plugin.
+
+"""
 import struct
 import socket
 import logging
@@ -24,18 +28,16 @@
 TYPE_VALUES = 0x0006
 TYPE_INTERVAL = 0x0007
 
-header = struct.Struct("!2H")
-number = struct.Struct("!Q")
+header = struct.Struct("!HH")
 short = struct.Struct("!H")
-double = struct.Struct("<d")
 char = struct.Struct("B")
 long_ = struct.Struct("!q")
 
 _value_formats = {
-    VALUE_COUNTER: number,
-    VALUE_GAUGE: double,
+    VALUE_COUNTER: struct.Struct("!Q"),
+    VALUE_GAUGE: struct.Struct("<d"),
     VALUE_DERIVE: long_,
-    VALUE_ABSOLUTE: number
+    VALUE_ABSOLUTE: struct.Struct("!Q")
 }
 
 
@@ -56,7 +58,10 @@ class Type(object):
 
     """
 
-    __slots__ = 'name', '_value_types', '_value_formats', '_message_template'
+    __slots__ = (
+        'name', '_value_types', '_value_formats',
+        '_message_template', '_field_names'
+    )
 
     def __init__(self, name, *db_template):
         """Contruct a new Type.
@@ -77,6 +82,7 @@ def __init__(self, name, *db_template):
 
         """
         self.name = name
+        self._field_names = [dsname for dsname, value_type in db_template]
         self._value_types = [value_type for dsname, value_type in db_template]
         self._value_formats = [
             _value_formats[value_type] for value_type in self._value_types]
@@ -87,7 +93,7 @@ def __init__(self, name, *db_template):
         for value_type in self._value_types:
             self._message_template += char.pack(value_type)
 
-    def encode_values(self, *values):
+    def _encode_values(self, *values):
         """Encode a series of values according to the type template."""
 
         msg = self._message_template
@@ -137,11 +143,80 @@ def send(self, connection, timestamp, *values):
             long_.pack(int(timestamp)) + \
             self._remainder_message_parts
 
-        payload = self.type.encode_values(*values)
+        payload = self.type._encode_values(*values)
 
         connection.send(header_ + payload)
 
 
+class MessageReceiver(object):
+    def __init__(self, *types):
+        self._receivers = {
+            TYPE_HOST: self._unpack_string,
+            TYPE_TIME: self._unpack_long,
+            TYPE_PLUGIN: self._unpack_string,
+            TYPE_PLUGIN_INSTANCE: self._unpack_string,
+            TYPE_TYPE: self._unpack_string,
+            TYPE_TYPE_INSTANCE: self._unpack_string,
+            TYPE_VALUES: self._unpack_values,
+            TYPE_INTERVAL: self._unpack_long,
+        }
+        self._types = {
+            type_.name: type_ for type_ in types
+        }
+
+    def receive(self, buf):
+        result = self._unpack_packet(buf)
+        type_name = result[TYPE_TYPE]
+        try:
+            type_ = self._types[type_name]
+        except KeyError:
+            log.warn("Type %s not known, skipping", type_name)
+        else:
+            result["type"] = type_
+            result["values"] = {
+                name: value
+                for name, value in zip(type_._field_names, result[TYPE_VALUES])
+            }
+            return result
+
+    def _unpack_packet(self, buf):
+        pos = 0
+        length = len(buf)
+        result = {}
+        while pos < length:
+            type_, len_ = header.unpack_from(buf, pos)
+
+            try:
+                fn = self._receivers[type_]
+            except KeyError:
+                log.warn("Message %s not known, skipping", type_)
+            else:
+                value = fn(type_, len_, buf[pos:])
+
+                result[type_] = value
+
+            pos += len_
+        return result
+
+    def _unpack_long(self, type_, length, buf):
+        return long_.unpack_from(buf, header.size)[0]
+
+    def _unpack_string(self, type_, length, buf):
+        return buf[header.size:length - 1]
+
+    def _unpack_values(self, type_, length, buf):
+        num = short.unpack_from(buf, header.size)[0]
+        types_start = header.size + short.size
+        values_pos = types_start + num * char.size
+        result = []
+        for pos in range(0, num * char.size, char.size):
+            value_type = char.unpack_from(buf, types_start + pos)[0]
+            struct_ = _value_formats[value_type]
+            result.append(struct_.unpack_from(buf, values_pos)[0])
+            values_pos += struct_.size
+        return result
+
+
 class ClientConnection(object):
     connections = {}
     create_mutex = threading.Lock()
@@ -183,3 +258,7 @@ def send(self, message):
             log.error("Error in socket.sendto", exc_info=True)
         finally:
             self._mutex.release()
+
+
+class ServerConnection(object):
+    pass

sqlalchemy_collectd/sender.py renamed to sqlalchemy_collectd/client/sender.py

@@ -0,0 +1,86 @@
+import mock
+import unittest
+
+from sqlalchemy_collectd import protocol
+
+
+class CollectDProtocolTest(unittest.TestCase):
+
+    def test_encode_type_values(self):
+        type_ = protocol.Type(
+            "my_type",
+            ("some_val", protocol.VALUE_GAUGE),
+            ("some_other_val", protocol.VALUE_DERIVE)
+        )
+
+        self.assertEqual(
+            self.value_block,
+            type_._encode_values(25.809, 450)
+        )
+
+    value_block = (
+        b'\x00\x06'  # TYPE_VALUES
+        b'\x00\x18'  # part length
+        b'\x00\x02'  # number of values
+        b'\x01\x02'  # dstype codes GAUGE, DERIVE
+        b'\xc9v\xbe\x9f\x1a\xcf9@'  # 8 bytes for 25.809
+        b'\x00\x00\x00\x00\x00\x00\x01\xc2'  # 8 bytes for 450
+    )
+
+    message = (
+        b'\x00\x00\x00\rsomehost\x00'  # TYPE_HOST
+        b'\x00\x01\x00\x0c\x00\x00\x00\x00Zt\xd8\x82'  # TYPE_TIME
+        b'\x00\x02\x00\x0fsomeplugin\x00'  # TYPE_PLUGIN
+        # TYPE_PLUGIN_INSTANCE
+        b'\x00\x03\x00\x17someplugininstance\x00'
+        b'\x00\x04\x00\x0cmy_type\x00'  # TYPE_TYPE
+        # TYPE_TIMESTAMP
+        b'\x00\x07\x00\x0c\x00\x00\x00\x00\x00\x00\x00\n'
+        b'\x00\x05\x00\x15sometypeinstance\x00'  # TYPE_TYPE_INSTANCE
+    ) + value_block
+
+    def test_message_construct(self):
+        type_ = protocol.Type(
+            "my_type",
+            ("some_val", protocol.VALUE_GAUGE),
+            ("some_other_val", protocol.VALUE_DERIVE)
+        )
+
+        sender = protocol.MessageSender(
+            type_, "somehost", "someplugin", "someplugininstance",
+            "sometypeinstance"
+        )
+
+        connection = mock.Mock()
+
+        sender.send(connection, 1517607042.95968, 25.809, 450)
+
+        self.assertEqual(
+            [mock.call(self.message)],
+            connection.send.mock_calls
+        )
+
+    def test_decode_packet(self):
+        type_ = protocol.Type(
+            "my_type",
+            ("some_val", protocol.VALUE_GAUGE),
+            ("some_other_val", protocol.VALUE_DERIVE)
+        )
+
+        message_receiver = protocol.MessageReceiver(type_)
+        result = message_receiver.receive(self.message)
+        self.assertEqual(
+            {
+                protocol.TYPE_HOST: 'somehost',
+                protocol.TYPE_TIME: 1517607042,
+                protocol.TYPE_PLUGIN: 'someplugin',
+                protocol.TYPE_PLUGIN_INSTANCE: 'someplugininstance',
+                protocol.TYPE_TYPE: 'my_type',
+                protocol.TYPE_TYPE_INSTANCE: 'sometypeinstance',
+                protocol.TYPE_VALUES: [25.809, 450],
+                protocol.TYPE_INTERVAL: 10,
+                "type": type_,
+                "values": {'some_other_val': 450, 'some_val': 25.809}
+            },
+            result
+        )

sqlalchemy_collectd/worker.py renamed to sqlalchemy_collectd/client/worker.py

@@ -1,27 +1,27 @@
-from . import collectd
+from . import protocol
 
-pool = collectd.Type(
+pool = protocol.Type(
 	"sqlalchemy_pool",
-	("numpools", collectd.VALUE_GAUGE),
-	("checkedout", collectd.VALUE_GAUGE),
-	("checkedin", collectd.VALUE_GAUGE),
-	("detached", collectd.VALUE_GAUGE),
-	("invalidated", collectd.VALUE_GAUGE),
-	("total", collectd.VALUE_GAUGE),
+	("numpools", protocol.VALUE_GAUGE),
+	("checkedout", protocol.VALUE_GAUGE),
+	("checkedin", protocol.VALUE_GAUGE),
+	("detached", protocol.VALUE_GAUGE),
+	("invalidated", protocol.VALUE_GAUGE),
+	("total", protocol.VALUE_GAUGE),
 )
 
-checkouts = collectd.Type(
-	"sqlalchemy_checkouts", ("count", collectd.VALUE_DERIVE))
+checkouts = protocol.Type(
+	"sqlalchemy_checkouts", ("count", protocol.VALUE_DERIVE))
 
-commits = collectd.Type(
-	"sqlalchemy_commits", ("count", collectd.VALUE_DERIVE))
+commits = protocol.Type(
+	"sqlalchemy_commits", ("count", protocol.VALUE_DERIVE))
 
-rollbacks = collectd.Type(
-	"sqlalchemy_rollbacks", ("count", collectd.VALUE_DERIVE))
+rollbacks = protocol.Type(
+	"sqlalchemy_rollbacks", ("count", protocol.VALUE_DERIVE))
 
-invalidated = collectd.Type(
-	"sqlalchemy_invalidated", ("count", collectd.VALUE_DERIVE))
+invalidated = protocol.Type(
+	"sqlalchemy_invalidated", ("count", protocol.VALUE_DERIVE))
 
-transactions = collectd.Type(
-	"sqlalchemy_transactions", ("count", collectd.VALUE_DERIVE))
+transactions = protocol.Type(
+	"sqlalchemy_transactions", ("count", protocol.VALUE_DERIVE))