Add support for all postgres explain options

sqlalchemy_utils/expressions.py

@@ -1,26 +1,66 @@
 import sqlalchemy as sa
 from sqlalchemy.sql import expression
-from sqlalchemy.sql.expression import Executable, ClauseElement, _literal_as_text
+from sqlalchemy.sql.expression import (
+    Executable,
+    ClauseElement,
+    _literal_as_text
+)
 from sqlalchemy.ext.compiler import compiles
 from sqlalchemy_utils.types import TSVectorType
 
 
 class explain(Executable, ClauseElement):
-    def __init__(self, stmt, analyze=False):
+    """
+    Define EXPLAIN element.
+
+    http://www.postgresql.org/docs/devel/static/sql-explain.html
+    """
+    def __init__(
+        self,
+        stmt,
+        analyze=False,
+        verbose=False,
+        costs=True,
+        buffers=False,
+        timing=True,
+        format='text'
+    ):
         self.statement = _literal_as_text(stmt)
         self.analyze = analyze
+        self.verbose = verbose
+        self.costs = costs
+        self.buffers = buffers
+        self.timing = timing
+        self.format = format
 
 
 class explain_analyze(explain):
-    def __init__(self, stmt):
+    def __init__(self, stmt, **kwargs):
-        super(explain_analyze, self).__init__(stmt, analyze=True)
+        super(explain_analyze, self).__init__(
+            stmt,
+            analyze=True,
+            **kwargs
+        )
 
 
 @compiles(explain, 'postgresql')
 def pg_explain(element, compiler, **kw):
     text = "EXPLAIN "
+    options = []
     if element.analyze:
-        text += "ANALYZE "
+        options.append('ANALYZE true')
+    if not element.timing:
+        options.append('TIMING false')
+    if element.buffers:
+        options.append('BUFFERS true')
+    if element.format != 'text':
+        options.append('FORMAT %s' % element.format)
+    if element.verbose:
+        options.append('VERBOSE true')
+    if not element.costs:
+        options.append('COSTS false')
+    if options:
+        text += '(%s) ' % ', '.join(options)
     text += compiler.process(element.statement)
     return text
 

sqlalchemy_utils/types/arrow.py

@@ -9,7 +9,7 @@ try:
 except:
     pass
 from sqlalchemy import types
-from sqlalchemy_utils import ImproperlyConfigured
+from sqlalchemy_utils.exceptions import ImproperlyConfigured
 from .scalar_coercible import ScalarCoercible
 
 

sqlalchemy_utils/types/color.py

@@ -1,6 +1,6 @@
 import six
 from sqlalchemy import types
-from sqlalchemy_utils import ImproperlyConfigured
+from sqlalchemy_utils.exceptions import ImproperlyConfigured
 from .scalar_coercible import ScalarCoercible
 
 colour = None

sqlalchemy_utils/types/ip_address.py

@@ -11,7 +11,7 @@ except ImportError:
 
 
 from sqlalchemy import types
-from sqlalchemy_utils import ImproperlyConfigured
+from sqlalchemy_utils.exceptions import ImproperlyConfigured
 from .scalar_coercible import ScalarCoercible
 
 

sqlalchemy_utils/types/password.py

@@ -1,6 +1,6 @@
 import six
 import weakref
-from sqlalchemy_utils import ImproperlyConfigured
+from sqlalchemy_utils.exceptions import ImproperlyConfigured
 from sqlalchemy import types
 from sqlalchemy.dialects import postgresql, oracle
 from .scalar_coercible import ScalarCoercible

sqlalchemy_utils/types/phone_number.py

@@ -1,6 +1,6 @@
 import six
 from sqlalchemy import types
-from sqlalchemy_utils import ImproperlyConfigured
+from sqlalchemy_utils.exceptions import ImproperlyConfigured
 from .scalar_coercible import ScalarCoercible
 
 

sqlalchemy_utils/types/timezone.py

@@ -1,6 +1,6 @@
 import six
 from sqlalchemy import types
-from sqlalchemy_utils import ImproperlyConfigured
+from sqlalchemy_utils.exceptions import ImproperlyConfigured
 from .scalar_coercible import ScalarCoercible
 
 

tests/test_expressions.py

@@ -27,22 +27,62 @@ class ExpressionTestCase(TestCase):
 
         self.Article = Article
 
+    def assert_startswith(self, query, query_part):
+        assert str(
+            query.compile(dialect=postgresql.dialect())
+        ).startswith(query_part)
+        # Check that query executes properly
+        self.session.execute(query)
+
 
 class TestExplain(ExpressionTestCase):
     def test_render_explain(self):
-        assert str(
+        self.assert_startswith(
-            explain(self.session.query(self.Article)).compile(
+            explain(self.session.query(self.Article)),
-                dialect=postgresql.dialect()
+            'EXPLAIN SELECT'
-            )
+        )
-        ).startswith('EXPLAIN SELECT')
 
     def test_render_explain_with_analyze(self):
-        assert str(
+        self.assert_startswith(
-            explain(self.session.query(self.Article), analyze=True)
+            explain(self.session.query(self.Article), analyze=True),
-            .compile(
+            'EXPLAIN (ANALYZE true) SELECT'
-                dialect=postgresql.dialect()
+        )
-            )
+
-        ).startswith('EXPLAIN ANALYZE SELECT')
+    def test_with_string_as_stmt_param(self):
+        self.assert_startswith(
+            explain('SELECT 1 FROM article'),
+            'EXPLAIN SELECT'
+        )
+
+    def test_format(self):
+        self.assert_startswith(
+            explain('SELECT 1 FROM article', format='json'),
+            'EXPLAIN (FORMAT json) SELECT'
+        )
+
+    def test_timing(self):
+        self.assert_startswith(
+            explain('SELECT 1 FROM article', analyze=True, timing=False),
+            'EXPLAIN (ANALYZE true, TIMING false) SELECT'
+        )
+
+    def test_verbose(self):
+        self.assert_startswith(
+            explain('SELECT 1 FROM article', verbose=True),
+            'EXPLAIN (VERBOSE true) SELECT'
+        )
+
+    def test_buffers(self):
+        self.assert_startswith(
+            explain('SELECT 1 FROM article', analyze=True, buffers=True),
+            'EXPLAIN (ANALYZE true, BUFFERS true) SELECT'
+        )
+
+    def test_costs(self):
+        self.assert_startswith(
+            explain('SELECT 1 FROM article', costs=False),
+            'EXPLAIN (COSTS false) SELECT'
+        )
 
 
 class TestExplainAnalyze(ExpressionTestCase):
@@ -52,7 +92,7 @@ class TestExplainAnalyze(ExpressionTestCase):
             .compile(
                 dialect=postgresql.dialect()
             )
-        ).startswith('EXPLAIN ANALYZE SELECT')
+        ).startswith('EXPLAIN (ANALYZE true) SELECT')
 
 
 class TestMatchTSVector(ExpressionTestCase):