# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

#
from __future__ import annotations

from collections import namedtuple
from datetime import timedelta
from unittest import mock
from unittest.mock import PropertyMock, patch
from urllib.parse import quote_plus

import pandas as pd
import polars as pl
import pytest
from databricks.sql.types import Row

from airflow.models import Connection
from airflow.providers.common.compat.sdk import AirflowException, AirflowOptionalProviderFeatureException
from airflow.providers.common.sql.hooks.handlers import fetch_all_handler
from airflow.providers.databricks.hooks.databricks_sql import (
    DatabricksSqlHook,
    _format_query_tag_value,
    _format_query_tags,
    create_timeout_thread,
)

TASK_ID = "databricks-sql-operator"
DEFAULT_CONN_ID = "databricks_default"
HOST = "xx.cloud.databricks.com"
HOST_WITH_SCHEME = "https://xx.cloud.databricks.com"
TOKEN = "token"
HTTP_PATH = "sql/protocolv1/o/1234567890123456/0123-456789-abcd123"
SCHEMA = "test_schema"
CATALOG = "test_catalog"


@pytest.fixture(autouse=True)
def create_connection(create_connection_without_db):
    create_connection_without_db(
        Connection(
            conn_id=DEFAULT_CONN_ID,
            conn_type="databricks",
            host=HOST,
            login=None,
            password=TOKEN,
            extra=None,
        )
    )


@pytest.fixture
def databricks_hook():
    return DatabricksSqlHook(sql_endpoint_name="Test")


@pytest.fixture
def mock_get_conn():
    # Start the patcher
    mock_patch = patch("airflow.providers.databricks.hooks.databricks_sql.DatabricksSqlHook.get_conn")
    mock_conn = mock_patch.start()
    # Use yield to provide the mock object
    yield mock_conn
    # Stop the patcher
    mock_patch.stop()


@pytest.fixture
def mock_get_requests():
    # Start the patcher
    mock_patch = patch("airflow.providers.databricks.hooks.databricks_base.requests")
    mock_requests = mock_patch.start()

    # Configure the mock object with the current API response format ("warehouses" key)
    mock_requests.codes.ok = 200
    mock_requests.get.return_value.json.return_value = {
        "warehouses": [
            {
                "id": "1264e5078741679a",
                "name": "Test",
                "odbc_params": {
                    "hostname": "xx.cloud.databricks.com",
                    "path": "/sql/1.0/endpoints/1264e5078741679a",
                },
            }
        ]
    }
    status_code_mock = PropertyMock(return_value=200)
    type(mock_requests.get.return_value).status_code = status_code_mock

    # Yield the mock object
    yield mock_requests

    # Stop the patcher after the test
    mock_patch.stop()


@pytest.fixture
def mock_timer():
    with patch("threading.Timer") as mock_timer:
        yield mock_timer


def test_sqlachemy_url_property():
    hook = DatabricksSqlHook(
        databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA
    )
    url = hook.sqlalchemy_url.render_as_string(hide_password=False)
    expected_url = (
        f"databricks://token:{TOKEN}@{HOST}?"
        f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}"
    )
    assert url == expected_url


def test_get_sqlalchemy_engine():
    hook = DatabricksSqlHook(
        databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA
    )
    engine = hook.get_sqlalchemy_engine()
    assert engine.url.render_as_string(hide_password=False) == (
        f"databricks://token:{TOKEN}@{HOST}?"
        f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}"
    )


def test_get_uri():
    hook = DatabricksSqlHook(
        databricks_conn_id=DEFAULT_CONN_ID, http_path=HTTP_PATH, catalog=CATALOG, schema=SCHEMA
    )
    uri = hook.get_uri()
    expected_uri = (
        f"databricks://token:{TOKEN}@{HOST}?"
        f"catalog={CATALOG}&http_path={quote_plus(HTTP_PATH)}&schema={SCHEMA}"
    )
    assert uri == expected_uri


def get_cursor_descriptions(fields: list[str]) -> list[tuple[str]]:
    return [(field,) for field in fields]


