add columns_clause_froms and related use cases

Added new attribute :attr:`_sql.Select.columns_clause_froms` that will
retrieve the FROM list implied by the columns clause of the
:class:`_sql.Select` statement. This differs from the old
:attr:`_sql.Select.froms` collection in that it does not perform any ORM
compilation steps, which necessarily deannotate the FROM elements and do
things like compute joinedloads etc., which makes it not an appropriate
candidate for the :meth:`_sql.Select.select_from` method. Additionally adds
a new parameter
:paramref:`_sql.Select.with_only_columns.maintain_column_froms` that
transfers this collection to :meth:`_sql.Select.select_from` before
replacing the columns collection.

In addition, the :attr:`_sql.Select.froms` is renamed to
:meth:`_sql.Select.get_final_froms`, to stress that this collection is not
a simple accessor and is instead calculated given the full state of the
object, which can be an expensive call when used in an ORM context.

Additionally fixes a regression involving the
:func:`_orm.with_only_columns` function to support applying criteria to
column elements that were replaced with either
:meth:`_sql.Select.with_only_columns` or :meth:`_orm.Query.with_entities` ,
which had broken as part of :ticket:`6503` released in 1.4.19.

Fixes: #6808

Change-Id: Ib5d66cce488bbaca06dab4f68fb5cdaa73e8823e

4 files changed, 187 insertions, 3 deletions

diff --git a/test/orm/test_core_compilation.py b/test/orm/test_core_compilation.py
index e730d9097..2adc43842 100644
--- a/test/orm/test_core_compilation.py
+++ b/test/orm/test_core_compilation.py
@@ -2,6 +2,7 @@ from sqlalchemy import bindparam
 from sqlalchemy import exc
 from sqlalchemy import func
 from sqlalchemy import insert
+from sqlalchemy import inspect
 from sqlalchemy import literal_column
 from sqlalchemy import null
 from sqlalchemy import or_
@@ -21,6 +22,7 @@ from sqlalchemy.orm import query_expression
 from sqlalchemy.orm import relationship
 from sqlalchemy.orm import undefer
 from sqlalchemy.orm import with_expression
+from sqlalchemy.orm import with_loader_criteria
 from sqlalchemy.orm import with_polymorphic
 from sqlalchemy.sql import and_
 from sqlalchemy.sql import sqltypes
@@ -30,6 +32,8 @@ from sqlalchemy.sql.selectable import LABEL_STYLE_TABLENAME_PLUS_COL
 from sqlalchemy.testing import assert_raises_message
 from sqlalchemy.testing import AssertsCompiledSQL
 from sqlalchemy.testing import eq_
+from sqlalchemy.testing import is_
+from sqlalchemy.testing.assertions import expect_raises_message
 from sqlalchemy.testing.fixtures import fixture_session
 from sqlalchemy.testing.util import resolve_lambda
 from .inheritance import _poly_fixtures
@@ -81,7 +85,7 @@ class SelectableTest(QueryTest, AssertsCompiledSQL):
 
         stmt = select(User).filter_by(name="ed")
 
-        eq_(stmt.froms, [self.tables.users])
+        eq_(stmt.get_final_froms(), [self.tables.users])
 
     def test_froms_join(self):
         User, Address = self.classes("User", "Address")
@@ -89,7 +93,7 @@ class SelectableTest(QueryTest, AssertsCompiledSQL):
 
         stmt = select(User).join(User.addresses)
 
