Skip to content

Commit 74b1d09

Browse files
committed
Support Interval Type
1 parent c01d58f commit 74b1d09

File tree

7 files changed

+287
-257
lines changed

7 files changed

+287
-257
lines changed

Pipfile.lock

Lines changed: 197 additions & 173 deletions
Some generated files are not rendered by default. Learn more about customizing how changed files appear on GitHub.

databend_sqlalchemy/connector.py

Lines changed: 2 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -48,6 +48,8 @@ def escape_item(self, item):
4848
return self.escape_number(item)
4949
elif isinstance(item, decimal.Decimal):
5050
return self.escape_number(item)
51+
elif isinstance(item, timedelta):
52+
return self.escape_string(f'{item.total_seconds()} seconds') + '::interval'
5153
elif isinstance(item, (datetime, date, time, timedelta)):
5254
return self.escape_string(item.strftime("%Y-%m-%d %H:%M:%S"))
5355
else:

databend_sqlalchemy/databend_dialect.py

Lines changed: 7 additions & 38 deletions
Original file line numberDiff line numberDiff line change
@@ -50,7 +50,9 @@
5050
)
5151
from sqlalchemy.engine import ExecutionContext, default
5252
from sqlalchemy.exc import DBAPIError, NoSuchTableError
53+
5354
from .dml import Merge
55+
from .types import INTERVAL
5456

