Implement support for functions as FROM with columns clause support

WIP

Fixes: #3566
Change-Id: I5b093b72533ef695293e737eb75850b9713e5e03

1 files changed, 250 insertions, 22 deletions

diff --git a/lib/sqlalchemy/sql/functions.py b/lib/sqlalchemy/sql/functions.py
index a9ea98d04..f493b08db 100644
--- a/lib/sqlalchemy/sql/functions.py
+++ b/lib/sqlalchemy/sql/functions.py
@@ -17,22 +17,25 @@ from . import sqltypes
 from . import util as sqlutil
 from .base import ColumnCollection
 from .base import Executable
+from .base import Generative
 from .base import HasMemoized
 from .elements import _type_from_args
 from .elements import BinaryExpression
 from .elements import BindParameter
 from .elements import Cast
 from .elements import ClauseList
+from .elements import ColumnClause
 from .elements import ColumnElement
 from .elements import Extract
 from .elements import FunctionFilter
 from .elements import Grouping
 from .elements import literal_column
+from .elements import NamedColumn
 from .elements import Over
 from .elements import WithinGroup
-from .selectable import Alias
 from .selectable import FromClause
 from .selectable import Select
+from .selectable import TableValuedAlias
 from .visitors import InternalTraversal
 from .visitors import TraversibleType
 from .. import util
@@ -63,7 +66,7 @@ def register_function(identifier, fn, package="_default"):
     reg[identifier] = fn
 
 
-class FunctionElement(Executable, ColumnElement, FromClause):
+class FunctionElement(Executable, ColumnElement, FromClause, Generative):
     """Base for SQL function-oriented constructs.
 
     .. seealso::
@@ -80,11 +83,19 @@ class FunctionElement(Executable, ColumnElement, FromClause):
 
     """
 
-    _traverse_internals = [("clause_expr", InternalTraversal.dp_clauseelement)]
+    _traverse_internals = [
+        ("clause_expr", InternalTraversal.dp_clauseelement),
+        ("_table_valued", InternalTraversal.dp_clauseelement_tuple),
+        ("_table_values_named", InternalTraversal.dp_boolean),
+        ("_with_ordinality", InternalTraversal.dp_boolean),
+    ]
 
     packagenames = ()
 
     _has_args = False
