Merge pull request #2172 from creideiki/sql-updates

SQL: add MSSQL support and allow storing only a subset of event fields

Author: sebix

CHANGELOG.md

@@ -20,6 +20,8 @@ CHANGELOG
 - `intelmq.lib.message`:
   - Fix and pre-compile the regular expression for harmonization key names and also check keys in the `extra.` namespace (PR#2059 by Sebastian Wagner, fixes #1807).
 - `intelmq.lib.bot.SQLBot` was replaced by an SQLMixin in `intelmq.lib.mixins.SQLMixin`. The Generic DB Lookup Expert bot and the SQLOutput bot were updated accordingly.
+  - Added support for MSSQL (PR#2171 by Karl-Johan Karlsson).
+  - Added optional reconnect delay parameter (PR#2171 by Karl-Johan Karlsson).
 - Added an ExpertBot class - it should be used by all expert bots as a parent class
 - Introduced a module for IntelMQ related datatypes `intelmq.lib.datatypes` which for now only contains an Enum listing the four bot types
 - Added a `bottype` attribute to CollectorBot, ParserBot, ExpertBot, OutputBot

docs/user/bots.rst

@@ -4006,7 +4006,7 @@ SQL
 * `lookup:` no
 * `public:` yes
 * `cache (redis db):` none
-* `description:` SQL is the bot responsible to send events to a PostgreSQL or SQLite Database, e.g. the IntelMQ :doc:`eventdb`
+* `description:` SQL is the bot responsible to send events to a PostgreSQL, SQLite, or MSSQL Database, e.g. the IntelMQ :doc:`eventdb`
 * `notes`: When activating autocommit, transactions are not used: http://initd.org/psycopg/docs/connection.html#connection.autocommit
 
 **Configuration Parameters**
@@ -4015,15 +4015,17 @@ The parameters marked with 'PostgreSQL' will be sent to libpq via psycopg2. Chec
 
 * `autocommit`: `psycopg's autocommit mode <http://initd.org/psycopg/docs/connection.html?#connection.autocommit>`_, optional, default True
 * `connect_timeout`: Database connect_timeout, optional, default 5 seconds
-* `engine`: 'postgresql' or 'sqlite'
-* `database`: PostgreSQL database or SQLite file
-* `host`: PostgreSQL host
+* `engine`: 'postgresql', 'sqlite', or 'mssql'
+* `database`: Database or SQLite file
+* `host`: Database host
 * `jsondict_as_string`: save JSONDict fields as JSON string, boolean. Default: true (like in versions before 1.1)
-* `port`: PostgreSQL port
-* `user`: PostgreSQL user
-* `password`: PostgreSQL password
-* `sslmode`: PostgreSQL sslmode, can be `'disable'`, `'allow'`, `'prefer'` (default), `'require'`, `'verify-ca'` or `'verify-full'`. See postgresql docs: https://www.postgresql.org/docs/current/static/libpq-connect.html#libpq-connect-sslmode
+* `port`: Database port
+* `user`: Database user
+* `password`: Database password
+* `sslmode`: Database sslmode, can be `'disable'`, `'allow'`, `'prefer'` (default), `'require'`, `'verify-ca'` or `'verify-full'`. See postgresql docs: https://www.postgresql.org/docs/current/static/libpq-connect.html#libpq-connect-sslmode
 * `table`: name of the database table into which events are to be inserted
+* `fields`: list of fields to read from the event. If None, read all fields
+* `reconnect_delay`: number of seconds to wait before reconnecting in case of an error
 
 **PostgreSQL**
 
@@ -4099,6 +4101,10 @@ Then, set the `database` parameter to the `your-db.db` file path.
 
 .. _intelmq.bots.outputs.stomp.output:
 
+**MSSQL**
+
+For MSSQL support, the library `pymssql>=2.2` is required.
+
 STOMP
 ^^^^^
 

intelmq/bots/outputs/sql/REQUIREMENTS.txt

@@ -2,3 +2,4 @@
 # SPDX-License-Identifier: AGPL-3.0-or-later
 
 psycopg2-binary>=2.5.5
+pymssql>=2.2

intelmq/bots/outputs/sql/output.py

@@ -17,6 +17,12 @@
 from intelmq.lib.mixins import SQLMixin
 
 
+def itemgetter_tuple(*items):
+    def g(obj):
+        return tuple(obj[item] for item in items)
+    return g
+
+
 class SQLOutputBot(OutputBot, SQLMixin):
     """Send events to a PostgreSQL or SQLite database"""
     autocommit = True
@@ -29,16 +35,21 @@ class SQLOutputBot(OutputBot, SQLMixin):
     sslmode = "require"
     table = 'events'
     user = "intelmq"
+    fields = None
 
     def init(self):
         super().init()
 
     def process(self):
         event = self.receive_message().to_dict(jsondict_as_string=self.jsondict_as_string)
 
-        keys = '", "'.join(event.keys())
-        values = list(event.values())
-        fvalues = len(values) * f'{self.format_char}, '
+        key_names = self.fields
+        if key_names is None:
+            key_names = event.keys()
+        valid_keys = [key for key in key_names if key in event]
+        keys = '", "'.join(valid_keys)
+        values = itemgetter_tuple(*valid_keys)(event)
+        fvalues = len(values) * '{0}, '.format(self.format_char)
         query = ('INSERT INTO {table} ("{keys}") VALUES ({values})'
                  ''.format(table=self.table, keys=keys, values=fvalues[:-2]))
 

intelmq/lib/mixins/sql.py

@@ -7,6 +7,8 @@
 """
 from intelmq.lib import exceptions
 
+from time import sleep
+
 
 class SQLMixin:
     """
@@ -19,10 +21,12 @@ class SQLMixin:
 
     POSTGRESQL = "postgresql"
     SQLITE = "sqlite"
+    MSSQL = "mssql"
     _default_engine = "postgresql"
     engine = None
     # overwrite the default value from the OutputBot
     message_jsondict_as_string = True
+    reconnect_delay = 0
 
     def __init__(self, *args, **kwargs):
         self._init_sql()
@@ -33,7 +37,8 @@ def _init_sql(self):
         self.logger.debug("Running SQL Mixin initialization.")
         self._engine_name = getattr(self, 'engine', self._default_engine).lower()
         engines = {SQLMixin.POSTGRESQL: (self._init_postgresql, "%s"),
-                   SQLMixin.SQLITE: (self._init_sqlite, "?")}
+                   SQLMixin.SQLITE: (self._init_sqlite, "?"),
+                   SQLMixin.MSSQL: (self._init_mssql, "%s")}
         for key, val in engines.items():
             if self._engine_name == key:
                 val[0]()
@@ -48,7 +53,7 @@ def _connect(self, engine, connect_args: dict, autocommitable: bool = False):
 
         try:
             self.con = self._engine.connect(**connect_args)
-            if autocommitable:  # psycopg2 has it, sqlite3 has not
+            if autocommitable:  # psycopg2 and mssql has it, sqlite3 has not
                 self.con.autocommit = getattr(self, 'autocommit', True)  # True prevents deadlocks
             self.cur = self.con.cursor()
         except (self._engine.Error, Exception):
@@ -86,6 +91,23 @@ def _init_sqlite(self):
                        }
                       )
 