# Serializable Row object similar to the one returned by the Hook
SerializableRow = namedtuple("Row", ["id", "value"])  # type: ignore[name-match]


@pytest.mark.parametrize(
    (
        "return_last",
        "split_statements",
        "sql",
        "execution_timeout",
        "cursor_calls",
        "cursor_descriptions",
        "cursor_results",
        "hook_descriptions",
        "hook_results",
    ),
    [
        pytest.param(
            True,
            False,
            "select * from test.test",
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row(id=1, value=2), Row(id=11, value=12)],),
            [[("id",), ("value",)]],
            [Row(id=1, value=2), Row(id=11, value=12)],
            id="The return_last set and no split statements set on single query in string",
        ),
        pytest.param(
            False,
            False,
            "select * from test.test;",
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row(id=1, value=2), Row(id=11, value=12)],),
            [[("id",), ("value",)]],
            [Row(id=1, value=2), Row(id=11, value=12)],
            id="The return_last not set and no split statements set on single query in string",
        ),
        pytest.param(
            True,
            True,
            "select * from test.test;",
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row(id=1, value=2), Row(id=11, value=12)],),
            [[("id",), ("value",)]],
            [Row(id=1, value=2), Row(id=11, value=12)],
            id="The return_last set and split statements set on single query in string",
        ),
        pytest.param(
            False,
            True,
            "select * from test.test;",
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row(id=1, value=2), Row(id=11, value=12)],),
            [[("id",), ("value",)]],
            [[Row(id=1, value=2), Row(id=11, value=12)]],
            id="The return_last not set and split statements set on single query in string",
        ),
        pytest.param(
            True,
            True,
            "select * from test.test;select * from test.test2;",
            None,
            ["select * from test.test", "select * from test.test2"],
            [["id", "value"], ["id2", "value2"]],
            [[Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]],
            [[("id2",), ("value2",)]],
            [Row(id=3, value=4), Row(id=13, value=14)],
            id="The return_last set and split statements set on multiple queries in string",
        ),
        pytest.param(
            False,
            True,
            "select * from test.test;select * from test.test2;",
            None,
            ["select * from test.test", "select * from test.test2"],
            [["id", "value"], ["id2", "value2"]],
            [[Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]],
            [[("id",), ("value",)], [("id2",), ("value2",)]],
            [
                [Row(id=1, value=2), Row(id=11, value=12)],
                [Row(id=3, value=4), Row(id=13, value=14)],
            ],
            id="The return_last not set and split statements set on multiple queries in string",
        ),
        pytest.param(
            True,
            True,
            ["select * from test.test;"],
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row(id=1, value=2), Row(id=11, value=12)],),
            [[("id",), ("value",)]],
            [[Row(id=1, value=2), Row(id=11, value=12)]],
            id="The return_last set on single query in list",
        ),
        pytest.param(
            False,
            True,
            ["select * from test.test;"],
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row(id=1, value=2), Row(id=11, value=12)],),
            [[("id",), ("value",)]],
            [[Row(id=1, value=2), Row(id=11, value=12)]],
            id="The return_last not set on single query in list",
        ),
        pytest.param(
            True,
            True,
            "select * from test.test;select * from test.test2;",
            None,
            ["select * from test.test", "select * from test.test2"],
            [["id", "value"], ["id2", "value2"]],
            [[Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]],
            [[("id2",), ("value2",)]],
            [Row(id=3, value=4), Row(id=13, value=14)],
            id="The return_last set on multiple queries in list",
        ),
        pytest.param(
            False,
            True,
            "select * from test.test;select * from test.test2;",
            None,
            ["select * from test.test", "select * from test.test2"],
            [["id", "value"], ["id2", "value2"]],
            [[Row(id=1, value=2), Row(id=11, value=12)], [Row(id=3, value=4), Row(id=13, value=14)]],
            [[("id",), ("value",)], [("id2",), ("value2",)]],
            [
                [Row(id=1, value=2), Row(id=11, value=12)],
                [Row(id=3, value=4), Row(id=13, value=14)],
            ],
            id="The return_last not set on multiple queries not set",
        ),
        pytest.param(
            True,
            False,
            "select * from test.test",
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row("id", "value")(1, 2)],),
            [[("id",), ("value",)]],
            [SerializableRow(1, 2)],
            id="Return a serializable row (tuple) from a row instance created in two step",
        ),
        pytest.param(
            True,
            False,
            "select * from test.test",
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([Row(id=1, value=2)],),
            [[("id",), ("value",)]],
            [SerializableRow(1, 2)],
            id="Return a serializable row (tuple) from a row instance created in one step",
        ),
        pytest.param(
            True,
            False,
            "select * from test.test",
            None,
            ["select * from test.test"],
            [["id", "value"]],
            ([],),
            [[("id",), ("value",)]],
            [],
            id="Empty list",
        ),
    ],
)
def test_query(
    return_last,
    split_statements,
    sql,
    execution_timeout,
    cursor_calls,
    cursor_descriptions,
    cursor_results,
    hook_descriptions,
    hook_results,
    mock_get_conn,
    mock_get_requests,
):
    connections = []
    cursors = []

    for index, cursor_description in enumerate(cursor_descriptions):
        conn = mock.MagicMock()
        cur = mock.MagicMock(
            rowcount=len(cursor_results[index]),
            description=get_cursor_descriptions(cursor_description),
        )
        cur.fetchall.return_value = cursor_results[index]
        conn.cursor.return_value = cur
        cursors.append(cur)
        connections.append(conn)
    mock_get_conn.side_effect = connections

    databricks_hook = DatabricksSqlHook(sql_endpoint_name="Test")
    results = databricks_hook.run(
        sql=sql, handler=fetch_all_handler, return_last=return_last, split_statements=split_statements
    )

    assert databricks_hook.descriptions == hook_descriptions
    assert databricks_hook.last_description == hook_descriptions[-1]
    assert results == hook_results

    for index, cur in enumerate(cursors):
        cur.execute.assert_has_calls([mock.call(cursor_calls[index])])
    cur.close.assert_called()


