1from sqlalchemy import Column
2from sqlalchemy import ForeignKey
3from sqlalchemy import Integer
4from sqlalchemy import join
5from sqlalchemy import lateral
6from sqlalchemy import String
7from sqlalchemy import Table
8from sqlalchemy import true
9from sqlalchemy.engine import default
10from sqlalchemy.sql import func
11from sqlalchemy.sql import select
12from sqlalchemy.testing import AssertsCompiledSQL
13from sqlalchemy.testing import fixtures
14
15
16class LateralTest(fixtures.TablesTest, AssertsCompiledSQL):
17    __dialect__ = default.DefaultDialect(supports_native_boolean=True)
18
19    run_setup_bind = None
20
21    run_create_tables = None
22
23    @classmethod
24    def define_tables(cls, metadata):
25        Table(
26            "people",
27            metadata,
28            Column("people_id", Integer, primary_key=True),
29            Column("age", Integer),
30            Column("name", String(30)),
31        )
32        Table(
33            "bookcases",
34            metadata,
35            Column("bookcase_id", Integer, primary_key=True),
36            Column(
37                "bookcase_owner_id", Integer, ForeignKey("people.people_id")
38            ),
39            Column("bookcase_shelves", Integer),
40            Column("bookcase_width", Integer),
41        )
42        Table(
43            "books",
44            metadata,
45            Column("book_id", Integer, primary_key=True),
46            Column(
47                "bookcase_id", Integer, ForeignKey("bookcases.bookcase_id")
48            ),
49            Column("book_owner_id", Integer, ForeignKey("people.people_id")),
50            Column("book_weight", Integer),
51        )
52
53    def test_standalone(self):
54        table1 = self.tables.people
55        subq = select([table1.c.people_id])
56
57        # alias name is not rendered because subquery is not
58        # in the context of a FROM clause
59        self.assert_compile(
60            lateral(subq, name="alias"),
61            "LATERAL (SELECT people.people_id FROM people)",
62        )
63
64        self.assert_compile(
65            subq.lateral(name="alias"),
66            "LATERAL (SELECT people.people_id FROM people)",
67        )
68
69    def test_select_from(self):
70        table1 = self.tables.people
71        subq = select([table1.c.people_id])
72
73        # in a FROM context, now you get "AS alias" and column labeling
74        self.assert_compile(
75            select([subq.lateral(name="alias")]),
76            "SELECT alias.people_id FROM LATERAL "
77            "(SELECT people.people_id AS people_id FROM people) AS alias",
78        )
79
80    def test_plain_join(self):
81        table1 = self.tables.people
82        table2 = self.tables.books
83        subq = select([table2.c.book_id]).where(
84            table2.c.book_owner_id == table1.c.people_id
85        )
86
87        # FROM books, people?  isn't this wrong?  No!  Because
88        # this is only a fragment, books isn't in any other FROM clause
89        self.assert_compile(
90            join(table1, lateral(subq, name="alias"), true()),
91            "people JOIN LATERAL (SELECT books.book_id AS book_id "
92            "FROM books, people WHERE books.book_owner_id = people.people_id) "
93            "AS alias ON true",
94        )
95
96        # put it in correct context, implicit correlation works fine
97        self.assert_compile(
98            select([table1]).select_from(
99                join(table1, lateral(subq, name="alias"), true())
100            ),
101            "SELECT people.people_id, people.age, people.name "
102            "FROM people JOIN LATERAL (SELECT books.book_id AS book_id "
103            "FROM books WHERE books.book_owner_id = people.people_id) "
104            "AS alias ON true",
105        )
106
107        # explicit correlation
108        subq = subq.correlate(table1)
109        self.assert_compile(
110            select([table1]).select_from(
111                join(table1, lateral(subq, name="alias"), true())
112            ),
113            "SELECT people.people_id, people.age, people.name "
114            "FROM people JOIN LATERAL (SELECT books.book_id AS book_id "
115            "FROM books WHERE books.book_owner_id = people.people_id) "
116            "AS alias ON true",
117        )
118
119    def test_join_lateral_w_select_subquery(self):
120        table1 = self.tables.people
121        table2 = self.tables.books
122
123        subq = (
124            select([table2.c.book_id])
125            .correlate(table1)
126            .where(table1.c.people_id == table2.c.book_owner_id)
127            .lateral()
128        )
129        stmt = select([table1, subq.c.book_id]).select_from(
130            table1.join(subq, true())
131        )
132
133        self.assert_compile(
134            stmt,
135            "SELECT people.people_id, people.age, people.name, anon_1.book_id "
136            "FROM people JOIN LATERAL (SELECT books.book_id AS book_id "
137            "FROM books "
138            "WHERE people.people_id = books.book_owner_id) AS anon_1 ON true",
139        )
140
141    def test_from_function(self):
142        bookcases = self.tables.bookcases
143        srf = lateral(func.generate_series(1, bookcases.c.bookcase_shelves))
144
145        self.assert_compile(
146            select([bookcases]).select_from(bookcases.join(srf, true())),
147            "SELECT bookcases.bookcase_id, bookcases.bookcase_owner_id, "
148            "bookcases.bookcase_shelves, bookcases.bookcase_width "
149            "FROM bookcases JOIN "
150            "LATERAL generate_series(:generate_series_1, "
151            "bookcases.bookcase_shelves) AS anon_1 ON true",
152        )
153