Add sqlalchemy collector

Beside the already available collectors, add a sqlalchemy based
collector. This is useful if you don't want to maintain another DB
solution and just use the (usually) already available database.
The driver currently implements the notify() and get_report() methods
so it is possible to store trace points and to get a single trace.

Change-Id: If91b35d4b97862c0ecf6677f4c6b95a09d411195

4 files changed, 151 insertions, 0 deletions

diff --git a/doc/source/user/collectors.rst b/doc/source/user/collectors.rst
index e163d57..5d48caa 100644
--- a/doc/source/user/collectors.rst
+++ b/doc/source/user/collectors.rst
@@ -39,3 +39,28 @@ Redis
       value. Defaults to: 0.1 seconds
     * sentinel_service_name: The name of the Sentinel service to use.
       Defaults to: "mymaster"
+
+SQLAlchemy
+----------
+
+The SQLAlchemy collector allows you to store profiling data into a database
+supported by SQLAlchemy.
+
+Usage
+=====
+To use the driver, the `connection_string` in the `[osprofiler]` config section
+needs to be set to a connection string that `SQLAlchemy understands`_
+For example::
+
+  [osprofiler]
+  connection_string = mysql+pymysql://username:password@192.168.192.81/profiler?charset=utf8
+
+where `username` is the database username, `password` is the database password,
+`192.168.192.81` is the database IP address and `profiler` is the database name.
+
+The database (in this example called `profiler`) needs to be created manually and
+the database user (in this example called `username`) needs to have priviliges
+to create tables and select and insert rows.
+
+
+.. _SQLAlchemy understands: https://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls
diff --git a/osprofiler/drivers/__init__.py b/osprofiler/drivers/__init__.py
index 37fdb69..022b094 100644
--- a/osprofiler/drivers/__init__.py
+++ b/osprofiler/drivers/__init__.py
@@ -5,3 +5,4 @@ from osprofiler.drivers import loginsight  # noqa
 from osprofiler.drivers import messaging  # noqa
 from osprofiler.drivers import mongodb  # noqa
 from osprofiler.drivers import redis_driver  # noqa
+from osprofiler.drivers import sqlalchemy_driver  # noqa
diff --git a/osprofiler/drivers/base.py b/osprofiler/drivers/base.py
index 6583a88..b85ffda 100644
--- a/osprofiler/drivers/base.py
+++ b/osprofiler/drivers/base.py
@@ -36,6 +36,12 @@ def get_driver(connection_string, *args, **kwargs):
               connection_string)
 
     backend = parsed_connection.scheme
+    # NOTE(toabctl): To be able to use the connection_string for as sqlalchemy
+    # connection string, transform the backend to the correct driver
+    # See https://docs.sqlalchemy.org/en/latest/core/engines.html#database-urls
+    if backend in ["mysql", "mysql+pymysql", "mysql+mysqldb",
+                   "postgresql", "postgresql+psycopg2"]:
+        backend = "sqlalchemy"
     for driver in _utils.itersubclasses(Driver):
         if backend == driver.get_name():
             return driver(connection_string, *args, **kwargs)
diff --git a/osprofiler/drivers/sqlalchemy_driver.py b/osprofiler/drivers/sqlalchemy_driver.py
new file mode 100644
index 0000000..c16a3ac
--- /dev/null
+++ b/osprofiler/drivers/sqlalchemy_driver.py
@@ -0,0 +1,119 @@
+# Copyright 2019 SUSE Linux GmbH
+# All Rights Reserved.
+#
+#    Licensed under the Apache License, Version 2.0 (the "License"); you may
+#    not use this file except in compliance with the License. You may obtain
+#    a copy of the License at
+#
+#         http://www.apache.org/licenses/LICENSE-2.0
+#
+#    Unless required by applicable law or agreed to in writing, software
+#    distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
+#    WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
+#    License for the specific language governing permissions and limitations
+#    under the License.
+
+import logging
+
+from oslo_serialization import jsonutils
+
+from osprofiler.drivers import base
+from osprofiler import exc
+
+LOG = logging.getLogger(__name__)
+
+
+class SQLAlchemyDriver(base.Driver):
+    def __init__(self, connection_str, project=None, service=None, host=None,
+                 **kwargs):
+        super(SQLAlchemyDriver, self).__init__(connection_str, project=project,
+                                               service=service, host=host)
+
+        try:
+            from sqlalchemy import create_engine
+            from sqlalchemy import Table, MetaData, Column
+            from sqlalchemy import String, JSON, Integer
+        except ImportError:
+            raise exc.CommandError(
+                "To use this command, you should install 'SQLAlchemy'")
+
+        self._engine = create_engine(connection_str)
+        self._conn = self._engine.connect()
+        self._metadata = MetaData()
+        self._data_table = Table(
+            "data", self._metadata,
+            Column("id", Integer, primary_key=True),
+            # timestamp - date/time of the trace point
+            Column("timestamp", String(26), index=True),
+            # base_id - uuid common for all notifications related to one trace
+            Column("base_id", String(255), index=True),
+            # parent_id - uuid of parent element in trace
+            Column("parent_id", String(255), index=True),
+            # trace_id - uuid of current element in trace
+            Column("trace_id", String(255), index=True),
+            Column("project", String(255), index=True),
+            Column("host", String(255), index=True),
+            Column("service", String(255), index=True),
+            # name - trace point name
+            Column("name", String(255), index=True),
+            Column("data", JSON)
+        )
+
+        # FIXME(toabctl): Not the best idea to create the table on every
+        # startup when using the sqlalchemy driver...
+        self._metadata.create_all(self._engine, checkfirst=True)
+
+    @classmethod
+    def get_name(cls):
+        return "sqlalchemy"
+
+    def notify(self, info, context=None):
+        """Write a notification the the database"""
+        data = info.copy()
+        base_id = data.pop("base_id", None)
+        timestamp = data.pop("timestamp", None)
+        parent_id = data.pop("parent_id", None)
+        trace_id = data.pop("trace_id", None)
+        project = data.pop("project", self.project)
+        host = data.pop("host", self.host)
+        service = data.pop("service", self.service)
+        name = data.pop("name", None)
+
+        ins = self._data_table.insert().values(
+            timestamp=timestamp,
+            base_id=base_id,
+            parent_id=parent_id,
+            trace_id=trace_id,
+            project=project,
+            service=service,
+            host=host,
+            name=name,
+            data=jsonutils.dumps(data)
+        )
+        try:
+            self._conn.execute(ins)
+        except Exception:
+            LOG.exception("Can not store osprofiler tracepoint {} "
+                          "(base_id {})".format(trace_id, base_id))
+
+    def get_report(self, base_id):
+        try:
+            from sqlalchemy.sql import select
+        except ImportError:
+            raise exc.CommandError(
+                "To use this command, you should install 'SQLAlchemy'")
+        stmt = select([self._data_table]).where(
+            self._data_table.c.base_id == base_id)
+        results = self._conn.execute(stmt).fetchall()
+        for n in results:
+            timestamp = n["timestamp"]
+            trace_id = n["trace_id"]
+            parent_id = n["parent_id"]
+            name = n["name"]
+            project = n["project"]
+            service = n["service"]
+            host = n["host"]
+            data = jsonutils.loads(n["data"])
+            self._append_results(trace_id, parent_id, name, project, service,
+                                 host, timestamp, data)
+        return self._parse_results()