@pytest.mark.parametrize(
    "empty_statement",
    [
        pytest.param([], id="Empty list"),
        pytest.param("", id="Empty string"),
        pytest.param("\n", id="Only EOL"),
    ],
)
def test_no_query(databricks_hook, empty_statement):
    with pytest.raises(ValueError, match="List of SQL statements is empty"):
        databricks_hook.run(sql=empty_statement)


@mock.patch("airflow.providers.common.sql.hooks.sql.send_sql_hook_lineage")
def test_run_hook_lineage(mock_send_lineage, mock_get_conn, mock_get_requests):
    """Ensure run() triggers send_sql_hook_lineage via base DbApiHook._run_command."""
    conn = mock.MagicMock()
    cur = mock.MagicMock(
        rowcount=1,
        description=[("id",), ("value",)],
    )
    cur.fetchall.return_value = [Row(id=1, value=2)]
    conn.cursor.return_value = cur
    mock_get_conn.return_value = conn

    sql = "SELECT 1"
    hook = DatabricksSqlHook(sql_endpoint_name="Test")
    hook.run(sql=sql, handler=fetch_all_handler)

    mock_send_lineage.assert_called()
    call_kw = mock_send_lineage.call_args.kwargs
    assert call_kw["context"] is hook
    assert call_kw["sql"] == sql
    assert call_kw["sql_parameters"] is None
    assert call_kw["cur"] is cur


@mock.patch("airflow.providers.common.sql.hooks.sql.send_sql_hook_lineage")
def test_insert_rows_hook_lineage(mock_send_lineage, mock_get_conn):
    conn = mock.MagicMock()
    cur = mock.MagicMock(rowcount=0)
    conn.cursor.return_value = cur
    mock_get_conn.return_value = conn

    table = "table"
    rows = [("hello",), ("world",)]
    hook = DatabricksSqlHook(sql_endpoint_name="Test")
    hook.insert_rows(table, rows)

    mock_send_lineage.assert_called()
    call_kw = mock_send_lineage.call_args.kwargs
    assert call_kw["context"] is hook
    assert call_kw["sql"] == "INSERT INTO table  VALUES (%s)"
    assert call_kw["row_count"] == 2


