test_util.py

import decimal
import datetime
from datetime import timezone, timedelta
import pytest
import pytz
from databricks.sql.utils import (
    convert_to_assigned_datatypes_in_column_table,
    parse_timestamp,
    ColumnTable,
    concat_table_chunks,
    serialize_query_tags,
)

try:
    import pyarrow
except ImportError:
    pyarrow = None


class TestUtils:
    def get_column_table_and_description(self):
        table_description = [
            ("id", "int", None, None, None, None, None),
            ("varchar_column", "string", None, None, None, None, None),
            ("boolean_column", "boolean", None, None, None, None, None),
            ("integer_column", "int", None, None, None, None, None),
            ("bigint_column", "bigint", None, None, None, None, None),
            ("smallint_column", "smallint", None, None, None, None, None),
            ("tinyint_column", "tinyint", None, None, None, None, None),
            ("float_column", "float", None, None, None, None, None),
            ("double_column", "double", None, None, None, None, None),
            ("decimal_column", "decimal", None, None, 10, 2, None),
            ("date_column", "date", None, None, None, None, None),
            ("timestamp_column", "timestamp", None, None, None, None, None),
            ("timestamp_ntz_column", "timestamp", None, None, None, None, None),
            ("timestamp_column_2", "timestamp", None, None, None, None, None),
            ("timestamp_column_3", "timestamp", None, None, None, None, None),
            ("timestamp_column_4", "timestamp", None, None, None, None, None),
            ("timestamp_column_5", "timestamp", None, None, None, None, None),
            ("timestamp_column_6", "timestamp", None, None, None, None, None),
            ("timestamp_column_7", "timestamp", None, None, None, None, None),
            ("binary_column", "binary", None, None, None, None, None),
            ("array_column", "array", None, None, None, None, None),
            ("map_column", "map", None, None, None, None, None),
            ("struct_column", "struct", None, None, None, None, None),
            ("variant_column", "string", None, None, None, None, None),
        ]

        column_table = [
            (9,),
            ("Test Varchar",),
            (True,),
            (123,),
            (9876543210,),
            (32000,),
            (120,),
            (1.23,),
            (4.56,),
            ("7890.12",),
            ("2023-12-31",),
            ("2023-12-31 12:30:00",),
            ("2023-12-31 12:30:00",),
            ("2021-09-30 11:27:35.123",),
            ("03/08/2024 02:30:15 PM",),
            ("08-Mar-2024 14:30:15",),
            ("2024-03-16T14:30:25.123",),
            ("2025-03-16T12:30:45+0530",),
            ("2025-03-16 12:30:45 +0530",),
            (b"\xde\xad\xbe\xef",),
            ('["item1","item2"]',),
            ('{"key1":"value1","key2":"value2"}',),
            ('{"name":"John","age":30}',),
            ('"semi-structured data"',),
        ]

        return column_table, table_description

    def test_convert_to_assigned_datatypes_in_column_table(self):
        column_table, description = self.get_column_table_and_description()
        converted_column_table = convert_to_assigned_datatypes_in_column_table(
            column_table, description
        )

        # (data , datatype)
        expected_convertion = [
            (9, int),
            ("Test Varchar", str),
            (True, bool),
            (123, int),
            (9876543210, int),
            (32000, int),
            (120, int),
            (1.23, float),
            (4.56, float),
            (decimal.Decimal("7890.12"), decimal.Decimal),
            (datetime.date(2023, 12, 31), datetime.date),
            (datetime.datetime(2023, 12, 31, 12, 30, 0), datetime.datetime),
            (datetime.datetime(2023, 12, 31, 12, 30, 0), datetime.datetime),
            (datetime.datetime(2021, 9, 30, 11, 27, 35, 123000), datetime.datetime),
            (datetime.datetime(2024, 3, 8, 14, 30, 15), datetime.datetime),
            (datetime.datetime(2024, 3, 8, 14, 30, 15), datetime.datetime),
            (datetime.datetime(2024, 3, 16, 14, 30, 25, 123000), datetime.datetime),
            (
                datetime.datetime(
                    2025,
                    3,
                    16,
                    12,
                    30,
                    45,
                    tzinfo=timezone(timedelta(hours=5, minutes=30)),
                ),
                datetime.datetime,
            ),
            (
                datetime.datetime(
                    2025,
                    3,
                    16,
                    12,
                    30,
                    45,
                    tzinfo=timezone(timedelta(hours=5, minutes=30)),
                ),
                datetime.datetime,
            ),
            (b"\xde\xad\xbe\xef", bytes),
            ('["item1","item2"]', str),
            ('{"key1":"value1","key2":"value2"}', str),
            ('{"name":"John","age":30}', str),
            ('"semi-structured data"', str),
        ]

        for index, entry in enumerate(converted_column_table):
            assert entry[0] == expected_convertion[index][0]
            assert isinstance(entry[0], expected_convertion[index][1])

    def test_concat_table_chunks_column_table(self):
        column_table1 = ColumnTable([[1, 2], [5, 6]], ["col1", "col2"])
        column_table2 = ColumnTable([[3, 4], [7, 8]], ["col1", "col2"])

        result_table = concat_table_chunks([column_table1, column_table2])

        assert result_table.column_table == [[1, 2, 3, 4], [5, 6, 7, 8]]
        assert result_table.column_names == ["col1", "col2"]

    @pytest.mark.skipif(pyarrow is None, reason="PyArrow is not installed")
    def test_concat_table_chunks_arrow_table(self):
        arrow_table1 = pyarrow.Table.from_pydict({"col1": [1, 2], "col2": [5, 6]})
        arrow_table2 = pyarrow.Table.from_pydict({"col1": [3, 4], "col2": [7, 8]})

        result_table = concat_table_chunks([arrow_table1, arrow_table2])
        assert result_table.column_names == ["col1", "col2"]
        assert result_table.column("col1").to_pylist() == [1, 2, 3, 4]
        assert result_table.column("col2").to_pylist() == [5, 6, 7, 8]

    def test_concat_table_chunks_empty(self):
        result_table = concat_table_chunks([])
        assert result_table == []

    def test_concat_table_chunks__incorrect_column_names_error(self):
        column_table1 = ColumnTable([[1, 2], [5, 6]], ["col1", "col2"])
        column_table2 = ColumnTable([[3, 4], [7, 8]], ["col1", "col3"])

        with pytest.raises(ValueError):
            concat_table_chunks([column_table1, column_table2])

    def test_serialize_query_tags_basic(self):
        """Test basic query tags serialization"""
        query_tags = {"team": "data-eng", "application": "etl"}
        result = serialize_query_tags(query_tags)
        assert result == "team:data-eng,application:etl"

    def test_serialize_query_tags_with_none_value(self):
        """Test query tags with None value (should omit colon and value)"""
        query_tags = {"key1": "value1", "key2": None, "key3": "value3"}
        result = serialize_query_tags(query_tags)
        assert result == "key1:value1,key2,key3:value3"

    def test_serialize_query_tags_with_special_chars(self):
        """Test query tags with special characters (colon, comma, backslash)"""
        query_tags = {
            "key1": "value:with:colons",
            "key2": "value,with,commas",
            "key3": r"value\with\backslashes",
        }
        result = serialize_query_tags(query_tags)
        assert (
            result
            == r"key1:value\:with\:colons,key2:value\,with\,commas,key3:value\\with\\backslashes"
        )

    def test_serialize_query_tags_with_mixed_special_chars(self):
        """Test query tags with mixed special characters"""
        query_tags = {"key1": r"a:b,c\d"}
        result = serialize_query_tags(query_tags)
        assert result == r"key1:a\:b\,c\\d"

    def test_serialize_query_tags_empty_dict(self):
        """Test serialization with empty dictionary"""
        query_tags = {}
        result = serialize_query_tags(query_tags)
        assert result is None

    def test_serialize_query_tags_none(self):
        """Test serialization with None input"""
        result = serialize_query_tags(None)
        assert result is None

    def test_serialize_query_tags_with_special_chars_in_key(self):
        """Test query tags with special characters in keys (only backslashes are escaped in keys)"""
        query_tags = {
            "key:with:colons": "value1",
            "key,with,commas": "value2",
            r"key\with\backslashes": "value3",
        }
        result = serialize_query_tags(query_tags)
        # Only backslashes are escaped in keys; colons and commas in keys are not escaped
        assert (
            result
            == r"key:with:colons:value1,key,with,commas:value2,key\\with\\backslashes:value3"
        )

    def test_serialize_query_tags_all_none_values(self):
        """Test query tags where all values are None"""
        query_tags = {"key1": None, "key2": None, "key3": None}
        result = serialize_query_tags(query_tags)
        assert result == "key1,key2,key3"