-        assert stmt.froms[0].compare(users.join(addresses))
+        assert stmt.get_final_froms()[0].compare(users.join(addresses))
 
     @testing.combinations(
         (
@@ -166,6 +170,115 @@ class SelectableTest(QueryTest, AssertsCompiledSQL):
         eq_(stmt.column_descriptions, expected)
 
 
+class ColumnsClauseFromsTest(QueryTest, AssertsCompiledSQL):
+    __dialect__ = "default"
+
+    def test_exclude_eagerloads(self):
+        User, Address = self.classes("User", "Address")
+
+        stmt = select(User).options(joinedload(User.addresses))
+
+        froms = stmt.columns_clause_froms
+
+        mapper = inspect(User)
+        is_(froms[0], inspect(User).__clause_element__())
+        eq_(
+            froms[0]._annotations,
+            {
+                "entity_namespace": mapper,
+                "parententity": mapper,
+                "parentmapper": mapper,
+            },
+        )
+        eq_(len(froms), 1)
+
+    def test_maintain_annotations_from_table(self):
+        User, Address = self.classes("User", "Address")
+
+        stmt = select(User)
+
+        mapper = inspect(User)
+        froms = stmt.columns_clause_froms
+        is_(froms[0], inspect(User).__clause_element__())
+        eq_(
+            froms[0]._annotations,
+            {
+                "entity_namespace": mapper,
+                "parententity": mapper,
+                "parentmapper": mapper,
+            },
+        )
+        eq_(len(froms), 1)
+
+    def test_maintain_annotations_from_annoated_cols(self):
+        User, Address = self.classes("User", "Address")
+
+        stmt = select(User.id)
+
+        mapper = inspect(User)
+        froms = stmt.columns_clause_froms
+        is_(froms[0], inspect(User).__clause_element__())
+        eq_(
+            froms[0]._annotations,
+            {
+                "entity_namespace": mapper,
+                "parententity": mapper,
+                "parentmapper": mapper,
+            },
+        )
+        eq_(len(froms), 1)
+
+    def test_with_only_columns_unknown_kw(self):
+        User, Address = self.classes("User", "Address")
+
+        stmt = select(User.id)
+
+        with expect_raises_message(TypeError, "unknown parameters: foo"):
+            stmt.with_only_columns(User.id, foo="bar")
+
+    @testing.combinations((True,), (False,))
+    def test_replace_into_select_from_maintains_existing(self, use_flag):
+        User, Address = self.classes("User", "Address")
+
+        stmt = select(User.id).select_from(Address)
+
+        if use_flag:
+            stmt = stmt.with_only_columns(
+                func.count(), maintain_column_froms=True
+            )
+        else:
+            stmt = stmt.select_from(
+                *stmt.columns_clause_froms
+            ).with_only_columns(func.count())
+
+        # Address is maintained in the FROM list
+        self.assert_compile(
+            stmt, "SELECT count(*) AS count_1 FROM addresses, users"
+        )
+
+    @testing.combinations((True,), (False,))
+    def test_replace_into_select_from_with_loader_criteria(self, use_flag):
+        User, Address = self.classes("User", "Address")
+
+        stmt = select(User.id).options(
+            with_loader_criteria(User, User.name == "ed")
+        )
+
+        if use_flag:
+            stmt = stmt.with_only_columns(
+                func.count(), maintain_column_froms=True
+            )
+        else:
+            stmt = stmt.select_from(
+                *stmt.columns_clause_froms
+            ).with_only_columns(func.count())
+
+        self.assert_compile(
+            stmt,
+            "SELECT count(*) AS count_1 FROM users WHERE users.name = :name_1",
+        )
+
+
 class JoinTest(QueryTest, AssertsCompiledSQL):
     __dialect__ = "default"
 
diff --git a/test/orm/test_relationship_criteria.py b/test/orm/test_relationship_criteria.py
index 683267b1c..f9b4335df 100644
--- a/test/orm/test_relationship_criteria.py
+++ b/test/orm/test_relationship_criteria.py
@@ -5,6 +5,7 @@ from sqlalchemy import Column
 from sqlalchemy import DateTime
 from sqlalchemy import event
 from sqlalchemy import ForeignKey
+from sqlalchemy import func
 from sqlalchemy import Integer
 from sqlalchemy import orm
 from sqlalchemy import select
@@ -25,6 +26,7 @@ from sqlalchemy.orm import with_loader_criteria
 from sqlalchemy.orm.decl_api import declared_attr
 from sqlalchemy.testing import eq_
 from sqlalchemy.testing.assertsql import CompiledSQL
+from sqlalchemy.testing.fixtures import fixture_session
 from test.orm import _fixtures
 
 
@@ -172,6 +174,39 @@ class LoaderCriteriaTest(_Fixtures, testing.AssertsCompiledSQL):
             "FROM users WHERE users.name != :name_1",
         )
 
+    def test_criteria_post_replace(self, user_address_fixture):
+        User, Address = user_address_fixture
+
+        stmt = (
+            select(User)
+            .select_from(User)
+            .options(with_loader_criteria(User, User.name != "name"))
+            .with_only_columns(func.count())
+        )
+
+        self.assert_compile(
+            stmt,
+            "SELECT count(*) AS count_1 FROM users "
+            "WHERE users.name != :name_1",
+        )
+
+    def test_criteria_post_replace_legacy(self, user_address_fixture):
+        User, Address = user_address_fixture
+
+        s = fixture_session()
+        stmt = (
+            s.query(User)
+            .select_from(User)
+            .options(with_loader_criteria(User, User.name != "name"))
+            .with_entities(func.count())
+        )
+
+        self.assert_compile(
+            stmt,
+            "SELECT count(*) AS count_1 FROM users "
+            "WHERE users.name != :name_1",
+        )
+
     def test_select_from_mapper_mapper_criteria(self, user_address_fixture):
         User, Address = user_address_fixture
 
diff --git a/test/sql/test_deprecations.py b/test/sql/test_deprecations.py
index 44135e373..9b74ab1fa 100644
--- a/test/sql/test_deprecations.py
+++ b/test/sql/test_deprecations.py
@@ -451,6 +451,17 @@ class SelectableTest(fixtures.TestBase, AssertsCompiledSQL):
             "deprecated"
         )
 
+    def test_froms_renamed(self):
+        t1 = table("t1", column("q"))
+
+        stmt = select(t1)
+
+        with testing.expect_deprecated(
+            r"The Select.froms attribute is moved to the "
+            r"Select.get_final_froms\(\) method."
+        ):
+            eq_(stmt.froms, [t1])
+
     def test_select_list_argument(self):
 
         with testing.expect_deprecated_20(
diff --git a/test/sql/test_selectable.py b/test/sql/test_selectable.py
index 5a94d4038..b76873490 100644
--- a/test/sql/test_selectable.py
+++ b/test/sql/test_selectable.py
@@ -1,5 +1,4 @@
 """Test various algorithmic properties of selectables."""
-
 from sqlalchemy import and_
 from sqlalchemy import bindparam
 from sqlalchemy import Boolean
@@ -590,6 +589,32 @@ class SelectableTest(
             "table1.col3, table1.colx FROM table1) AS anon_1",
         )
 
+    @testing.combinations(
+        (
+            [table1.c.col1],
+            [table1.join(table2)],
+            [table1.join(table2)],
+            [table1],
+        ),
+        ([table1], [table2], [table1, table2], [table1]),
+        (
+            [table1.c.col1, table2.c.col1],
+            [],
+            [table1, table2],
+            [table1, table2],
+        ),
+    )
+    def test_froms_accessors(
+        self, cols_expr, select_from, exp_final_froms, exp_cc_froms
+    ):
+        """tests for #6808"""
+        s1 = select(*cols_expr).select_from(*select_from)
+
+        for ff, efp in util.zip_longest(s1.get_final_froms(), exp_final_froms):
+            assert ff.compare(efp)
+
+        eq_(s1.columns_clause_froms, exp_cc_froms)
+
     def test_scalar_subquery_from_subq_same_source(self):
         s1 = select(table1.c.col1)