python-spanner-orm/spanner_orm/query.py at main · 7dracoder/python-spanner-orm · GitHub

History

213 lines (177 loc) · 7.34 KB

Raw

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

135

136

137

138

139

140

141

142

143

144

145

146

147

148

149

150

151

152

153

154

155

156

157

158

159

160

161

162

163

164

165

166

167

168

169

170

171

172

173

174

175

176

177

178

179

180

181

182

183

184

185

186

187

188

189

190

191

192

193

194

195

196

197

198

199

200

201

202

203

204

205

206

207

208

209

210

211

212

213

# Licensed under the Apache License, Version 2.0 (the "License");

# you may not use this file except in compliance with the License.

# You may obtain a copy of the License at

# https://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software

# distributed under the License is distributed on an "AS IS" BASIS,

# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

# See the License for the specific language governing permissions and

# limitations under the License.

"""Helps build SQL for complex Spanner queries."""

import abc

from typing import Any, Dict, Generic, Iterable, List, Sequence, Tuple, Type, TypeVar

from spanner_orm import condition

from spanner_orm import error

ResultType = TypeVar('ResultType')

class SpannerQuery(abc.ABC, Generic[ResultType]):

"""Helps build SQL for complex Spanner queries."""

def __init__(self, model: Type[Any],

conditions: Iterable[condition.Condition]):

self.param_offset = 0

self._model = model

self._conditions = conditions

self._sql = ''

self._parameters = {}

self._types = {}

self._build()

def _next_param_index(self) -> int:

return self.param_offset + len(self._parameters)

def parameters(self) -> Dict[str, Any]:

return self._parameters

def sql(self) -> str:

return self._sql

def types(self) -> Dict[str, Any]:

return self._types

@abc.abstractmethod

def process_results(self, results: List[Sequence[Any]]) -> ResultType:

pass

def _segments(self,

segment_type: condition.Segment) -> List[condition.Condition]:

segments = [

condition for condition in self._conditions

if condition.segment() == segment_type

]

for segment in segments:

segment.bind(self._model)

return segments

def _build(self) -> None:

"""Builds the Spanner query from the given model and conditions."""

segment_builders = [

self._select, self._from, self._where, self._order, self._limit

]

self._sql, self._parameters, self._types = '', {}, {}

for segment_builder in segment_builders:

segment_sql, segment_parameters, segment_types = segment_builder()

self._sql += segment_sql

self._parameters.update(segment_parameters)

self._types.update(segment_types)

@abc.abstractmethod

def _select(self) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:

"""Processes the SELECT segment of the SQL query."""

pass

def _from(self) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:

"""Processes the FROM segment of the SQL query."""

froms = self._segments(condition.Segment.FROM)

index_sql = ''

if froms:

if len(froms) != 1:

raise error.SpannerError('Only one index can be specified')

force_index = froms[0]

index_sql = force_index.sql()

sql = ' FROM {}{}'.format(self._model.table, index_sql)

return (sql, {}, {})

def _where(self) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:

"""Processes the WHERE segment of the SQL query."""

sql, sql_parts, parameters, types = '', [], {}, {}

wheres = self._segments(condition.Segment.WHERE)

for where in wheres:

where.suffix = str(self._next_param_index() + len(parameters))

sql_parts.append(where.sql())

parameters.update(where.params())

types.update(where.types())

if sql_parts:

sql = ' WHERE {}'.format(' AND '.join(sql_parts))

return (sql, parameters, types)

def _order(self) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:

"""Processes the ORDER BY segment of the SQL query."""

sql, parameters, types = '', {}, {}

orders = self._segments(condition.Segment.ORDER_BY)

if orders:

if len(orders) != 1:

raise error.SpannerError('Only one order condition may be specified')

order = orders[0]

order.suffix = str(self._next_param_index())

sql = ' ' + order.sql()

parameters = order.params()

types = order.types()

return (sql, parameters, types)

def _limit(self) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:

"""Processes the LIMIT segment of the SQL query."""

sql, parameters, types = '', {}, {}

limits = self._segments(condition.Segment.LIMIT)

if limits:

if len(limits) != 1:

raise error.SpannerError('Only one limit condition may be specified')

limit = limits[0]

limit.suffix = str(self._next_param_index())

sql = ' ' + limit.sql()

parameters = limit.params()

types = limit.types()

return (sql, parameters, types)

class CountQuery(SpannerQuery[int]):

"""Handles COUNT Spanner queries."""

def __init__(self, model: Type[Any],

conditions: Iterable[condition.Condition]):

super().__init__(model, conditions)

for c in conditions:

if c.segment() not in [condition.Segment.WHERE, condition.Segment.FROM]:

raise error.SpannerError('Only conditions that affect the WHERE or '

'FROM clauses are allowed for count queries')

def _select(self) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:

return ('SELECT COUNT(*)', {}, {})

def process_results(self, results: List[Sequence[Any]]) -> int:

return int(results[0][0])

class SelectQuery(SpannerQuery[List[Type[Any]]]):

"""Handles SELECT Spanner queries."""

def __init__(self, model: Type[Any],

conditions: Iterable[condition.Condition]):

self._model = model

self._conditions = conditions

self._joins = self._segments(condition.Segment.JOIN)

self._subqueries = [

_SelectSubQuery(join.destination, join.conditions)

for join in self._joins

if isinstance(join, condition.IncludesCondition)

]

super().__init__(model, conditions)

def _select_prefix(self) -> str:

return 'SELECT'

def _select(self) -> Tuple[str, Dict[str, Any], Dict[str, Any]]:

parameters, types = {}, {}

columns = [

'{alias}.{column}'.format(

alias=self._model.column_prefix, column=column)

for column in self._model.columns

]

for subquery in self._subqueries:

subquery.param_offset = self._next_param_index()

columns.append('ARRAY({subquery})'.format(subquery=subquery.sql()))

parameters.update(subquery.parameters())

types.update(subquery.types())

return ('{prefix} {columns}'.format(

prefix=self._select_prefix(),

columns=', '.join(columns)), parameters, types)

def process_results(self, results: List[Sequence[Any]]) -> List[Type[Any]]:

return [self._process_row(result) for result in results]

def _process_row(self, row: Sequence[Any]) -> Type[Any]:

"""Parses a row of results from a Spanner query based on the conditions."""

values = dict(zip(self._model.columns, row))

join_values = row[len(self._model.columns):]

for join, subquery, join_value in zip(self._joins, self._subqueries,

join_values):

models = subquery.process_results(join_value)

if join.single:

if len(models) > 1:

raise error.SpannerError(

'Multiple objects returned for relationship marked as single')

values[join.relation_name] = models[0] if models else None

else:

values[join.relation_name] = models

return self._model(values, persisted=True)

class _SelectSubQuery(SelectQuery):

def _select_prefix(self) -> str:

return 'SELECT AS STRUCT'

python-spanner-orm/spanner_orm/query.py at main · 7dracoder/python-spanner-orm · GitHub

Expand file tree

Breadcrumbs

query.py

Latest commit

History

Breadcrumbs

query.py

File metadata and controls