clickhouse-sqlalchemy/clickhouse_sqlalchemy/orm/query.py at 90800654ffd9438e87d36ba10a4d78613f3242c4 · xzkostyan/clickhouse-sqlalchemy · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
from sqlalchemy import exc
from sqlalchemy.sql.base import _generative
from sqlalchemy.orm.query import Query as BaseQuery

from ..ext.clauses import (
    ArrayJoin,
    LeftArrayJoin,
    LimitByClause,
    sample_clause,
)


class Query(BaseQuery):
    _with_cube = False
    _with_rollup = False
    _with_totals = False
    _final = None
    _sample = None
    _limit_by = None
    _array_join = None

    def _statement_20(self, for_statement=False, use_legacy_query_style=True):
        orig_smt = super(Query, self)._statement_20(for_statement=for_statement,
                                                    use_legacy_query_style=use_legacy_query_style)

        orig_smt._with_cube = self._with_cube
        orig_smt._with_rollup = self._with_rollup
        orig_smt._with_totals = self._with_totals
        orig_smt._final_clause = self._final
        orig_smt._sample_clause = sample_clause(self._sample)
        orig_smt._limit_by_clause = self._limit_by
        orig_smt._array_join = self._array_join

        return orig_smt


    def _compile_context(self, *args, **kwargs):
        context = super(Query, self)._compile_context(*args, **kwargs)
        query = context.query

        query._with_cube = self._with_cube
        query._with_rollup = self._with_rollup
        query._with_totals = self._with_totals
        query._final_clause = self._final
        query._sample_clause = sample_clause(self._sample)
        query._limit_by_clause = self._limit_by
        query._array_join = self._array_join

        return context

    @_generative
    def with_cube(self):
        if not self._group_by_clauses:
            raise exc.InvalidRequestError(
                "Query.with_cube() can be used only with specified "
                "GROUP BY, call group_by()"
            )
        if self._with_rollup:
            raise exc.InvalidRequestError(
                "Query.with_cube() and Query.with_rollup() are mutually "
                "exclusive"
            )

        self._with_cube = True

    @_generative
    def with_rollup(self):
        if not self._group_by_clauses:
            raise exc.InvalidRequestError(
                "Query.with_rollup() can be used only with specified "
                "GROUP BY, call group_by()"
            )
        if self._with_cube:
            raise exc.InvalidRequestError(
                "Query.with_cube() and Query.with_rollup() are mutually "
                "exclusive"
            )

        self._with_rollup = True

    @_generative
    def with_totals(self):
        if not self._group_by_clauses:
            raise exc.InvalidRequestError(
                "Query.with_totals() can be used only with specified "
                "GROUP BY, call group_by()"
            )

        self._with_totals = True

    def _add_array_join(self, columns, left):
        join_type = ArrayJoin if not left else LeftArrayJoin
        self._array_join = join_type(*columns)

    @_generative
    def array_join(self, *columns, **kwargs):
        left = kwargs.get("left", False)
        self._add_array_join(columns, left=left)

    @_generative
    def left_array_join(self, *columns):
        self._add_array_join(columns, left=True)

    @_generative
    def final(self):
        self._final = True

    @_generative
    def sample(self, sample):
        self._sample = sample

    @_generative
    def limit_by(self, by_clauses, limit, offset=None):
        self._limit_by = LimitByClause(by_clauses, limit, offset)

    def join(self, *props, **kwargs):
        spec = {
            'type': kwargs.pop('type', None),
            'strictness': kwargs.pop('strictness', None),
            'distribution': kwargs.pop('distribution', None)
        }
        rv = super(Query, self).join(*props, **kwargs)

        x = rv._legacy_setup_joins[-1]
        x_spec = dict(spec)
        # use 'full' key to pass extra flags
        x_spec['full'] = x[-1]['full']
        x[-1]['full'] = x_spec

        return rv

    def outerjoin(self, *props, **kwargs):
        kwargs['type'] = kwargs.get('type') or 'LEFT OUTER'
        return self.join(*props, **kwargs)