import datetime

from .predicates import compare_dbr_versions, is_thrift_v5_plus, pysql_has_version

class TimestampTestsMixin:

timestamp_and_expected_results = [

('2021-09-30 11:27:35.123+04:00', datetime.datetime(2021, 9, 30, 7, 27, 35, 123000)),

('2021-09-30 11:27:35+04:00', datetime.datetime(2021, 9, 30, 7, 27, 35)),

('2021-09-30 11:27:35.123', datetime.datetime(2021, 9, 30, 11, 27, 35, 123000)),

('2021-09-30 11:27:35', datetime.datetime(2021, 9, 30, 11, 27, 35)),

('2021-09-30 11:27', datetime.datetime(2021, 9, 30, 11, 27)),

('2021-09-30 11', datetime.datetime(2021, 9, 30, 11)),

('2021-09-30', datetime.datetime(2021, 9, 30)),

('2021-09', datetime.datetime(2021, 9, 1)),

('2021', datetime.datetime(2021, 1, 1)),

('9999-12-31T15:59:59', datetime.datetime(9999, 12, 31, 15, 59, 59)),

('9999-99-31T15:59:59', None),

]

date_and_expected_results = [

('2021-09-30', datetime.date(2021, 9, 30)),

('2021-09', datetime.date(2021, 9, 1)),

('2021', datetime.date(2021, 1, 1)),

('9999-12-31', datetime.date(9999, 12, 31)),

('9999-99-31', None),

]

def should_add_timezone(self):

return pysql_has_version(">=", 2) and is_thrift_v5_plus(self.arguments)

def maybe_add_timezone_to_timestamp(self, ts):

"""If we're using DBR >= 10.2, then we expect back aware timestamps, so add timezone to `ts`

Otherwise we have naive timestamps, so no change is needed

"""

if ts and self.should_add_timezone():

return ts.replace(tzinfo=datetime.timezone.utc)

else:

return ts

def assertTimestampsEqual(self, result, expected):

self.assertEqual(result, self.maybe_add_timezone_to_timestamp(expected))

def multi_query(self, n_rows=10):

row_sql = "SELECT " + ", ".join(

["TIMESTAMP('{}')".format(ts) for (ts, _) in self.timestamp_and_expected_results])

query = " UNION ALL ".join([row_sql for _ in range(n_rows)])

expected_matrix = [[dt for (_, dt) in self.timestamp_and_expected_results]

for _ in range(n_rows)]

return query, expected_matrix

def test_timestamps(self):

with self.cursor({"session_configuration": {"ansi_mode": False}}) as cursor:

for (timestamp, expected) in self.timestamp_and_expected_results:

cursor.execute("SELECT TIMESTAMP('{timestamp}')".format(timestamp=timestamp))

result = cursor.fetchone()[0]

self.assertTimestampsEqual(result, expected)

def test_multi_timestamps(self):

with self.cursor({"session_configuration": {"ansi_mode": False}}) as cursor:

query, expected = self.multi_query()

cursor.execute(query)

result = cursor.fetchall()

# We list-ify the rows because PyHive will return a tuple for a row

self.assertEqual([list(r) for r in result],

[[self.maybe_add_timezone_to_timestamp(ts) for ts in r]

for r in expected])

def test_dates(self):

with self.cursor({"session_configuration": {"ansi_mode": False}}) as cursor:

for (date, expected) in self.date_and_expected_results:

cursor.execute("SELECT DATE('{date}')".format(date=date))

result = cursor.fetchone()[0]

self.assertEqual(result, expected)