|
3 | 3 | import time |
4 | 4 | import re |
5 | 5 | import math |
6 | | -import datetime |
| 6 | +from datetime import datetime, timedelta |
7 | 7 | from decimal import Decimal |
8 | 8 | from parameterized import parameterized |
9 | 9 |
|
@@ -52,33 +52,33 @@ def __next__(self) -> str: |
52 | 52 |
|
53 | 53 | class DateTimeFaker: |
54 | 54 | MANUAL_FAKES = [ |
55 | | - datetime.datetime.fromisoformat("2020-01-01 15:10:10"), |
56 | | - datetime.datetime.fromisoformat("2020-02-01 09:09:09"), |
57 | | - datetime.datetime.fromisoformat("2022-03-01 15:10:01.139"), |
58 | | - datetime.datetime.fromisoformat("2022-04-01 15:10:02.020409"), |
59 | | - datetime.datetime.fromisoformat("2022-05-01 15:10:03.003030"), |
60 | | - datetime.datetime.fromisoformat("2022-06-01 15:10:05.009900"), |
| 55 | + datetime.fromisoformat("2020-01-01 15:10:10"), |
| 56 | + datetime.fromisoformat("2020-02-01 09:09:09"), |
| 57 | + datetime.fromisoformat("2022-03-01 15:10:01.139"), |
| 58 | + datetime.fromisoformat("2022-04-01 15:10:02.020409"), |
| 59 | + datetime.fromisoformat("2022-05-01 15:10:03.003030"), |
| 60 | + datetime.fromisoformat("2022-06-01 15:10:05.009900"), |
61 | 61 | ] |
62 | 62 |
|
63 | 63 | def __init__(self, max): |
64 | 64 | self.max = max |
65 | 65 |
|
66 | 66 | def __iter__(self): |
67 | 67 | iter = DateTimeFaker(self.max) |
68 | | - iter.prev = datetime.datetime(2000, 1, 1, 0, 0, 0, 0) |
| 68 | + iter.prev = datetime(2000, 1, 1, 0, 0, 0, 0) |
69 | 69 | iter.i = 0 |
70 | 70 | return iter |
71 | 71 |
|
72 | 72 | def __len__(self): |
73 | 73 | return self.max |
74 | 74 |
|
75 | | - def __next__(self) -> datetime.datetime: |
| 75 | + def __next__(self) -> datetime: |
76 | 76 | if self.i < len(self.MANUAL_FAKES): |
77 | 77 | fake = self.MANUAL_FAKES[self.i] |
78 | 78 | self.i += 1 |
79 | 79 | return fake |
80 | 80 | elif self.i < self.max: |
81 | | - self.prev = self.prev + datetime.timedelta(seconds=3, microseconds=571) |
| 81 | + self.prev = self.prev + timedelta(seconds=3, microseconds=571) |
82 | 82 | self.i += 1 |
83 | 83 | return self.prev |
84 | 84 | else: |
@@ -373,7 +373,7 @@ def _insert_to_table(conn, table, values): |
373 | 373 | for j, sample in values: |
374 | 374 | if isinstance(sample, (float, Decimal, int)): |
375 | 375 | value = str(sample) |
376 | | - elif isinstance(sample, datetime.datetime) and isinstance(conn, db.Presto): |
| 376 | + elif isinstance(sample, datetime) and isinstance(conn, db.Presto): |
377 | 377 | value = f"timestamp '{sample}'" |
378 | 378 | else: |
379 | 379 | value = f"'{sample}'" |
@@ -422,6 +422,11 @@ def test_types(self, source_db, target_db, source_type, target_type, type_catego |
422 | 422 | _insert_to_table(src_conn, src_table, enumerate(sample_values, 1)) |
423 | 423 |
|
424 | 424 | values_in_source = PaginatedTable(src_table, src_conn) |
| 425 | + if source_db is db.Presto: |
| 426 | + if source_type.startswith("decimal"): |
| 427 | + values_in_source = [(a, Decimal(b)) for a, b in values_in_source] |
| 428 | + elif source_type.startswith("timestamp"): |
| 429 | + values_in_source = [(a, datetime.fromisoformat(b.rstrip(" UTC"))) for a, b in values_in_source] |
425 | 430 |
|
426 | 431 | _drop_table_if_exists(dst_conn, dst_table) |
427 | 432 | dst_conn.query(f"CREATE TABLE {dst_table}(id int, col {target_type})", None) |
|
0 commit comments