-
Notifications
You must be signed in to change notification settings - Fork 0
Expand file tree
/
Copy pathsqlcontext.py
More file actions
143 lines (112 loc) · 3.65 KB
/
Copy pathsqlcontext.py
File metadata and controls
143 lines (112 loc) · 3.65 KB
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
135
136
137
138
139
140
141
142
143
"""
Constructs that provide the context to evaluate a query.
"""
from dataclasses import dataclass
from enum import Enum
from typing import Any, Iterator, Optional
from .common import S, Symbol
from .prettier import (
Doc,
QuoteContext,
call_expr,
assg_expr,
list_expr,
resolve,
to_doc,
)
__all__ = [
"SQLTable",
"SQLCatalog",
"ValuesTable",
"VarStyle",
"LimitStyle",
"SQLDialect",
]
class SQLTable:
name: Symbol
columns: list[Symbol]
schema: Optional[Symbol]
def __init__(
self, name: Symbol, columns: list[Symbol], schema: Optional[Symbol] = None
) -> None:
self.name = S(name)
self.columns = [S(col) for col in columns]
self.schema = None if schema is None else S(schema)
def __repr__(self) -> str:
ctx = QuoteContext()
return resolve(to_doc(self, ctx), 80)
def pretty_repr(self, ctx: QuoteContext) -> Doc:
name = "SQLTable"
args = []
args.append(str(self.name))
if self.schema is not None:
args.append(assg_expr("schema", str(self.schema)))
if ctx.limit:
args.append("...")
else:
args.append(
assg_expr("columns", list_expr([str(col) for col in self.columns]))
)
return call_expr(name, args)
class ValuesTable:
"""To represent queries of the type - `FROM (VALUES ...) AS t(col1, col2, ...)`"""
columns: tuple[Symbol]
data: list[tuple]
def __init__(self, columns: tuple[str], data: list[tuple]) -> None:
for row in data:
assert len(row) == len(columns)
self.columns = tuple(S(c) for c in columns)
self.data = data
class SQLCatalog:
"""SQL catalog to capture the structure of a table like object for constructing queries"""
tables: dict[Symbol, SQLTable]
dialect: "SQLDialect"
def __init__(
self, dialect: "SQLDialect", tables: Optional[dict[Symbol, SQLTable]] = None
) -> None:
self.dialect = dialect
self.tables = {} if tables is None else tables
def __repr__(self) -> str:
ctx = QuoteContext()
return resolve(to_doc(self, ctx), 80)
def pretty_repr(self, ctx: QuoteContext) -> Doc:
name = "SQLCatalog"
args = []
args.append(assg_expr("dialect", str(self.dialect)))
for t_name, table in self.tables.items():
args.append(assg_expr(str(t_name), to_doc(table, ctx)))
return call_expr(name, args)
def get(self, key: Symbol) -> Optional[SQLTable]:
"""get the table with the given name"""
return self.tables.get(key, None)
def __getitem__(self, key: Symbol) -> SQLTable:
"""get the table with the given name"""
return self.tables[key]
def __len__(self) -> int:
return len(self.tables)
def __iter__(self) -> Iterator[tuple[Symbol, SQLTable]]:
for name, table in self.tables.items():
yield name, table
class VarStyle(Enum):
NAMED = 1
NUMBERED = 2
POSITIONAL = 3
class LimitStyle(Enum):
REGULAR = 1
FETCH_FIRST_KIND = 2
@dataclass(repr=False)
class SQLDialect:
name: str = "default"
var_style: VarStyle = VarStyle.NAMED
var_prefix: str = "?"
id_quotes: tuple[str, str] = ('"', '"')
has_bool_literals: bool = True
limit_style: LimitStyle = LimitStyle.REGULAR
has_recursive_annotation: bool = True
has_as_columns: bool = True
has_datetime_types: bool = True
values_row_constructor: Optional[str] = None
values_column_prefix: Optional[str] = "column"
values_column_index: int = 1
def __repr__(self) -> str:
return f"SQLDialect(:{self.name})"