class TestParseTimestamp:
    def test_no_format_uses_dateutil(self):
        result = parse_timestamp("2023-12-31 12:30:00")
        assert result == datetime.datetime(2023, 12, 31, 12, 30, 0)

    def test_matching_format_uses_strptime(self):
        fmt = "%Y-%m-%d %H:%M:%S.%f"
        result = parse_timestamp("2023-12-31 12:30:00.123000", fmt)
        assert result == datetime.datetime(2023, 12, 31, 12, 30, 0, 123000, tzinfo=pytz.UTC)

    def test_non_matching_format_falls_back_to_dateutil(self):
        fmt = "%Y-%m-%d %H:%M:%S.%f"
        # This doesn't match the format, so should fall back to dateutil
        result = parse_timestamp("08-Mar-2024 14:30:15", fmt)
        assert result == datetime.datetime(2024, 3, 8, 14, 30, 15)

    def test_convert_column_table_with_timestamp_format(self):
        description = [
            ("ts_col", "timestamp", None, None, None, None, None),
        ]
        column_table = [("2023-12-31 12:30:00.000000",)]
        fmt = "%Y-%m-%d %H:%M:%S.%f"
        result = convert_to_assigned_datatypes_in_column_table(
            column_table, description, timestamp_format=fmt
        )
        assert result[0][0] == datetime.datetime(2023, 12, 31, 12, 30, 0, tzinfo=pytz.UTC)

    def test_convert_column_table_without_timestamp_format(self):
        description = [
            ("ts_col", "timestamp", None, None, None, None, None),
        ]
        column_table = [("2023-12-31 12:30:00",)]
        result = convert_to_assigned_datatypes_in_column_table(
            column_table, description
        )
        assert result[0][0] == datetime.datetime(2023, 12, 31, 12, 30, 0)