5557
RESERVED_WORDS = {
5658
'Error', 'EOI', 'Whitespace', 'Comment', 'CommentBlock', 'Ident', 'ColumnPosition', 'LiteralString',
@@ -232,44 +234,8 @@ def process(value):
232234
return process
233235

234236

235-
class DatabendInterval(sqltypes.Interval):
236-
"""Stores interval as a datetime relative to epoch, see base implementation."""
237-
238-
_reg = re.compile(r"(\d+)-(\d+)-(\d+) (\d+):(\d+):(\d+)")
239-
240-
def result_processor(self, dialect, coltype):
241-
def process(value):
242-
if value is None:
243-
return None
244-
if isinstance(value, str):
245-
m = self._reg.match(value)
246-
if not m:
247-
raise ValueError(
248-
"could not parse %r as a datetime value" % (value,)
249-
)
250-
groups = m.groups()
251-
dt = datetime.datetime(*[
252-
int(groups[0] or self.epoch.year),
253-
int(groups[1] or self.epoch.month),
254-
int(groups[2] or self.epoch.day),
255-
int(groups[3] or 0),
256-
int(groups[4] or 0),
257-
int(groups[5] or 0),
258-
])
259-
else:
260-
dt = value
261-
return dt - self.epoch
262-
263-
return process
264-
265-
def literal_processor(self, dialect):
266-
def process(value):
267-
if value is not None:
268-
d = self.epoch + value
269-
interval_str = d.isoformat(" ", timespec="microseconds")
270-
return f"'{interval_str}'"
271-
272-
return process
237+
class DatabendInterval(INTERVAL):
238+
render_bind_cast = True
273239

274240

275241
# Type converters
@@ -526,6 +492,9 @@ def visit_JSON(self, type_, **kw):
526492
def visit_TIME(self, type_, **kw):
527493
return "DATETIME"
528494

495+
def visit_INTERVAL(self, type, **kw):
496+
return "INTERVAL"
497+
529498

530499
class DatabendDDLCompiler(compiler.DDLCompiler):
531500

databend_sqlalchemy/provision.py

Lines changed: 6 additions & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -1,7 +1,7 @@
11

22
from sqlalchemy.testing.provision import create_db
33
from sqlalchemy.testing.provision import drop_db
4-
from sqlalchemy.testing.provision import configure_follower
4+
from sqlalchemy.testing.provision import configure_follower, update_db_opts
55

66

77
@create_db.for_db("databend")
@@ -36,3 +36,8 @@ def _databend_drop_db(cfg, eng, ident):
3636
def _databend_configure_follower(config, ident):
3737
config.test_schema = "%s_test_schema" % ident
3838
config.test_schema_2 = "%s_test_schema_2" % ident
39+
40+
# Uncomment to debug SQL Statements in tests
41+
# @update_db_opts.for_db("databend")
42+
# def _mssql_update_db_opts(db_url, db_opts):
43+
# db_opts["echo"] = True

databend_sqlalchemy/types.py

Lines changed: 74 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -0,0 +1,74 @@
1+
2+
from __future__ import annotations
3+
4+
import datetime as dt
5+
from typing import Optional, Type, Any
6+
7+
from sqlalchemy.engine.interfaces import Dialect
8+
from sqlalchemy.sql import sqltypes
9+
from sqlalchemy.sql import type_api
10+
11+
12+
# ToDo - This is perhaps how numeric should be defined
13+
# class NUMERIC(sqltypes.Numeric):
14+
# def result_processor(self, dialect, type_):
15+
#
16+
# orig = super().result_processor(dialect, type_)
17+
#
18+
# def process(value):
19+
# if value is not None:
20+
# if self.decimal_return_scale:
21+
# value = decimal.Decimal(f'{value:.{self.decimal_return_scale}f}')
22+
# else:
23+
# value = decimal.Decimal(value)
24+
# if orig:
25+
# return orig(value)
26+
# return value
27+
#
28+
# return process
29+
30+
31+
class INTERVAL(type_api.NativeForEmulated, sqltypes._AbstractInterval):
32+
"""Databend INTERVAL type."""
33+
34+
__visit_name__ = "INTERVAL"
35+
native = True
36+
37+
def __init__(
38+
self, precision: Optional[int] = None, fields: Optional[str] = None
39+
) -> None:
40+
"""Construct an INTERVAL.
41+
42+
:param precision: optional integer precision value
43+
:param fields: string fields specifier. allows storage of fields
44+
to be limited, such as ``"YEAR"``, ``"MONTH"``, ``"DAY TO HOUR"``,
45+
etc.
46+
47+
"""
48+
self.precision = precision
49+
self.fields = fields
50+
51+
@classmethod
52+
def adapt_emulated_to_native(
53+
cls, interval: sqltypes.Interval, **kw: Any # type: ignore[override]
54+
) -> INTERVAL:
55+
return INTERVAL(precision=interval.second_precision)
56+
57+
@property
58+
def _type_affinity(self) -> Type[sqltypes.Interval]:
59+
return sqltypes.Interval
60+
61+
def as_generic(self, allow_nulltype: bool = False) -> sqltypes.Interval:
62+
return sqltypes.Interval(native=True, second_precision=self.precision)
63+
64+
@property
65+
def python_type(self) -> Type[dt.timedelta]:
66+
return dt.timedelta
67+
68+
def literal_processor(
69+
self, dialect: Dialect
70+
) -> Optional[type_api._LiteralProcessorType[dt.timedelta]]:
71+
def process(value: dt.timedelta) -> str:
72+
return f"to_interval('{value.total_seconds()} seconds')"
73+
74+
return process

setup.cfg

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -27,7 +27,7 @@ project_urls =
2727
[options]
2828
packages = find:
2929
install_requires =
30-
databend-driver>=0.16.0
30+
databend-driver>=0.25.0
3131
sqlalchemy>=1.4
3232
python_requires = >=3.7
3333
package_dir =

tests/test_sqlalchemy.py

Lines changed: 0 additions & 44 deletions
Original file line numberDiff line numberDiff line change
@@ -16,13 +16,10 @@
1616
from sqlalchemy.testing.suite import JoinTest as _JoinTest
1717
from sqlalchemy.testing.suite import BizarroCharacterFKResolutionTest as _BizarroCharacterFKResolutionTest
1818
from sqlalchemy.testing.suite import ServerSideCursorsTest as _ServerSideCursorsTest
19-
from sqlalchemy.testing.suite import IntervalTest as _IntervalTest
20-
from sqlalchemy.testing.suite import PrecisionIntervalTest as _PrecisionIntervalTest
2119
from sqlalchemy.testing.suite import EnumTest as _EnumTest
2220
from sqlalchemy import types as sql_types
2321
from sqlalchemy import testing, select
2422
from sqlalchemy.testing import config, eq_
25-
from databend_sqlalchemy.databend_dialect import DatabendInterval
2623

2724

2825
class ComponentReflectionTest(_ComponentReflectionTest):
@@ -283,47 +280,6 @@ def test_roundtrip_fetchmany(self):
283280
pass
284281

285282

286-
class IntervalTest(_IntervalTest):
287-
__backend__ = True
288-
datatype = DatabendInterval
289-
290-
@testing.skip("databend") # Skipped because cannot figure out the literal() part
291-
def test_arithmetic_operation_literal_interval(self, connection):
292-
pass
293-
294-
@testing.skip("databend") # Skipped because cannot figure out the literal() part
295-
def test_arithmetic_operation_table_interval_and_literal_interval(
296-
self, connection, arithmetic_table_fixture
297-
):
298-
pass
299-
300-
@testing.skip("databend") # Skipped because cannot figure out the literal() part
301-
def test_arithmetic_operation_table_date_and_literal_interval(
302-
self, connection, arithmetic_table_fixture
303-
):
304-
pass
305-
306-
307-
class PrecisionIntervalTest(_PrecisionIntervalTest):
308-
__backend__ = True
309-
datatype = DatabendInterval
310-
311-
@testing.skip("databend") # Skipped because cannot figure out the literal() part
312-
def test_arithmetic_operation_literal_interval(self, connection):
313-
pass
314-
315-
@testing.skip("databend") # Skipped because cannot figure out the literal() part
316-
def test_arithmetic_operation_table_interval_and_literal_interval(
317-
self, connection, arithmetic_table_fixture
318-
):
319-
pass
320-
321-
@testing.skip("databend") # Skipped because cannot figure out the literal() part
322-
def test_arithmetic_operation_table_date_and_literal_interval(
323-
self, connection, arithmetic_table_fixture
324-
):
325-
pass
326-
327283
class EnumTest(_EnumTest):
328284
__backend__ = True
329285

0 commit comments

Comments
 (0)