@mock.patch("airflow.providers.common.sql.hooks.sql.send_sql_hook_lineage")
@mock.patch("airflow.providers.common.sql.hooks.sql.DbApiHook._get_pandas_df")
def test_get_df_hook_lineage(mock_get_pandas_df, mock_send_lineage, mock_get_conn):
    sql = "SELECT 1"
    parameters = ("x",)
    hook = DatabricksSqlHook(sql_endpoint_name="Test")
    hook.get_df(sql, parameters=parameters)

    mock_send_lineage.assert_called_once()
    call_kw = mock_send_lineage.call_args.kwargs
    assert call_kw["context"] is hook
    assert call_kw["sql"] == sql
    assert call_kw["sql_parameters"] == parameters


@mock.patch("airflow.providers.common.sql.hooks.sql.send_sql_hook_lineage")
@mock.patch("airflow.providers.common.sql.hooks.sql.DbApiHook._get_pandas_df_by_chunks")
def test_get_df_by_chunks_hook_lineage(mock_get_pandas_df_by_chunks, mock_send_lineage, mock_get_conn):
    sql = "SELECT 1"
    parameters = ("x",)
    hook = DatabricksSqlHook(sql_endpoint_name="Test")
    hook.get_df_by_chunks(sql, parameters=parameters, chunksize=1)

    mock_send_lineage.assert_called_once()
    call_kw = mock_send_lineage.call_args.kwargs
    assert call_kw["context"] is hook
    assert call_kw["sql"] == sql
    assert call_kw["sql_parameters"] == parameters


@pytest.mark.parametrize(
    ("row_objects", "fields_names"),
    [
        pytest.param(Row("count(1)")(9714), ("_0",)),
        pytest.param(Row("1//@:()")("data"), ("_0",)),
        pytest.param(Row("class")("data"), ("_0",)),
        pytest.param(Row("1_wrong", "2_wrong")(1, 2), ("_0", "_1")),
    ],
)
def test_incorrect_column_names(row_objects, fields_names):
    """Ensure that column names can be used as namedtuple attribute.

    namedtuple do not accept special characters and reserved python keywords
    as column name. This test ensure that such columns are renamed.
    """
    result = DatabricksSqlHook()._make_common_data_structure(row_objects)
    assert result._fields == fields_names


@pytest.mark.parametrize(
    ("sql", "execution_timeout", "cursor_descriptions", "cursor_results"),
    [
        (
            "select * from test.test",
            timedelta(microseconds=0),
            ("id", "value"),
            (Row(id=1, value=2), Row(id=11, value=12)),
        )
    ],
)
def test_execution_timeout_exceeded(
    mock_get_conn,
    mock_get_requests,
    sql,
    execution_timeout,
    cursor_descriptions,
    cursor_results,
):
    with (
        patch(
            "airflow.providers.databricks.hooks.databricks_sql.create_timeout_thread"
        ) as mock_create_timeout_thread,
        patch.object(DatabricksSqlHook, "_run_command") as mock_run_command,
    ):
        conn = mock.MagicMock()
        cur = mock.MagicMock(
            rowcount=len(cursor_results),
            description=get_cursor_descriptions(cursor_descriptions),
        )

        mock_event = mock.MagicMock()
        mock_event.is_set.return_value = True  # simulate timeout

        mock_timer = mock.MagicMock()

        mock_create_timeout_thread.return_value = (mock_timer, mock_event)

        mock_run_command.side_effect = Exception("Mocked exception")

        cur.fetchall.return_value = cursor_results
        conn.cursor.return_value = cur
        mock_get_conn.side_effect = [conn]

        with pytest.raises(AirflowException) as exc_info:
            DatabricksSqlHook(sql_endpoint_name="Test").run(
                sql=sql,
                execution_timeout=execution_timeout,
                handler=fetch_all_handler,
            )

        assert "Timeout threshold exceeded" in str(exc_info.value)