+    def _init_mssql(self):
+        try:
+            import pymssql
+        except ImportError:
+            raise exceptions.MissingDependencyError("pymssql")
+
+        self._connect(pymssql,
+                      {"server": self.host,
+                       "user": self.user,
+                       "password": self.password,
+                       "database": self.database,
+                       "login_timeout": getattr(self, 'connect_timeout', 5),
+                       "port": self.port,
+                       "as_dict": True
+                       },
+                      autocommitable=True)
+
     def execute(self, query: str, values: tuple, rollback=False):
         try:
             self.logger.debug('Executing %r.', (query, values))
@@ -102,13 +124,19 @@ def execute(self, query: str, values: tuple, rollback=False):
                 except self._engine.OperationalError:
                     self.logger.exception('Executed rollback command '
                                           'after failed query execution.')
+                    if self.reconnect_delay > 0:
+                        sleep(self.reconnect_delay)
                     self._init_sql()
                 except Exception:
                     self.logger.exception('Cursor has been closed, connecting '
                                           'again.')
+                    if self.reconnect_delay > 0:
+                        sleep(self.reconnect_delay)
                     self._init_sql()
             else:
                 self.logger.exception('Database connection problem, connecting again.')
+                if self.reconnect_delay > 0:
+                    sleep(self.reconnect_delay)
                 self._init_sql()
         else:
             return True