+    _table_valued = None
+    _table_values_named = False
+    _with_ordinality = False
 
     def __init__(self, *clauses, **kwargs):
         r"""Construct a :class:`.FunctionElement`.
@@ -123,6 +134,171 @@ class FunctionElement(Executable, ColumnElement, FromClause):
             self, multiparams, params, execution_options
         )
 
+    def scalar_table_valued(self, name, type_=None):
+        """Return a column expression that's against this
+        :class:`_functions.FunctionElement` as a scalar
+        table-valued expression.
+
+        The returned expression is similar to that returned by a single column
+        accessed off of a :meth:`_functions.FunctionElement.table_valued`
+        construct, except no FROM clause is generated; the function is rendered
+        in the similar way as a scalar subquery.
+
+        E.g.::
+
+            >>> from sqlalchemy import func, select
+            >>> fn = func.jsonb_each("{'k', 'v'}").scalar_table_valued("key")
+            >>> print(select(fn))
+            SELECT (jsonb_each(:jsonb_each_1)).key
+
+        .. versionadded:: 1.4.0b2
+
+        .. seealso::
+
+            :meth:`_functions.FunctionElement.table_valued`
+
+            :meth:`_functions.FunctionElement.named_table_valued`
+
+            :meth:`_functions.FunctionElement.alias`
+
+            :meth:`_functions.FunctionElement.column_valued`
+
+        """  # noqa E501
+
+        return ScalarFunctionColumn(self, name, type_)
+
+    def table_valued(self, *expr, **kw):
+        """Return a :class:`_sql.TableValuedAlias` representation of this
+        :class:`_functions.FunctionElement` with table-valued expressions added.
+
+        e.g.::
+
+            >>> fn = (
+            ...     func.generate_series(1, 5).
+            ...     table_valued("value", "start", "stop", "step")
+            ... )
+
+            >>> print(select(fn))
+            SELECT anon_1.value, anon_1.start, anon_1.stop, anon_1.step
+            FROM generate_series(:generate_series_1, :generate_series_2) AS anon_1
+
+            >>> print(select(fn.c.value, fn.c.stop).where(fn.c.value > 2))
+            SELECT anon_1.value, anon_1.stop
+            FROM generate_series(:generate_series_1, :generate_series_2) AS anon_1
+            WHERE anon_1.value > :value_1
+
+        A WITH ORDINALITY expression may be generated by passing the keyword
+        argument "with_ordinality"::
+
+            >>> fn = func.generate_series(4, 1, -1).table_valued("gen", with_ordinality="ordinality")
+            >>> print(select(fn))
+
+        .. versionadded:: 1.4.0b2
+
+        .. seealso::
+
+            :meth:`_functions.FunctionElement.table_valued`
+
+            :meth:`_functions.FunctionElement.named_table_valued`
+
+            :meth:`_functions.FunctionElement.alias`
+
+
+        """  # noqa 501
+
+        new_func = self._generate()
+        new_func._table_valued = [
+            coercions.expect(roles.StrAsPlainColumnRole, elem) for elem in expr
+        ]
+
+        with_ordinality = kw.pop("with_ordinality", None)
+        if with_ordinality:
+            new_func._table_valued += (
+                coercions.expect(roles.StrAsPlainColumnRole, with_ordinality),
+            )
+            new_func._with_ordinality = True
+
+        return new_func.alias()
+
+    def named_table_valued(self, *expr, **kw):
+        """Return a :class:`_sql.TableValuedAlias` representation of this
+        :class:`_functions.FunctionElement` with named table-valued
+        expressions added.
+
+        E.g.::
+
+
+            >>> fn = (
+            ...     func.json_to_recordset(
+            ...             '[{"a":1,"b":"foo"},{"a":"2","c":"bar"}]'
+            ...     )
+            ...     .named_table_valued(column("a", Integer), column("b", String))
+            ... )
+            >>> print(select(fn.c.a, fn.c.b))
+            SELECT anon_1.a, anon_1.b
+            FROM json_to_recordset(:json_to_recordset_1) AS anon_1(a INTEGER, b VARCHAR)
+
+        A WITH ORDINALITY expression may be generated by passing the keyword
+        argument "with_ordinality"::
+
+            >>> fn = (
+            ...     func.json_object_keys('{"a1":"1","a2":"2","a3":"3"}')
+            ...     .named_table_valued("keys", with_ordinality="n")
+            ... )
+
+            >>> print(select(fn))
+
+        .. seealso::
+
+            :meth:`_functions.FunctionElement.table_valued`
+
+            :meth:`_functions.FunctionElement.named_table_valued`
+
+            :meth:`_functions.FunctionElement.alias`
+
+
+        """  # noqa E501
+
+        new_func = self._generate()
+        new_func._table_valued = [
+            coercions.expect(roles.StrAsPlainColumnRole, elem) for elem in expr
+        ]
+        with_ordinality = kw.pop("with_ordinality", None)
+
+        if with_ordinality:
+            new_func._table_valued += (
+                coercions.expect(roles.StrAsPlainColumnRole, with_ordinality),
+            )
+            new_func._with_ordinality = True
+
+        new_func._table_values_named = True
+        return new_func.alias()
+
+    def column_valued(self, name=None):
+        """Return this :class:`_function.FunctionElement` as a column expression that
+        selects from itself as a FROM clause.
+
+        E.g.::
+
+            >>> from sqlalchemy import select, func
+            >>> gs = func.generate_series(1, 5, -1).column_valued()
+            >>> print(select(gs))
+            SELECT anon_1
+            FROM generate_series(:generate_series_1, :generate_series_2, :generate_series_3) AS anon_1
+
+        This is shorthand for::
+
+            gs = func.generate_series(1, 5, -1).alias().column
+
+
+        .. seealso::
+
+            :meth:`_functions.FunctionElement.alias`
+
+        """  # noqa 501
+
+        return self.alias(name=name).column
+
     @property
     def columns(self):
         r"""The set of columns exported by this :class:`.FunctionElement`.