@pytest.mark.parametrize(
    "cursor_descriptions",
    [(("id", "value"),)],
)
def test_create_timeout_thread(mock_get_conn, mock_get_requests, cursor_descriptions):

    cur = mock.MagicMock(
        rowcount=1,
        description=get_cursor_descriptions(cursor_descriptions),
    )

    timeout = timedelta(seconds=1)

    timer, event = create_timeout_thread(cur=cur, execution_timeout=timeout)

    assert timer is not None
    assert event is not None
    assert not event.is_set()

    timer.cancel()


@pytest.mark.parametrize(
    "cursor_descriptions",
    [(("id", "value"),)],
)
def test_create_timeout_thread_no_timeout(
    mock_get_conn,
    mock_get_requests,
    mock_timer,
    cursor_descriptions,
):
    cur = mock.MagicMock(
        rowcount=1,
        description=get_cursor_descriptions(cursor_descriptions),
    )

    timer, timeout_event = create_timeout_thread(
        cur=cur,
        execution_timeout=None,
    )

    mock_timer.assert_not_called()
    assert timer is None
    assert timeout_event is None


def test_get_openlineage_default_schema_with_no_schema_set():
    hook = DatabricksSqlHook()
    assert hook.get_openlineage_default_schema() == "default"


def test_get_openlineage_default_schema_with_schema_set():
    hook = DatabricksSqlHook(schema="my-schema")
    assert hook.get_openlineage_default_schema() == "my-schema"


def test_get_openlineage_database_specific_lineage_with_no_query_id():
    hook = DatabricksSqlHook()
    hook.query_ids = []

    result = hook.get_openlineage_database_specific_lineage(None)
    assert result is None


@mock.patch("airflow.providers.databricks.utils.openlineage.emit_openlineage_events_for_databricks_queries")
def test_get_openlineage_database_specific_lineage_with_single_query_id(mock_emit):
    from airflow.providers.common.compat.openlineage.facet import ExternalQueryRunFacet
    from airflow.providers.openlineage.extractors import OperatorLineage

    hook = DatabricksSqlHook()
    hook.query_ids = ["query1"]
    hook.get_connection = mock.MagicMock()
    hook.get_openlineage_database_info = lambda x: mock.MagicMock(authority="auth", scheme="scheme")

    ti = mock.MagicMock()

    result = hook.get_openlineage_database_specific_lineage(ti)
    mock_emit.assert_called_once_with(
        **{
            "hook": hook,
            "query_ids": ["query1"],
            "query_source_namespace": "scheme://auth",
            "task_instance": ti,
            "query_for_extra_metadata": True,
        }
    )
    assert result == OperatorLineage(
        run_facets={"externalQuery": ExternalQueryRunFacet(externalQueryId="query1", source="scheme://auth")}
    )


@mock.patch("airflow.providers.databricks.utils.openlineage.emit_openlineage_events_for_databricks_queries")
def test_get_openlineage_database_specific_lineage_with_multiple_query_ids(mock_emit):
    hook = DatabricksSqlHook()
    hook.query_ids = ["query1", "query2"]
    hook.get_connection = mock.MagicMock()
    hook.get_openlineage_database_info = lambda x: mock.MagicMock(authority="auth", scheme="scheme")

    ti = mock.MagicMock()

    result = hook.get_openlineage_database_specific_lineage(ti)
    mock_emit.assert_called_once_with(
        **{
            "hook": hook,
            "query_ids": ["query1", "query2"],
            "query_source_namespace": "scheme://auth",
            "task_instance": ti,
            "query_for_extra_metadata": True,
        }
    )
    assert result is None


