Add analyze function

CHANGES.rst

@@ -8,6 +8,7 @@ Here you can see the full list of changes between each SQLAlchemy-Utils release.
 ^^^^^^^^^^^^^^^^^^^^
 
 - Added explain and explain_analyze expressions
+- Added analyze function
 
 
 0.26.16 (2014-09-09)

docs/database_helpers.rst

@@ -5,6 +5,12 @@ Database helpers
 .. module:: sqlalchemy_utils.functions
 
 
+analyze
+^^^^^^^
+
+.. autofunction:: analyze
+
+
 database_exists
 ^^^^^^^^^^^^^^^
 

sqlalchemy_utils/__init__.py

@@ -4,6 +4,7 @@ from .decorators import generates
 from .exceptions import ImproperlyConfigured
 from .expression_parser import ExpressionParser
 from .functions import (
+    analyze,
     create_database,
     create_mock_engine,
     database_exists,
@@ -80,6 +81,7 @@ __version__ = '0.26.16'
 
 __all__ = (
     aggregated,
+    analyze,
     auto_delete_orphans,
     batch_fetch,
     coercion_listener,

sqlalchemy_utils/functions/__init__.py

@@ -3,6 +3,7 @@ from .mock import create_mock_engine, mock_engine
 from .render import render_expression, render_statement
 from .sort_query import sort_query, QuerySorterException
 from .database import (
+    analyze,
     create_database,
     database_exists,
     drop_database,

sqlalchemy_utils/functions/database.py

@@ -1,9 +1,84 @@
-from sqlalchemy.engine.url import make_url
-import sqlalchemy as sa
-from sqlalchemy.exc import ProgrammingError, OperationalError
 import os
 from copy import copy
 
+import sqlalchemy as sa
+from sqlalchemy.engine.url import make_url
+from sqlalchemy.exc import ProgrammingError, OperationalError
+from sqlalchemy_utils.expressions import explain_analyze
+
+
+class PlanAnalysis(object):
+    def __init__(self, plan):
+        self.plan = plan
+
+    @property
+    def node_types(self):
+        types = [self.plan['Node Type']]
+        if 'Plans' in self.plan:
+            for plan in self.plan['Plans']:
+                analysis = PlanAnalysis(plan)
+                types.extend(analysis.node_types)
+        return types
+
+
+class QueryAnalysis(object):
+    def __init__(self, result_set):
+        self.plan = result_set[0]['Plan']
+        self.runtime = result_set[0]['Total Runtime']
+
+    @property
+    def node_types(self):
+        return list(PlanAnalysis(self.plan).node_types)
+
+    def __repr__(self):
+        return '<QueryAnalysis runtime=%r>' % self.runtime
+
+
+def analyze(conn, query):
+    """
+    Analyze query using given connection and return :class:`QueryAnalysis`
+    object. Analysis is performed using database specific EXPLAIN ANALYZE
+    construct and then examining the results into structured format. Currently
+    only PostgreSQL is supported.
+
+
+    Getting query runtime (in database level) ::
+
+
+        from sqlalchemy_utils import analyze
+
+
+        analysis = analyze(conn, 'SELECT * FROM article')
+        analysis.runtime  # runtime as milliseconds
+
+
+    Analyze can be very useful when testing that query doesn't issue a
+    sequential scan (scanning all rows in table). You can for example write
+    simple performance tests this way.::
+
+
+        query = (
+            session.query(Article.name)
+            .order_by(Article.name)
+            .limit(10)
+        )
+        analysis = analyze(self.connection, query)
+        analysis.node_types  # [u'Limit', u'Index Only Scan']
+
+        assert 'Seq Scan' not in analysis.node_types
+
+
+    .. versionadded: 0.26.17
+
+    :param conn: SQLAlchemy Connection object
+    :param query: SQLAlchemy Query object or query as a string
+    """
+    return QueryAnalysis(
+        conn.execute(
+            explain_analyze(query, buffers=True, format='json')
+        ).scalar()
+    )
+
 
 def escape_like(string, escape_char='*'):
     """

tests/__init__.py

@@ -96,7 +96,7 @@ class TestCase(object):
         class Article(self.Base):
             __tablename__ = 'article'
             id = sa.Column(sa.Integer, primary_key=True)
-            name = sa.Column(sa.Unicode(255))
+            name = sa.Column(sa.Unicode(255), index=True)
             category_id = sa.Column(sa.Integer, sa.ForeignKey(Category.id))
 
             category = sa.orm.relationship(

tests/functions/test_analyze.py

@@ -0,0 +1,29 @@
+from sqlalchemy_utils import analyze
+from tests import TestCase
+
+
+class TestAnalyzeWithPostgres(TestCase):
+    dns = 'postgres://postgres@localhost/sqlalchemy_utils_test'
+
+    def test_runtime(self):
+        query = self.session.query(self.Article)
+        assert analyze(self.connection, query).runtime
+
+    def test_node_types_with_join(self):
+        query = (
+            self.session.query(self.Article)
+            .join(self.Article.category)
+        )
+        analysis = analyze(self.connection, query)
+        assert analysis.node_types == [
+            u'Hash Join', u'Seq Scan', u'Hash', u'Seq Scan'
+        ]
+
+    def test_node_types_with_index_only_scan(self):
+        query = (
+            self.session.query(self.Article.name)
+            .order_by(self.Article.name)
+            .limit(10)
+        )
+        analysis = analyze(self.connection, query)
+        assert analysis.node_types == [u'Limit', u'Index Only Scan']