@@ -142,8 +318,15 @@ class FunctionElement(Executable, ColumnElement, FromClause):
 
 
         """
-        col = self.label(None)
-        return ColumnCollection(columns=[(col.key, col)])
+        if self._table_valued:
+            cols = [
+                ColumnClause(elem) if isinstance(elem, str) else elem
+                for elem in self._table_valued
+            ]
+        else:
+            cols = [self.label(None)]
+
+        return ColumnCollection(columns=[(col.key, col) for col in cols])
 
     @HasMemoized.memoized_attribute
     def clauses(self):
@@ -305,37 +488,65 @@ class FunctionElement(Executable, ColumnElement, FromClause):
 
         return None
 
-    def alias(self, name=None, flat=False):
+    def alias(self, name=None):
         r"""Produce a :class:`_expression.Alias` construct against this
         :class:`.FunctionElement`.
 
+        .. note::
+
+            The :meth:`_functions.FunctionElement.alias` method is part of the
+            mechanism by which "table valued" SQL functions are created.
+            However, most use cases are covered by higher level methods on
+            :class:`_functions.FunctionElement` including
+            :meth:`_functions.FunctionElement.table_valued`,
+            :meth:`_functions.FunctionElement.named_table_valued`, and
+            :meth:`_functions.FunctionElement.column_valued`.
+
         This construct wraps the function in a named alias which
         is suitable for the FROM clause, in the style accepted for example
-        by PostgreSQL.
+        by PostgreSQL.  A column expression is also provided using the
+        special ``.column`` attribute, which may
+        be used to refer to the output of the function as a scalar value
+        in the columns or where clause, for a backend such as PostgreSQL.
+
+        For a full table-valued expression, use the
+        :meth:`_function.FunctionElement.table_valued` method first to
+        establish named columns.
 
         e.g.::
 
-            from sqlalchemy.sql import column
+            >>> from sqlalchemy import func, select, column
+            >>> data_view = func.unnest([1, 2, 3]).alias("data_view")
+            >>> print(select(data_view.column))
+            SELECT data_view
+            FROM unnest(:unnest_1) AS data_view
 
-            stmt = select(column('data_view')).\
-                select_from(SomeTable).\
-                select_from(func.unnest(SomeTable.data).alias('data_view')
-            )
+        The :meth:`_functions.FunctionElement.column_valued` method provides
+        a shortcut for the above pattern::
+
+            >>> data_view = func.unnest([1, 2, 3]).column_valued("data_view")
+            >>> print(select(data_view))
+            SELECT data_view
+            FROM unnest(:unnest_1) AS data_view
 
-        Would produce:
+        .. versionadded:: 1.4.0b2  Added the ``.column`` accessor
 
-        .. sourcecode:: sql
+        .. seealso::
 
-            SELECT data_view
-            FROM sometable, unnest(sometable.data) AS data_view
+            :meth:`_functions.FunctionElement.table_valued`
+
+            :meth:`_functions.FunctionElement.named_table_valued`
+
+            :meth:`_functions.FunctionElement.scalar_table_valued`
+
+            :meth:`_functions.FunctionElement.column_valued`
 
-        .. versionadded:: 0.9.8 The :meth:`.FunctionElement.alias` method
-           is now supported.  Previously, this method's behavior was
-           undefined and did not behave consistently across versions.
 
         """
 
-        return Alias._construct(self, name)
+        return TableValuedAlias._construct(
+            self, name, named=self._table_values_named
+        )
 
     def select(self):
         """Produce a :func:`_expression.select` construct
@@ -441,6 +652,24 @@ class FunctionAsBinary(BinaryExpression):
         self.sql_function.clauses.clauses[self.right_index - 1] = value
 
 
+class ScalarFunctionColumn(NamedColumn):
+    __visit_name__ = "scalar_function_column"
+
+    _traverse_internals = [
+        ("name", InternalTraversal.dp_anon_name),
+        ("type", InternalTraversal.dp_type),
+        ("fn", InternalTraversal.dp_clauseelement),
+    ]
+
+    is_literal = False
+    table = None
+
+    def __init__(self, fn, name, type_=None):
+        self.fn = fn
+        self.name = name
+        self.type = sqltypes.to_instance(type_)
+
+
 class _FunctionGenerator(object):
     """Generate SQL function expressions.
 
@@ -586,10 +815,9 @@ class Function(FunctionElement):
 
         func.mypackage.some_function(col1, col2)
 
-
     .. seealso::
 
-        :ref:`coretutorial_functions`
+        :ref:`tutorial_functions` - in the :ref:`unified_tutorial`
 
         :data:`.func` - namespace which produces registered or ad-hoc
         :class:`.Function` instances.