@mock.patch("importlib.metadata.version", return_value="1.99.0")
def test_get_openlineage_database_specific_lineage_with_old_openlineage_provider(mock_version):
    hook = DatabricksSqlHook()
    hook.query_ids = ["query1", "query2"]
    hook.get_connection = mock.MagicMock()
    hook.get_openlineage_database_info = lambda x: mock.MagicMock(authority="auth", scheme="scheme")

    expected_err = (
        "OpenLineage provider version `1.99.0` is lower than required `2.5.0`, "
        "skipping function `emit_openlineage_events_for_databricks_queries` execution"
    )
    with pytest.raises(AirflowOptionalProviderFeatureException, match=expected_err):
        hook.get_openlineage_database_specific_lineage(mock.MagicMock())


@pytest.mark.parametrize(
    ("df_type", "df_class", "description"),
    [
        pytest.param("pandas", pd.DataFrame, [(("col",))], id="pandas-dataframe"),
        pytest.param(
            "polars",
            pl.DataFrame,
            [(("col", None, None, None, None, None, None))],
            id="polars-dataframe",
        ),
    ],
)
def test_get_df(df_type, df_class, description):
    hook = DatabricksSqlHook()
    statement = "SQL"
    column = "col"
    result_sets = [("row1",), ("row2",)]

    with mock.patch(
        "airflow.providers.databricks.hooks.databricks_sql.DatabricksSqlHook.get_conn"
    ) as mock_get_conn:
        if df_type == "pandas":
            # Setup for pandas test case
            mock_cursor = mock.MagicMock()
            mock_cursor.description = description
            mock_cursor.fetchall.return_value = result_sets
            mock_get_conn.return_value.cursor.return_value = mock_cursor
        else:
            # Setup for polars test case
            mock_execute = mock.MagicMock()
            mock_execute.description = description
            mock_execute.fetchall.return_value = result_sets

            mock_cursor = mock.MagicMock()
            mock_cursor.execute.return_value = mock_execute
            mock_get_conn.return_value.cursor.return_value = mock_cursor

        df = hook.get_df(statement, df_type=df_type)
        mock_cursor.execute.assert_called_once_with(statement)

        if df_type == "pandas":
            mock_cursor.fetchall.assert_called_once_with()
            assert df.columns[0] == column
            assert df.iloc[0, 0] == "row1"
            assert df.iloc[1, 0] == "row2"
        else:
            mock_execute.fetchall.assert_called_once_with()
            assert df.columns[0] == column
            assert df.row(0)[0] == result_sets[0][0]
            assert df.row(1)[0] == result_sets[1][0]

        assert isinstance(df, df_class)


