Merge "Add TABLESAMPLE clause support."

9 files changed, 206 insertions, 2 deletions

diff --git a/doc/build/changelog/changelog_11.rst b/doc/build/changelog/changelog_11.rst
index 7194ac116..534c3993c 100644
--- a/doc/build/changelog/changelog_11.rst
+++ b/doc/build/changelog/changelog_11.rst
@@ -22,6 +22,17 @@
     :version: 1.1.0b1
 
     .. change::
+        :tags: feature, sql
+        :tickets: 3718
+
+        Added TABLESAMPLE support via the new :meth:`.FromClause.tablesample`
+        method and standalone function.  Pull request courtesy Ilja Everilä.
+
+        .. seealso::
+
+            :ref:`change_3718`
+
+    .. change::
         :tags: feature, orm, ext
 
         A new ORM extension :ref:`indexable_toplevel` is added, which allows
diff --git a/doc/build/changelog/migration_11.rst b/doc/build/changelog/migration_11.rst
index 766b06f2e..f6e04225d 100644
--- a/doc/build/changelog/migration_11.rst
+++ b/doc/build/changelog/migration_11.rst
@@ -1150,6 +1150,32 @@ selectable, e.g. lateral correlation::
 
 :ticket:`2857`
 
+.. _change_3718:
+
+Support for TABLESAMPLE
+-----------------------
+
+The SQL standard TABLESAMPLE can be rendered using the
+:meth:`.FromClause.tablesample` method, which returns a :class:`.TableSample`
+construct similar to an alias::
+
+    from sqlalchemy import func
+
+    selectable = people.tablesample(
+                func.bernoulli(1),
+                name='alias',
+                seed=func.random())
+    stmt = select([selectable.c.people_id])
+
+Assuming ``people`` with a column ``people_id``, the above
+statement would render as::
+
+    SELECT alias.people_id FROM
+    people AS alias TABLESAMPLE bernoulli(:bernoulli_1)
+    REPEATABLE (random())
+
+:ticket:`3718`
+
 .. _change_3216:
 
 The ``.autoincrement`` directive is no longer implicitly enabled for a composite primary key column
diff --git a/doc/build/core/selectable.rst b/doc/build/core/selectable.rst
index 263594d24..f4860165f 100644
--- a/doc/build/core/selectable.rst
+++ b/doc/build/core/selectable.rst
@@ -33,6 +33,8 @@ elements are themselves :class:`.ColumnElement` subclasses).
 
 .. autofunction:: sqlalchemy.sql.expression.table
 
+.. autofunction:: tablesample
+
 .. autofunction:: union
 
 .. autofunction:: union_all
@@ -93,6 +95,10 @@ elements are themselves :class:`.ColumnElement` subclasses).
    :members:
    :inherited-members:
 
+.. autoclass:: TableSample
+   :members:
+   :inherited-members:
+
 .. autoclass:: TextAsFrom
    :members:
    :inherited-members:
diff --git a/lib/sqlalchemy/__init__.py b/lib/sqlalchemy/__init__.py
index b1d240edf..eca9f9c7f 100644
--- a/lib/sqlalchemy/__init__.py
+++ b/lib/sqlalchemy/__init__.py
@@ -45,6 +45,7 @@ from .sql import (
     select,
     subquery,
     table,
+    tablesample,
     text,
     true,
     tuple_,
diff --git a/lib/sqlalchemy/sql/__init__.py b/lib/sqlalchemy/sql/__init__.py
index 7f7abacc2..a9b842e1a 100644
--- a/lib/sqlalchemy/sql/__init__.py
+++ b/lib/sqlalchemy/sql/__init__.py
@@ -18,6 +18,7 @@ from .expression import (
     Select,
     Selectable,
     TableClause,
+    TableSample,
     Update,
     alias,
     and_,
@@ -59,6 +60,7 @@ from .expression import (
     select,
     subquery,
     table,
+    tablesample,
     text,
     true,
     True_,
diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index 94c7db20a..5e537dfdc 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -1351,6 +1351,17 @@ class SQLCompiler(Compiled):
         kw['lateral'] = True
         return "LATERAL %s" % self.visit_alias(lateral, **kw)
 
+    def visit_tablesample(self, tablesample, asfrom=False, **kw):
+        text = "%s TABLESAMPLE %s" % (
+            self.visit_alias(tablesample, asfrom=True, **kw),
+            tablesample._get_method()._compiler_dispatch(self, **kw))
+
+        if tablesample.seed is not None:
+            text += " REPEATABLE (%s)" % (
+                tablesample.seed._compiler_dispatch(self, **kw))
+
+        return text
+
     def get_render_as_alias_suffix(self, alias_name_text):
         return " AS " + alias_name_text
 
diff --git a/lib/sqlalchemy/sql/expression.py b/lib/sqlalchemy/sql/expression.py
index 97f74d4e4..cbb123ec6 100644
--- a/lib/sqlalchemy/sql/expression.py
+++ b/lib/sqlalchemy/sql/expression.py
@@ -26,7 +26,8 @@ __all__ = [
     'nullslast',
     'or_', 'outparam', 'outerjoin', 'over', 'select', 'subquery',
     'table', 'text',
-    'tuple_', 'type_coerce', 'union', 'union_all', 'update', 'within_group']
+    'tuple_', 'type_coerce', 'union', 'union_all', 'update', 'within_group',
+    'TableSample', 'tablesample']
 
 
 from .visitors import Visitable
@@ -49,7 +50,7 @@ from .base import ColumnCollection, Generative, Executable, \
 from .selectable import Alias, Join, Select, Selectable, TableClause, \
     CompoundSelect, CTE, FromClause, FromGrouping, Lateral, SelectBase, \
     alias, GenerativeSelect, subquery, HasCTE, HasPrefixes, HasSuffixes, \
-    lateral, Exists, ScalarSelect, TextAsFrom
+    lateral, Exists, ScalarSelect, TextAsFrom, TableSample, tablesample
 
 
 from .dml import Insert, Update, Delete, UpdateBase, ValuesBase
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index 741aa9bea..9770b11bb 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -183,6 +183,51 @@ def lateral(selectable, name=None):
     return _interpret_as_from(selectable).lateral(name=name)
 
 
+def tablesample(selectable, sampling, name=None, seed=None):
+    """Return a :class:`.TableSample` object.
+
+    :class:`.TableSample` is an :class:`.Alias` subclass that represents
+    a table with the TABLESAMPLE clause applied to it.
+    :func:`~.expression.tablesample`
+    is also available from the :class:`.FromClause` class via the
+    :meth:`.FromClause.tablesample` method.
+
+    The TABLESAMPLE clause allows selecting a randomly selected approximate
+    percentage of rows from a table. It supports multiple sampling methods,
+    most commonly BERNOULLI and SYSTEM.
+
+    e.g.::
+
+        from sqlalchemy import func
+
+        selectable = people.tablesample(
+                    func.bernoulli(1),
+                    name='alias',
+                    seed=func.random())
+        stmt = select([selectable.c.people_id])
+
+    Assuming ``people`` with a column ``people_id``, the above
+    statement would render as::
+
+        SELECT alias.people_id FROM
+        people AS alias TABLESAMPLE bernoulli(:bernoulli_1)
+        REPEATABLE (random())
+
+    .. versionadded:: 1.1
+
+    :param sampling: a ``float`` percentage between 0 and 100 or
+        :class:`.functions.Function`.
+
+    :param name: optional alias name
+
+    :param seed: any real-valued SQL expression.  When specified, the
+     REPEATABLE sub-clause is also rendered.
+
+    """
+    return _interpret_as_from(selectable).tablesample(
+        sampling, name=name, seed=seed)
+
+
 class Selectable(ClauseElement):
     """mark a class as being selectable"""
     __visit_name__ = 'selectable'
@@ -474,6 +519,21 @@ class FromClause(Selectable):
         """
         return Lateral(self, name)
 
+    def tablesample(self, sampling, name=None, seed=None):
+        """Return a TABLESAMPLE alias of this :class:`.FromClause`.
+
+        The return value is the :class:`.TableSample` construct also
+        provided by the top-level :func:`~.expression.tablesample` function.
+
+        .. versionadded:: 1.1
+
+        .. seealso::
+
+            :func:`~.expression.tablesample` - usage guidelines and parameters
+
+        """
+        return TableSample(self, sampling, name, seed)
+
     def is_derived_from(self, fromclause):
         """Return True if this FromClause is 'derived' from the given
         FromClause.
@@ -1268,6 +1328,38 @@ class Lateral(Alias):
     __visit_name__ = 'lateral'
 
 
+class TableSample(Alias):
+    """Represent a TABLESAMPLE clause.
+
+    This object is constructed from the :func:`~.expression.tablesample` module
+    level function as well as the :meth:`.FromClause.tablesample` method available
+    on all :class:`.FromClause` subclasses.
+
+    .. versionadded:: 1.1
+
+    .. seealso::
+
+        :func:`~.expression.tablesample`
+
+    """
+
+    __visit_name__ = 'tablesample'
+
+    def __init__(self, selectable, sampling,
+                 name=None,
+                 seed=None):
+        self.sampling = sampling
+        self.seed = seed
+        super(TableSample, self).__init__(selectable, name=name)
+
+    @util.dependencies("sqlalchemy.sql.functions")
+    def _get_method(self, functions):
+        if isinstance(self.sampling, functions.Function):
+            return self.sampling
+        else:
+            return functions.func.system(self.sampling)
+
+
 class CTE(Generative, HasSuffixes, Alias):
     """Represent a Common Table Expression.
 
diff --git a/test/sql/test_tablesample.py b/test/sql/test_tablesample.py
new file mode 100644
index 000000000..b2dddaf8c
--- /dev/null
+++ b/test/sql/test_tablesample.py
@@ -0,0 +1,54 @@
+from sqlalchemy.testing import fixtures
+from sqlalchemy.testing import AssertsCompiledSQL, assert_raises_message
+from sqlalchemy.sql import select, func, text
+from sqlalchemy.engine import default
+from sqlalchemy import exc
+from sqlalchemy import Table, Integer, String, Column
+from sqlalchemy import tablesample
+
+
+class TableSampleTest(fixtures.TablesTest, AssertsCompiledSQL):
+    __dialect__ = default.DefaultDialect(supports_native_boolean=True)
+
+    run_setup_bind = None
+
+    run_create_tables = None
+
+    @classmethod
+    def define_tables(cls, metadata):
+        Table('people', metadata,
+              Column('people_id', Integer, primary_key=True),
+              Column('age', Integer),
+              Column('name', String(30)))
+
+    def test_standalone(self):
+        table1 = self.tables.people
+
+        # no special alias handling even though clause is not in the
+        # context of a FROM clause
+        self.assert_compile(
+            tablesample(table1, 1, name='alias'),
+            'people AS alias TABLESAMPLE system(:system_1)'
+        )
+
+        self.assert_compile(
+            table1.tablesample(1, name='alias'),
+            'people AS alias TABLESAMPLE system(:system_1)'
+        )
+
+        self.assert_compile(
+            tablesample(table1, func.bernoulli(1), name='alias',
+                        seed=func.random()),
+            'people AS alias TABLESAMPLE bernoulli(:bernoulli_1) '
+            'REPEATABLE (random())'
+        )
+
+    def test_select_from(self):
+        table1 = self.tables.people
+
+        self.assert_compile(
+            select([table1.tablesample(text('1'), name='alias').c.people_id]),
+            'SELECT alias.people_id FROM '
+            'people AS alias TABLESAMPLE system(1)'
+        )
+