Changed how TableSample handles method and arg.pr/277

Instead of passing a float percentage and a string method,
accept either a non Function value wrapped in func.system,
or a Function, which is rendered as is.

2 files changed, 17 insertions, 15 deletions

diff --git a/lib/sqlalchemy/sql/compiler.py b/lib/sqlalchemy/sql/compiler.py
index ab1b42b2f..dbf919cf3 100644
--- a/lib/sqlalchemy/sql/compiler.py
+++ b/lib/sqlalchemy/sql/compiler.py
@@ -1325,10 +1325,9 @@ class SQLCompiler(Compiled):
         return "LATERAL %s" % self.visit_alias(lateral, **kw)
 
     def visit_tablesample(self, tablesample, asfrom=False, **kw):
-        text = "%s TABLESAMPLE %s(%s)" % (
+        text = "%s TABLESAMPLE %s" % (
             self.visit_alias(tablesample, asfrom=True, **kw),
-            tablesample.method,
-            tablesample.arg)
+            tablesample._get_method()._compiler_dispatch(self, **kw))
 
         if tablesample.seed is not None:
             text += " REPEATABLE (%s)" % (
diff --git a/lib/sqlalchemy/sql/selectable.py b/lib/sqlalchemy/sql/selectable.py
index 0bb590bd0..c554429be 100644
--- a/lib/sqlalchemy/sql/selectable.py
+++ b/lib/sqlalchemy/sql/selectable.py
@@ -183,7 +183,7 @@ def lateral(selectable, name=None):
     return selectable.lateral(name=name)
 
 
-def tablesample(selectable, arg, name=None, method=None, seed=None):
+def tablesample(selectable, sampling, name=None, seed=None):
     """Return a :class:`.TableSample` object.
 
     :class:`.TableSample` is an :class:`.Alias` subclass that represents
@@ -193,15 +193,13 @@ def tablesample(selectable, arg, name=None, method=None, seed=None):
     percentage of rows from a table. It supports multiple sampling methods,
     most commonly BERNOULLI and SYSTEM.
 
-    :param arg: a ``float`` percentage between 0 and 100.
-
-    :param method: string name of the method to use.
-        Commonly accepted methods are ``"BERNOULLI"`` and ``"SYSTEM"``.
+    :param sampling: a ``float`` percentage between 0 and 100 or
+        :class:`.functions.Function`.
 
     :param seed: any real-valued SQL expression.
 
     """
-    return selectable.sample(arg, name=name, method=method, seed=seed)
+    return selectable.sample(sampling, name=name, seed=seed)
 
 
 class Selectable(ClauseElement):
@@ -471,14 +469,14 @@ class FromClause(Selectable):
         """
         return Lateral(self, name)
 
-    def sample(self, arg, name=None, method=None, seed=None):
+    def sample(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.
 
         """
-        return TableSample(self, arg, name, method, seed)
+        return TableSample(self, sampling, name, seed)
 
     def is_derived_from(self, fromclause):
         """Return True if this FromClause is 'derived' from the given
@@ -1285,15 +1283,20 @@ class TableSample(Alias):
 
     __visit_name__ = 'tablesample'
 
-    def __init__(self, selectable, arg,
+    def __init__(self, selectable, sampling,
                  name=None,
-                 method=None,
                  seed=None):
-        self.arg = arg
-        self.method = method or 'SYSTEM'
+        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.