class TestGetSqlEndpointByName:
    """Tests for _get_sql_endpoint_by_name with both 'warehouses' and legacy 'endpoints' API response keys."""

    @patch("airflow.providers.databricks.hooks.databricks_base.requests")
    def test_resolve_warehouse_name_with_warehouses_key(self, mock_requests):
        """Test that the current API response format with 'warehouses' key works."""
        mock_requests.codes.ok = 200
        mock_requests.get.return_value.json.return_value = {
            "warehouses": [
                {
                    "id": "abc123",
                    "name": "My Warehouse",
                    "odbc_params": {
                        "hostname": "xx.cloud.databricks.com",
                        "path": "/sql/1.0/warehouses/abc123",
                    },
                }
            ]
        }
        type(mock_requests.get.return_value).status_code = PropertyMock(return_value=200)

        hook = DatabricksSqlHook(sql_endpoint_name="My Warehouse")
        endpoint = hook._get_sql_endpoint_by_name("My Warehouse")
        assert endpoint["id"] == "abc123"
        assert endpoint["odbc_params"]["path"] == "/sql/1.0/warehouses/abc123"

    @patch("airflow.providers.databricks.hooks.databricks_base.requests")
    def test_resolve_warehouse_name_with_legacy_endpoints_key(self, mock_requests):
        """Test that the legacy API response format with 'endpoints' key still works."""
        mock_requests.codes.ok = 200
        mock_requests.get.return_value.json.return_value = {
            "endpoints": [
                {
                    "id": "def456",
                    "name": "Legacy Endpoint",
                    "odbc_params": {
                        "hostname": "xx.cloud.databricks.com",
                        "path": "/sql/1.0/endpoints/def456",
                    },
                }
            ]
        }
        type(mock_requests.get.return_value).status_code = PropertyMock(return_value=200)

        hook = DatabricksSqlHook(sql_endpoint_name="Legacy Endpoint")
        endpoint = hook._get_sql_endpoint_by_name("Legacy Endpoint")
        assert endpoint["id"] == "def456"
        assert endpoint["odbc_params"]["path"] == "/sql/1.0/endpoints/def456"

    @patch("airflow.providers.databricks.hooks.databricks_base.requests")
    def test_resolve_warehouse_name_not_found(self, mock_requests):
        """Test that a clear error is raised when the warehouse name doesn't match any warehouse."""
        mock_requests.codes.ok = 200
        mock_requests.get.return_value.json.return_value = {
            "warehouses": [
                {
                    "id": "abc123",
                    "name": "Some Other Warehouse",
                    "odbc_params": {"path": "/sql/1.0/warehouses/abc123"},
                }
            ]
        }
        type(mock_requests.get.return_value).status_code = PropertyMock(return_value=200)

        hook = DatabricksSqlHook(sql_endpoint_name="Nonexistent Warehouse")
        with pytest.raises(ValueError, match="Can't find Databricks SQL warehouse with name"):
            hook._get_sql_endpoint_by_name("Nonexistent Warehouse")

    @patch("airflow.providers.databricks.hooks.databricks_base.requests")
    def test_resolve_warehouse_name_empty_response(self, mock_requests):
        """Test that a clear error is raised when the API returns no warehouses."""
        mock_requests.codes.ok = 200
        mock_requests.get.return_value.json.return_value = {}
        type(mock_requests.get.return_value).status_code = PropertyMock(return_value=200)

        hook = DatabricksSqlHook(sql_endpoint_name="Test")
        with pytest.raises(RuntimeError, match="Can't list Databricks SQL warehouses"):
            hook._get_sql_endpoint_by_name("Test")


@mock.patch("airflow.providers.databricks.hooks.databricks_sql.sql.connect")
def test_get_conn_passes_query_tags_via_session_configuration(mock_connect, mock_get_requests):
    """query_tags must be injected into session_configuration['QUERY_TAGS'], not sql.connect(query_tags=)."""
    hook = DatabricksSqlHook(
        databricks_conn_id=DEFAULT_CONN_ID,
        http_path=HTTP_PATH,
        query_tags={"airflow_dag_id": "dag_1", "airflow_task_id": "task_1"},
    )

    hook.get_conn()

    mock_connect.assert_called_once()
    session_cfg = mock_connect.call_args.kwargs["session_configuration"]
    assert session_cfg is not None
    assert "QUERY_TAGS" in session_cfg
    query_tags_str = session_cfg["QUERY_TAGS"]
    assert "airflow_dag_id:dag_1" in query_tags_str
    assert "airflow_task_id:task_1" in query_tags_str
    assert "query_tags" not in mock_connect.call_args.kwargs


@mock.patch("airflow.providers.databricks.hooks.databricks_sql.sql.connect")
def test_get_conn_merges_query_tags_with_existing_session_configuration(mock_connect, mock_get_requests):
    """Existing QUERY_TAGS in session_configuration must be preserved and new tags appended."""
    hook = DatabricksSqlHook(
        databricks_conn_id=DEFAULT_CONN_ID,
        http_path=HTTP_PATH,
        session_configuration={"QUERY_TAGS": "existing_tag:existing_value"},
        query_tags={"airflow_dag_id": "dag_1"},
    )

    hook.get_conn()

    mock_connect.assert_called_once()
    session_cfg = mock_connect.call_args.kwargs["session_configuration"]
    query_tags_str = session_cfg["QUERY_TAGS"]
    assert "existing_tag:existing_value" in query_tags_str
    assert "airflow_dag_id:dag_1" in query_tags_str


@mock.patch("airflow.providers.databricks.hooks.databricks_sql.sql.connect")
def test_get_conn_no_query_tags(mock_connect, mock_get_requests):
    """When no query_tags are provided, session_configuration should not gain a QUERY_TAGS key."""
    hook = DatabricksSqlHook(
        databricks_conn_id=DEFAULT_CONN_ID,
        http_path=HTTP_PATH,
    )

    hook.get_conn()

    mock_connect.assert_called_once()
    session_cfg = mock_connect.call_args.kwargs.get("session_configuration")
    assert session_cfg is None or "QUERY_TAGS" not in session_cfg


class TestFormatQueryTags:
    def test_simple_values(self):
        result = _format_query_tags({"dag_id": "my_dag", "task_id": "my_task"})
        assert "dag_id:my_dag" in result
        assert "task_id:my_task" in result

    def test_none_values_omitted(self):
        result = _format_query_tags({"dag_id": "my_dag", "map_index": None})
        assert "dag_id:my_dag" in result
        assert "map_index" not in result

    def test_empty_dict_returns_empty_string(self):
        assert _format_query_tags({}) == ""

    def test_value_escaping_comma(self):
        result = _format_query_tag_value("a,b")
        assert result == "a\\,b"

    def test_value_escaping_colon(self):
        result = _format_query_tag_value("a:b")
        assert result == "a\\:b"

    def test_value_escaping_backslash(self):
        result = _format_query_tag_value("a\\b")
        assert result == "a\\\\b"

    def test_value_truncated_at_128_chars(self):
        long_value = "x" * 200
        result = _format_query_tag_value(long_value)
        assert len(result) == 128

    def test_format_query_tags_roundtrip(self):
        tags = {"airflow_dag_id": "dag:1", "airflow_run_id": "run,2"}
        result = _format_query_tags(tags)
        assert "airflow_dag_id:dag\\:1" in result
        assert "airflow_run_id:run\\,2" in result


class TestDatabricksSqlHookQueryTagsParamOrder:
    """Ensure moving query_tags after caller preserves positional backward compatibility."""

    def test_query_tags_keyword_sets_field(self):
        """query_tags kwarg must be stored on the instance."""
        with patch(
            "airflow.providers.databricks.hooks.databricks_sql.BaseDatabricksHook.__init__",
            return_value=None,
        ) as mock_base_init:
            hook = DatabricksSqlHook.__new__(DatabricksSqlHook)
            DatabricksSqlHook.__init__(
                hook,
                DEFAULT_CONN_ID,
                query_tags={"key": "val"},
            )
            assert hook.query_tags == {"key": "val"}
            # caller is forwarded to BaseDatabricksHook.__init__; verify the default was passed
            assert mock_base_init.call_args.kwargs.get("caller") == "DatabricksSqlHook"

    def test_caller_positional_not_confused_with_query_tags(self):
        """Passing caller as the 8th positional arg must not end up in query_tags."""
        with patch(
            "airflow.providers.databricks.hooks.databricks_sql.BaseDatabricksHook.__init__",
            return_value=None,
        ) as mock_base_init:
            hook = DatabricksSqlHook.__new__(DatabricksSqlHook)
            # positional order: conn_id, http_path, sql_endpoint, session_cfg,
            #                   http_headers, catalog, schema, caller
            DatabricksSqlHook.__init__(
                hook,
                DEFAULT_CONN_ID,  # databricks_conn_id
                None,  # http_path
                None,  # sql_endpoint_name
                None,  # session_configuration
                None,  # http_headers
                None,  # catalog
                None,  # schema
                "CustomCaller",  # caller (8th positional)
            )
            # caller is forwarded to BaseDatabricksHook.__init__; verify it was not
            # confused with query_tags (which comes after caller)
            assert mock_base_init.call_args.kwargs.get("caller") == "CustomCaller"
            assert hook.query_tags is None
