test_transformation.py

import agate
from datetime import datetime
import json
import logging
import typing as t
from pathlib import Path
from unittest.mock import patch

from sqlmesh.dbt.util import DBT_VERSION

import pytest
from dbt.adapters.base import BaseRelation

if DBT_VERSION >= (1, 4, 0):
    from dbt.exceptions import CompilationError
else:
    from dbt.exceptions import CompilationException as CompilationError  # type: ignore
import time_machine
from pytest_mock.plugin import MockerFixture
from sqlglot import exp, parse_one
from sqlmesh.core import dialect as d
from sqlmesh.core.environment import EnvironmentNamingInfo
from sqlmesh.core.macros import RuntimeStage
from sqlmesh.core.renderer import render_statements
from sqlmesh.core.audit import StandaloneAudit
from sqlmesh.core.context import Context
from sqlmesh.core.console import get_console
from sqlmesh.core.model import (
    EmbeddedKind,
    FullKind,
    IncrementalByTimeRangeKind,
    IncrementalByUniqueKeyKind,
    IncrementalUnmanagedKind,
    ManagedKind,
    SqlModel,
    ViewKind,
)
from sqlmesh.core.model.kind import (
    SCDType2ByColumnKind,
    SCDType2ByTimeKind,
    OnDestructiveChange,
    OnAdditiveChange,
)
from sqlmesh.core.state_sync.db.snapshot import _snapshot_to_json
from sqlmesh.dbt.builtin import _relation_info_to_relation
from sqlmesh.dbt.common import Dependencies
from sqlmesh.dbt.column import (
    ColumnConfig,
    column_descriptions_to_sqlmesh,
    column_types_to_sqlmesh,
)
from sqlmesh.dbt.context import DbtContext
from sqlmesh.dbt.model import Materialization, ModelConfig
from sqlmesh.dbt.project import Project
from sqlmesh.dbt.relation import Policy
from sqlmesh.dbt.seed import SeedConfig
from sqlmesh.dbt.target import (
    BigQueryConfig,
    DuckDbConfig,
    SnowflakeConfig,
    ClickhouseConfig,
    PostgresConfig,
)
from sqlmesh.dbt.test import TestConfig
from sqlmesh.utils.errors import ConfigError, MacroEvalError, SQLMeshError
from sqlmesh.utils.jinja import MacroReference

pytestmark = [pytest.mark.dbt, pytest.mark.slow]


def test_model_name(dbt_dummy_postgres_config: PostgresConfig):
    context = DbtContext()
    context._target = dbt_dummy_postgres_config
    assert ModelConfig(schema="foo", path="models/bar.sql").canonical_name(context) == "foo.bar"
    assert (
        ModelConfig(schema="foo", path="models/bar.sql", alias="baz").canonical_name(context)
        == "foo.baz"
    )
    assert (
        ModelConfig(
            database="dbname", schema="foo", path="models/bar.sql", alias="baz"
        ).canonical_name(context)
        == "foo.baz"
    )
    assert (
        ModelConfig(
            database="other", schema="foo", path="models/bar.sql", alias="baz"
        ).canonical_name(context)
        == "other.foo.baz"
    )


def test_materialization():
    context = DbtContext()
    context.project_name = "Test"
    context.target = DuckDbConfig(name="target", schema="foo")

    with patch.object(get_console(), "log_warning") as mock_logger:
        model_config = ModelConfig(
            name="model", alias="model", schema="schema", materialized="materialized_view"
        )

    assert (
        "SQLMesh does not support the 'materialized_view' model materialization. Falling back to the 'view' materialization."
        in mock_logger.call_args[0][0]
    )
    assert model_config.materialized == "view"

    # clickhouse "dictionary" materialization
    with pytest.raises(ConfigError):
        ModelConfig(name="model", alias="model", schema="schema", materialized="dictionary")


def test_model_kind():
    context = DbtContext()
    context.project_name = "Test"
    context.target = DuckDbConfig(name="target", schema="foo")

    assert ModelConfig(materialized=Materialization.TABLE).model_kind(context) == FullKind()
    assert ModelConfig(materialized=Materialization.VIEW).model_kind(context) == ViewKind()
    assert ModelConfig(materialized=Materialization.EPHEMERAL).model_kind(context) == EmbeddedKind()
    assert ModelConfig(
        materialized=Materialization.SNAPSHOT,
        unique_key=["id"],
        updated_at="updated_at",
        strategy="timestamp",
    ).model_kind(context) == SCDType2ByTimeKind(
        unique_key=["id"],
        valid_from_name="dbt_valid_from",
        valid_to_name="dbt_valid_to",
        updated_at_as_valid_from=True,
        updated_at_name="updated_at",
        dialect="duckdb",
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.ALLOW,
    )
    assert ModelConfig(
        materialized=Materialization.SNAPSHOT,
        unique_key=["id"],
        strategy="check",
        check_cols=["foo"],
    ).model_kind(context) == SCDType2ByColumnKind(
        unique_key=["id"],
        valid_from_name="dbt_valid_from",
        valid_to_name="dbt_valid_to",
        columns=["foo"],
        execution_time_as_valid_from=True,
        dialect="duckdb",
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.ALLOW,
    )
    assert ModelConfig(
        materialized=Materialization.SNAPSHOT,
        unique_key=["id"],
        strategy="check",
        check_cols=["foo"],
        dialect="bigquery",
    ).model_kind(context) == SCDType2ByColumnKind(
        unique_key=["id"],
        valid_from_name="dbt_valid_from",
        valid_to_name="dbt_valid_to",
        columns=["foo"],
        execution_time_as_valid_from=True,
        dialect="bigquery",
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.ALLOW,
    )

    assert ModelConfig(materialized=Materialization.INCREMENTAL, time_column="foo").model_kind(
        context
    ) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        forward_only=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )
    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        incremental_strategy="delete+insert",
        forward_only=False,
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )
    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        incremental_strategy="insert_overwrite",
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        forward_only=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )
    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        unique_key=["bar"],
        dialect="bigquery",
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="bigquery",
        forward_only=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL, unique_key=["bar"], incremental_strategy="merge"
    ).model_kind(context) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="duckdb",
        forward_only=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    dbt_incremental_predicate = "DBT_INTERNAL_DEST.session_start > dateadd(day, -7, current_date)"
    expected_sqlmesh_predicate = parse_one(
        "__MERGE_TARGET__.session_start > DATEADD(day, -7, CURRENT_DATE)"
    )
    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        unique_key=["bar"],
        incremental_strategy="merge",
        dialect="postgres",
        incremental_predicates=[dbt_incremental_predicate],
    ).model_kind(context) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="postgres",
        forward_only=True,
        disable_restatement=False,
        merge_filter=expected_sqlmesh_predicate,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(materialized=Materialization.INCREMENTAL, unique_key=["bar"]).model_kind(
        context
    ) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="duckdb",
        forward_only=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL, unique_key=["bar"], full_refresh=False
    ).model_kind(context) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="duckdb",
        forward_only=True,
        disable_restatement=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL, unique_key=["bar"], full_refresh=True
    ).model_kind(context) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="duckdb",
        forward_only=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL, unique_key=["bar"], disable_restatement=True
    ).model_kind(context) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="duckdb",
        forward_only=True,
        disable_restatement=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        unique_key=["bar"],
        disable_restatement=True,
        full_refresh=True,
    ).model_kind(context) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="duckdb",
        forward_only=True,
        disable_restatement=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        unique_key=["bar"],
        disable_restatement=True,
        full_refresh=False,
        auto_restatement_cron="0 0 * * *",
    ).model_kind(context) == IncrementalByUniqueKeyKind(
        unique_key=["bar"],
        dialect="duckdb",
        forward_only=True,
        disable_restatement=True,
        auto_restatement_cron="0 0 * * *",
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    # Test incompatibile incremental strategies
    for incremental_strategy in ("delete+insert", "insert_overwrite", "append"):
        assert ModelConfig(
            materialized=Materialization.INCREMENTAL,
            unique_key=["bar"],
            incremental_strategy=incremental_strategy,
        ).model_kind(context) == IncrementalByUniqueKeyKind(
            unique_key=["bar"],
            dialect="duckdb",
            forward_only=True,
            disable_restatement=False,
            on_destructive_change=OnDestructiveChange.IGNORE,
            on_additive_change=OnAdditiveChange.IGNORE,
        )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL, time_column="foo", incremental_strategy="merge"
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        forward_only=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        incremental_strategy="merge",
        full_refresh=True,
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        forward_only=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        incremental_strategy="merge",
        full_refresh=False,
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        forward_only=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        incremental_strategy="append",
        disable_restatement=True,
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        forward_only=True,
        disable_restatement=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        incremental_strategy="insert_overwrite",
        partition_by={"field": "bar"},
        forward_only=False,
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        time_column="foo",
        incremental_strategy="insert_overwrite",
        partition_by={"field": "bar"},
        forward_only=False,
        auto_restatement_cron="0 0 * * *",
        auto_restatement_intervals=3,
    ).model_kind(context) == IncrementalByTimeRangeKind(
        time_column="foo",
        dialect="duckdb",
        forward_only=False,
        auto_restatement_cron="0 0 * * *",
        auto_restatement_intervals=3,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        incremental_strategy="insert_overwrite",
        partition_by={"field": "bar"},
    ).model_kind(context) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(materialized=Materialization.INCREMENTAL).model_kind(
        context
    ) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(materialized=Materialization.INCREMENTAL, forward_only=False).model_kind(
        context
    ) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        disable_restatement=False,
        forward_only=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL, incremental_strategy="append"
    ).model_kind(context) == IncrementalUnmanagedKind(
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL, incremental_strategy="append", full_refresh=None
    ).model_kind(context) == IncrementalUnmanagedKind(
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        incremental_strategy="insert_overwrite",
        partition_by={"field": "bar", "data_type": "int64"},
    ).model_kind(context) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        incremental_strategy="insert_overwrite",
        partition_by={"field": "bar", "data_type": "int64"},
        full_refresh=False,
    ).model_kind(context) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        disable_restatement=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        incremental_strategy="insert_overwrite",
        partition_by={"field": "bar", "data_type": "int64"},
        disable_restatement=True,
        full_refresh=True,
    ).model_kind(context) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        disable_restatement=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        incremental_strategy="insert_overwrite",
        partition_by={"field": "bar", "data_type": "int64"},
        disable_restatement=True,
    ).model_kind(context) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        disable_restatement=True,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert ModelConfig(
        materialized=Materialization.INCREMENTAL,
        incremental_strategy="insert_overwrite",
        auto_restatement_cron="0 0 * * *",
    ).model_kind(context) == IncrementalUnmanagedKind(
        insert_overwrite=True,
        auto_restatement_cron="0 0 * * *",
        disable_restatement=False,
        on_destructive_change=OnDestructiveChange.IGNORE,
        on_additive_change=OnAdditiveChange.IGNORE,
    )

    assert (
        ModelConfig(materialized=Materialization.DYNAMIC_TABLE, target_lag="1 hour").model_kind(
            context
        )
        == ManagedKind()
    )


def test_model_kind_snapshot_bigquery():
    context = DbtContext()
    context.project_name = "Test"
    context.target = BigQueryConfig(name="target", schema="foo", project="bar")

    assert ModelConfig(
        materialized=Materialization.SNAPSHOT,
        unique_key=["id"],
        updated_at="updated_at",
        strategy="timestamp",
    ).model_kind(context) == SCDType2ByTimeKind(
        unique_key=["id"],
        valid_from_name="dbt_valid_from",
        valid_to_name="dbt_valid_to",
        updated_at_as_valid_from=True,
        updated_at_name="updated_at",
        time_data_type=exp.DataType.build("TIMESTAMPTZ"),
        dialect="bigquery",
        on_destructive_change=OnDestructiveChange.IGNORE,
    )

    # time_data_type is bigquery version even though model dialect is DuckDB
    # because model target is BigQuery
    assert ModelConfig(
        materialized=Materialization.SNAPSHOT,
        unique_key=["id"],
        updated_at="updated_at",
        strategy="timestamp",
        dialect="duckdb",
    ).model_kind(context) == SCDType2ByTimeKind(
        unique_key=["id"],
        valid_from_name="dbt_valid_from",
        valid_to_name="dbt_valid_to",
        updated_at_as_valid_from=True,
        updated_at_name="updated_at",
        time_data_type=exp.DataType.build("TIMESTAMPTZ"),  # bigquery version
        dialect="duckdb",
        on_destructive_change=OnDestructiveChange.IGNORE,
    )


def test_model_columns():
    model = ModelConfig(
        alias="test",
        target_schema="foo",
        table_name="bar",
        sql="SELECT * FROM baz",
        columns={
            "ADDRESS": ColumnConfig(
                name="address", data_type="text", description="Business address"
            ),
            "ZIPCODE": ColumnConfig(
                name="zipcode", data_type="varchar(5)", description="Business zipcode"
            ),
            "DATE": ColumnConfig(
                name="date", data_type="timestamp_ntz", description="Contract date"
            ),
        },
    )

    expected_column_types = {
        "ADDRESS": exp.DataType.build("text"),
        "ZIPCODE": exp.DataType.build("varchar(5)"),
        "DATE": exp.DataType.build("timestamp_ntz", dialect="snowflake"),
    }
    expected_column_descriptions = {
        "ADDRESS": "Business address",
        "ZIPCODE": "Business zipcode",
        "DATE": "Contract date",
    }

    assert column_types_to_sqlmesh(model.columns, "snowflake") == expected_column_types
    assert column_descriptions_to_sqlmesh(model.columns) == expected_column_descriptions

    context = DbtContext()
    context.project_name = "Foo"
    context.target = SnowflakeConfig(
        name="target", schema="test", database="test", account="foo", user="bar", password="baz"
    )
    sqlmesh_model = model.to_sqlmesh(context)
    assert sqlmesh_model.columns_to_types == expected_column_types
    assert sqlmesh_model.column_descriptions == expected_column_descriptions


def test_seed_columns():
    seed = SeedConfig(
        name="foo",
        package="package",
        path=Path("examples/sushi_dbt/seeds/waiter_names.csv"),
        columns={
            "id": ColumnConfig(name="id", data_type="text", description="The ID"),
            "name": ColumnConfig(name="name", data_type="text", description="The name"),
        },
    )

    expected_column_types = {
        "id": exp.DataType.build("text"),
        "name": exp.DataType.build("text"),
    }
    expected_column_descriptions = {
        "id": "The ID",
        "name": "The name",
    }

    context = DbtContext()
    context.project_name = "Foo"
    context.target = DuckDbConfig(name="target", schema="test")
    sqlmesh_seed = seed.to_sqlmesh(context)
    assert sqlmesh_seed.columns_to_types == expected_column_types
    assert sqlmesh_seed.column_descriptions == expected_column_descriptions


def test_seed_column_types():
    seed = SeedConfig(
        name="foo",
        package="package",
        path=Path("examples/sushi_dbt/seeds/waiter_names.csv"),
        column_types={
            "id": "text",
            "name": "text",
        },
        columns={
            "name": ColumnConfig(name="name", description="The name"),
        },
        quote_columns=True,
    )

    expected_column_types = {
        "id": exp.DataType.build("text"),
        "name": exp.DataType.build("text"),
    }
    expected_column_descriptions = {
        "name": "The name",
    }

    context = DbtContext()
    context.project_name = "Foo"
    context.target = DuckDbConfig(name="target", schema="test")
    sqlmesh_seed = seed.to_sqlmesh(context)

    assert sqlmesh_seed.columns_to_types == expected_column_types
    assert sqlmesh_seed.column_descriptions == expected_column_descriptions


def test_seed_column_inference(tmp_path):
    seed_csv = tmp_path / "seed.csv"
    with open(seed_csv, "w", encoding="utf-8") as fd:
        fd.write("int_col,double_col,datetime_col,date_col,boolean_col,text_col\n")
        fd.write("1,1.2,2021-01-01 00:00:00,2021-01-01,true,foo\n")
        fd.write("2,2.3,2021-01-02 00:00:00,2021-01-02,false,bar\n")
        fd.write("null,,null,,,null\n")

    seed = SeedConfig(
        name="test_model",
        package="package",
        path=Path(seed_csv),
    )

    context = DbtContext()
    context.project_name = "Foo"
    context.target = DuckDbConfig(name="target", schema="test")
    sqlmesh_seed = seed.to_sqlmesh(context)
    assert sqlmesh_seed.columns_to_types == {
        "int_col": exp.DataType.build("int")
        if DBT_VERSION >= (1, 8, 0)
        else exp.DataType.build("double"),
        "double_col": exp.DataType.build("double"),
        "datetime_col": exp.DataType.build("datetime"),
        "date_col": exp.DataType.build("date"),
        "boolean_col": exp.DataType.build("boolean"),
        "text_col": exp.DataType.build("text"),
    }


def test_seed_single_whitespace_is_na(tmp_path):
    seed_csv = tmp_path / "seed.csv"
    with open(seed_csv, "w", encoding="utf-8") as fd:
        fd.write("col_a, col_b\n")
        fd.write(" ,1\n")
        fd.write("2, \n")

    seed = SeedConfig(
        name="test_model",
        package="foo",
        path=Path(seed_csv),
    )

    context = DbtContext()
    context.project_name = "foo"
    context.target = DuckDbConfig(name="target", schema="test")
    sqlmesh_seed = seed.to_sqlmesh(context)
    assert sqlmesh_seed.columns_to_types == {
        "col_a": exp.DataType.build("int"),
        "col_b": exp.DataType.build("int"),
    }

    df = next(sqlmesh_seed.render_seed())
    assert df["col_a"].to_list() == [None, 2]
    assert df["col_b"].to_list() == [1, None]


def test_seed_partial_column_inference(tmp_path):
    seed_csv = tmp_path / "seed.csv"
    with open(seed_csv, "w", encoding="utf-8") as fd:
        fd.write("int_col,double_col,datetime_col,boolean_col\n")
        fd.write("1,1.2,2021-01-01 00:00:00,true\n")
        fd.write("2,2.3,2021-01-02 00:00:00,false\n")
        fd.write("null,,null,\n")

    seed = SeedConfig(
        name="test_model",
        package="package",
        path=Path(seed_csv),
        column_types={
            "double_col": "double",
        },
        columns={
            "int_col": ColumnConfig(
                name="int_col", data_type="int", description="Description with type."
            ),
            "datetime_col": ColumnConfig(
                name="datetime_col", description="Description without type."
            ),
            "boolean_col": ColumnConfig(name="boolean_col"),
        },
    )

    expected_column_types = {
        "int_col": exp.DataType.build("int"),
        "double_col": exp.DataType.build("double"),
        "datetime_col": exp.DataType.build("datetime"),
        "boolean_col": exp.DataType.build("boolean"),
    }

    expected_column_descriptions = {
        "int_col": "Description with type.",
        "datetime_col": "Description without type.",
    }

    context = DbtContext()
    context.project_name = "Foo"
    context.target = DuckDbConfig(name="target", schema="test")
    sqlmesh_seed = seed.to_sqlmesh(context)
    assert sqlmesh_seed.columns_to_types == expected_column_types
    assert sqlmesh_seed.column_descriptions == expected_column_descriptions

    # Check that everything still lines up
    seed_df = next(sqlmesh_seed.render_seed())
    assert list(seed_df.columns) == list(sqlmesh_seed.columns_to_types.keys())


def test_seed_column_order(tmp_path):
    seed_csv = tmp_path / "seed.csv"

    with open(seed_csv, "w", encoding="utf-8") as fd:
        fd.writelines("\n".join(["id,name", "0,Toby", "1,Tyson", "2,Ryan"]))

    seed = SeedConfig(
        name="test_model",
        package="package",
        path=Path(seed_csv),
        columns={
            "id": ColumnConfig(name="id"),
            "name": ColumnConfig(name="name", data_type="varchar"),
        },
    )

    context = DbtContext()
    context.project_name = "Foo"
    context.target = DuckDbConfig(name="target", schema="test")
    sqlmesh_seed = seed.to_sqlmesh(context)

    # Check that everything still lines up
    seed_df = next(sqlmesh_seed.render_seed())
    assert list(seed_df.columns) == list(sqlmesh_seed.columns_to_types.keys())


def test_agate_integer_cast():
    # Not all dbt versions have agate.Integer
    if DBT_VERSION < (1, 7, 0):
        pytest.skip("agate.Integer not available")

    from sqlmesh.dbt.seed import Integer

    agate_integer = Integer(null_values=("null", ""))
    assert agate_integer.cast("1") == 1
    assert agate_integer.cast(1) == 1
    assert agate_integer.cast("null") is None
    assert agate_integer.cast("") is None

    with pytest.raises(agate.exceptions.CastError):
        agate_integer.cast("1.2")

    with pytest.raises(agate.exceptions.CastError):
        agate_integer.cast(1.2)

    with pytest.raises(agate.exceptions.CastError):
        agate_integer.cast(datetime.now())


@pytest.mark.xdist_group("dbt_manifest")
def test_model_dialect(sushi_test_project: Project, assert_exp_eq):
    model_config = ModelConfig(
        name="model",
        package_name="package",
        schema="sushi",
        alias="table",
        sql="SELECT 1 AS `one` FROM {{ schema }}",
    )
    context = sushi_test_project.context

    # cannot parse model sql without specifying bigquery dialect
    with pytest.raises(ConfigError):
        model_config.to_sqlmesh(context).render_query_or_raise().sql()

    model_config = ModelConfig(
        name="model",
        package_name="package",
        schema="sushi",
        alias="table",
        sql="SELECT 1 AS `one` FROM {{ schema }}",
        dialect="bigquery",
    )
    assert_exp_eq(
        model_config.to_sqlmesh(context).render_query_or_raise().sql(),
        'SELECT 1 AS "one" FROM "sushi" AS "sushi"',
    )


@pytest.mark.xdist_group("dbt_manifest")
@pytest.mark.parametrize(
    "model_fqn", ['"memory"."sushi"."waiters"', '"memory"."sushi"."waiter_names"']
)
def test_hooks(sushi_test_dbt_context: Context, model_fqn: str):
    engine_adapter = sushi_test_dbt_context.engine_adapter
    waiters = sushi_test_dbt_context.models[model_fqn]

    logger = logging.getLogger("sqlmesh.dbt.builtin")
    with patch.object(logger, "debug") as mock_logger:
        engine_adapter.execute(
            waiters.render_pre_statements(
                engine_adapter=engine_adapter, execution_time="2023-01-01"
            )
        )
    assert "pre-hook" in mock_logger.call_args[0][0]

    with patch.object(logger, "debug") as mock_logger:
        engine_adapter.execute(
            waiters.render_post_statements(
                engine_adapter=sushi_test_dbt_context.engine_adapter, execution_time="2023-01-01"
            )
        )
    assert "post-hook" in mock_logger.call_args[0][0]


@pytest.mark.xdist_group("dbt_manifest")
def test_target_jinja(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ target.name }}") == "in_memory"
    assert context.render("{{ target.schema }}") == "sushi"
    assert context.render("{{ target.type }}") == "duckdb"
    # Path and Profile name are not included in serializable fields
    assert context.render("{{ target.path }}") == "None"
    assert context.render("{{ target.profile_name }}") == "None"


@pytest.mark.xdist_group("dbt_manifest")
def test_project_name_jinja(sushi_test_project: Project):
    context = sushi_test_project.context
    assert context.render("{{ project_name }}") == "sushi"


@pytest.mark.xdist_group("dbt_manifest")
def test_schema_jinja(sushi_test_project: Project, assert_exp_eq):
    model_config = ModelConfig(
        name="model",
        package_name="package",
        schema="sushi",
        alias="table",
        sql="SELECT 1 AS one FROM {{ schema }}",
    )
    context = sushi_test_project.context
    assert_exp_eq(
        model_config.to_sqlmesh(context).render_query_or_raise().sql(),
        'SELECT 1 AS "one" FROM "sushi" AS "sushi"',
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_config_jinja(sushi_test_project: Project):
    hook = "{{ config(alias='bar') }} {{ config.alias }}"
    model_config = ModelConfig(
        name="model",
        package_name="package",
        schema="sushi",
        sql="""SELECT 1 AS one FROM foo""",
        alias="model",
        **{"pre-hook": hook},
    )
    context = sushi_test_project.context
    model = t.cast(SqlModel, model_config.to_sqlmesh(context))
    assert hook in model.pre_statements[0].sql()
    assert model.render_pre_statements()[0].sql() == '"bar"'


@pytest.mark.xdist_group("dbt_manifest")
def test_model_this(assert_exp_eq, sushi_test_project: Project):
    model_config = ModelConfig(
        name="model",
        package_name="package",
        schema="schema",
        alias="test",
        sql="SELECT 1 AS one FROM {{ this.identifier }}",
    )
    context = sushi_test_project.context
    assert_exp_eq(
        model_config.to_sqlmesh(context).render_query_or_raise().sql(),
        'SELECT 1 AS "one" FROM "test" AS "test"',
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_test_this(assert_exp_eq, sushi_test_project: Project):
    test_config = TestConfig(
        name="test",
        alias="alias",
        database="database",
        schema_="schema",
        standalone=True,
        sql="SELECT 1 AS one FROM {{ this.identifier }}",
    )
    context = sushi_test_project.context
    audit = t.cast(StandaloneAudit, test_config.to_sqlmesh(context))
    assert_exp_eq(
        audit.render_audit_query().sql(),
        'SELECT 1 AS "one" FROM "test" AS "test"',
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_test_dialect(assert_exp_eq, sushi_test_project: Project):
    test_config = TestConfig(
        name="test",
        alias="alias",
        database="database",
        schema_="schema",
        standalone=True,
        sql="SELECT 1 AS `one` FROM {{ this.identifier }}",
    )
    context = sushi_test_project.context

    # can't parse test sql without specifying bigquery as default dialect
    with pytest.raises(ConfigError):
        audit = t.cast(StandaloneAudit, test_config.to_sqlmesh(context))
        audit.render_audit_query().sql()

    test_config.dialect_ = "bigquery"
    audit = t.cast(StandaloneAudit, test_config.to_sqlmesh(context))
    assert_exp_eq(
        audit.render_audit_query().sql(),
        'SELECT 1 AS "one" FROM "test" AS "test"',
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_statement(sushi_test_project: Project, runtime_renderer: t.Callable):
    context = sushi_test_project.context
    assert context.target
    engine_adapter = context.target.to_sqlmesh().create_engine_adapter()
    renderer = runtime_renderer(context, engine_adapter=engine_adapter)
    assert (
        renderer(
            "{% set test_var = 'SELECT 1' %}{% call statement('something', fetch_result=True) %} {{ test_var }} {% endcall %}{{ load_result('something').table }}",
        )
        == """| column | data_type |
| ------ | --------- |
| 1      | Integer   |
"""
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_run_query(sushi_test_project: Project, runtime_renderer: t.Callable):
    context = sushi_test_project.context
    assert context.target
    engine_adapter = context.target.to_sqlmesh().create_engine_adapter()
    renderer = runtime_renderer(context, engine_adapter=engine_adapter)
    assert (
        renderer(
            """{% set results = run_query('SELECT 1 UNION ALL SELECT 2') %}{% for val in results.columns[0] %}{{ val }} {% endfor %}"""
        )
        == "1 2 "
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_logging(sushi_test_project: Project, runtime_renderer: t.Callable):
    context = sushi_test_project.context
    assert context.target
    engine_adapter = context.target.to_sqlmesh().create_engine_adapter()
    renderer = runtime_renderer(context, engine_adapter=engine_adapter)

    logger = logging.getLogger("sqlmesh.dbt.builtin")

    # Test log with info=False (default), should only log to file with debug and not to console
    with (
        patch.object(logger, "debug") as mock_debug,
        patch.object(logger, "info") as mock_info,
        patch.object(get_console(), "log_status_update") as mock_console,
    ):
        assert renderer('{{ log("foo") }}') == ""
        mock_debug.assert_called_once()
        assert "foo" in mock_debug.call_args[0][0]
        mock_info.assert_not_called()
        mock_console.assert_not_called()

    # Test log with info=True, should log to info and also call log_status_update
    with (
        patch.object(logger, "debug") as mock_debug,
        patch.object(logger, "info") as mock_info,
        patch.object(get_console(), "log_status_update") as mock_console,
    ):
        assert renderer('{{ log("output to be logged with info", info=true) }}') == ""
        mock_info.assert_called_once()
        assert "output to be logged with info" in mock_info.call_args[0][0]
        mock_debug.assert_not_called()
        mock_console.assert_called_once()
        assert "output to be logged with info" in mock_console.call_args[0][0]

    # Test print function as well, should use debug
    with (
        patch.object(logger, "debug") as mock_logger,
        patch.object(get_console(), "log_status_update") as mock_console,
    ):
        assert renderer('{{ print("bar") }}') == ""
        assert "bar" in mock_logger.call_args[0][0]
        mock_console.assert_not_called()


@pytest.mark.xdist_group("dbt_manifest")
def test_exceptions(sushi_test_project: Project):
    context = sushi_test_project.context

    logger = logging.getLogger("sqlmesh.dbt.builtin")
    with patch.object(logger, "warning") as mock_logger:
        assert context.render('{{ exceptions.warn("Warning") }}') == ""
    assert "Warning" in mock_logger.call_args[0][0]

    with pytest.raises(CompilationError, match="Error"):
        context.render('{{ exceptions.raise_compiler_error("Error") }}')


@pytest.mark.xdist_group("dbt_manifest")
def test_modules(sushi_test_project: Project):
    context = sushi_test_project.context

    # datetime
    assert context.render("{{ modules.datetime.date(2022, 12, 25) }}") == "2022-12-25"

    # pytz
    try:
        assert "UTC" in context.render("{{ modules.pytz.all_timezones }}")
    except AttributeError as error:
        assert "object has no attribute 'pytz'" in str(error)

    # re
    assert context.render("{{ modules.re.search('(?<=abc)def', 'abcdef').group(0) }}") == "def"

    # itertools
    itertools_jinja = "{% for num in modules.itertools.accumulate([5]) %}{{ num }}{% endfor %}"
    assert context.render(itertools_jinja) == "5"


@pytest.mark.xdist_group("dbt_manifest")
def test_flags(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ flags.FULL_REFRESH }}") == "None"
    assert context.render("{{ flags.STORE_FAILURES }}") == "None"
    assert context.render("{{ flags.WHICH }}") == "parse"


@pytest.mark.xdist_group("dbt_manifest")
def test_relation(sushi_test_project: Project):
    context = sushi_test_project.context

    assert (
        context.render("{{ api.Relation }}")
        == "<class 'dbt.adapters.duckdb.relation.DuckDBRelation'>"
    )

    jinja = (
        "{% set relation = api.Relation.create(schema='sushi', identifier='waiters') %}"
        "{{ relation.schema }} {{ relation.identifier}}"
    )

    assert context.render(jinja) == "sushi waiters"


@pytest.mark.xdist_group("dbt_manifest")
def test_column(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ api.Column }}") == "<class 'dbt.adapters.base.column.Column'>"

    jinja = (
        "{% set col = api.Column('foo', 'integer') %}{{ col.is_integer() }} {{ col.is_string()}}"
    )

    assert context.render(jinja) == "True False"


@pytest.mark.xdist_group("dbt_manifest")
def test_quote(sushi_test_project: Project):
    context = sushi_test_project.context

    jinja = "{{ adapter.quote('foo') }} {{ adapter.quote('bar') }}"
    assert context.render(jinja) == '"foo" "bar"'


@pytest.mark.xdist_group("dbt_manifest")
def test_as_filters(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ True | as_bool }}") == "True"
    with pytest.raises(MacroEvalError, match="Failed to convert 'invalid' into boolean."):
        context.render("{{ 'invalid' | as_bool }}")

    assert context.render("{{ 123 | as_number }}") == "123"
    with pytest.raises(MacroEvalError, match="Failed to convert 'invalid' into number."):
        context.render("{{ 'invalid' | as_number }}")

    assert context.render("{{ None | as_text }}") == ""

    assert context.render("{{ None | as_native }}") == "None"


@pytest.mark.xdist_group("dbt_manifest")
def test_set(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ set([1, 1, 2]) }}") == "{1, 2}"
    assert context.render("{{ set(1) }}") == "None"

    assert context.render("{{ set_strict([1, 1, 2]) }}") == "{1, 2}"
    with pytest.raises(TypeError):
        assert context.render("{{ set_strict(1) }}")


@pytest.mark.xdist_group("dbt_manifest")
def test_json(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ tojson({'key': 'value'}) }}") == """{"key": "value"}"""
    assert context.render("{{ tojson(set([1])) }}") == "None"

    assert context.render("""{{ fromjson('{"key": "value"}') }}""") == "{'key': 'value'}"
    assert context.render("""{{ fromjson('invalid') }}""") == "None"


@pytest.mark.xdist_group("dbt_manifest")
def test_yaml(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ toyaml({'key': 'value'}) }}").strip() == "key: value"
    assert context.render("{{ toyaml(invalid) }}", invalid=lambda: "") == "None"

    assert context.render("""{{ fromyaml('key: value') }}""") == "{'key': 'value'}"


@pytest.mark.xdist_group("dbt_manifest")
def test_zip(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ zip([1, 2], ['a', 'b']) }}") == "[(1, 'a'), (2, 'b')]"
    assert context.render("{{ zip(12, ['a', 'b']) }}") == "None"

    assert context.render("{{ zip_strict([1, 2], ['a', 'b']) }}") == "[(1, 'a'), (2, 'b')]"
    with pytest.raises(TypeError):
        context.render("{{ zip_strict(12, ['a', 'b']) }}")


@pytest.mark.xdist_group("dbt_manifest")
def test_dbt_version(sushi_test_project: Project):
    context = sushi_test_project.context

    assert context.render("{{ dbt_version }}").startswith("1.")


@pytest.mark.xdist_group("dbt_manifest")
def test_dbt_on_run_start_end(sushi_test_project: Project):
    # Validate perservation of dbt's order of execution
    assert sushi_test_project.packages["sushi"].on_run_start["sushi-on-run-start-0"].index == 0
    assert sushi_test_project.packages["sushi"].on_run_start["sushi-on-run-start-1"].index == 1
    assert sushi_test_project.packages["sushi"].on_run_end["sushi-on-run-end-0"].index == 0
    assert sushi_test_project.packages["sushi"].on_run_end["sushi-on-run-end-1"].index == 1
    assert (
        sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-0"].index == 0
    )
    assert (
        sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-1"].index == 1
    )
    assert sushi_test_project.packages["customers"].on_run_end["customers-on-run-end-0"].index == 0
    assert sushi_test_project.packages["customers"].on_run_end["customers-on-run-end-1"].index == 1

    assert (
        sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-0"].sql
        == "CREATE TABLE IF NOT EXISTS to_be_executed_first (col VARCHAR);"
    )
    assert (
        sushi_test_project.packages["customers"].on_run_start["customers-on-run-start-1"].sql
        == "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);"
    )
    assert (
        sushi_test_project.packages["customers"].on_run_end["customers-on-run-end-1"].sql
        == "{{ packaged_tables(schemas) }}"
    )

    assert (
        sushi_test_project.packages["sushi"].on_run_start["sushi-on-run-start-0"].sql
        == "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);"
    )
    assert (
        sushi_test_project.packages["sushi"].on_run_end["sushi-on-run-end-0"].sql
        == "{{ create_tables(schemas) }}"
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_parsetime_adapter_call(
    assert_exp_eq, sushi_test_project: Project, sushi_test_dbt_context: Context
):
    model_config = ModelConfig(
        name="model",
        package_name="package",
        alias="test",
        schema="sushi",
        sql="""
            {% set results = run_query('select 1 as one') %}
            SELECT {{ results.columns[0].values()[0] }} AS one FROM {{ this.identifier }}
        """,
    )
    context = sushi_test_project.context

    sqlmesh_model = model_config.to_sqlmesh(context)
    assert sqlmesh_model.render_query() is None
    assert sqlmesh_model.columns_to_types is None
    assert not sqlmesh_model.annotated
    with pytest.raises(SQLMeshError):
        sqlmesh_model.ctas_query()

    engine_adapter = sushi_test_dbt_context.engine_adapter
    assert_exp_eq(
        sqlmesh_model.render_query_or_raise(engine_adapter=engine_adapter).sql(),
        'SELECT 1 AS "one" FROM "test" AS "test"',
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_partition_by(sushi_test_project: Project):
    context = sushi_test_project.context
    context.target = BigQueryConfig(name="production", database="main", schema="sushi")
    model_config = ModelConfig(
        name="model",
        alias="model",
        schema="test",
        package_name="package",
        materialized="table",
        unique_key="ds",
        partition_by={"field": "ds", "granularity": "month"},
        sql="""SELECT 1 AS one, ds FROM foo""",
    )
    date_trunc_expr = model_config.to_sqlmesh(context).partitioned_by[0]
    assert date_trunc_expr.sql(dialect="bigquery") == "DATE_TRUNC(`ds`, MONTH)"
    assert date_trunc_expr.sql() == "DATE_TRUNC('MONTH', \"ds\")"

    model_config.partition_by = {"field": "`ds`", "data_type": "datetime", "granularity": "day"}
    datetime_trunc_expr = model_config.to_sqlmesh(context).partitioned_by[0]
    assert datetime_trunc_expr.sql(dialect="bigquery") == "datetime_trunc(`ds`, DAY)"
    assert datetime_trunc_expr.sql() == 'DATETIME_TRUNC("ds", DAY)'

    model_config.partition_by = {"field": "ds", "data_type": "timestamp", "granularity": "day"}
    timestamp_trunc_expr = model_config.to_sqlmesh(context).partitioned_by[0]
    assert timestamp_trunc_expr.sql(dialect="bigquery") == "timestamp_trunc(`ds`, DAY)"
    assert timestamp_trunc_expr.sql() == 'TIMESTAMP_TRUNC("ds", DAY)'

    model_config.partition_by = {
        "field": "one",
        "data_type": "int64",
        "range": {"start": 0, "end": 10, "interval": 2},
    }
    assert (
        model_config.to_sqlmesh(context).partitioned_by[0].sql()
        == 'RANGE_BUCKET("one", GENERATE_SERIES(0, 10, 2))'
    )

    model_config.partition_by = {"field": "ds", "data_type": "date", "granularity": "day"}
    assert model_config.to_sqlmesh(context).partitioned_by == [exp.to_column("ds", quoted=True)]


@pytest.mark.xdist_group("dbt_manifest")
def test_partition_by_none(sushi_test_project: Project):
    context = sushi_test_project.context
    context.target = BigQueryConfig(name="production", database="main", schema="sushi")
    model_config = ModelConfig(
        name="model",
        alias="model",
        schema="test",
        package_name="package",
        materialized="table",
        unique_key="ds",
        partition_by=None,
        sql="""SELECT 1 AS one, ds FROM foo""",
    )
    assert model_config.partition_by is None


@pytest.mark.xdist_group("dbt_manifest")
def test_relation_info_to_relation():
    assert _relation_info_to_relation(
        {"quote_policy": {}},
        BaseRelation,
        Policy(database=True, schema=True, identifier=True),
    ).quote_policy == Policy(database=True, schema=True, identifier=True)

    assert _relation_info_to_relation(
        {"quote_policy": {"database": None, "schema": None, "identifier": None}},
        BaseRelation,
        Policy(database=True, schema=True, identifier=True),
    ).quote_policy == Policy(database=True, schema=True, identifier=True)

    assert _relation_info_to_relation(
        {"quote_policy": {"database": False, "schema": None, "identifier": None}},
        BaseRelation,
        Policy(database=True, schema=True, identifier=True),
    ).quote_policy == Policy(database=False, schema=True, identifier=True)

    assert _relation_info_to_relation(
        {"quote_policy": {"database": False}},
        BaseRelation,
        Policy(database=True, schema=True, identifier=True),
    ).quote_policy == Policy(database=False, schema=True, identifier=True)

    assert _relation_info_to_relation(
        {"quote_policy": {"database": False, "schema": False, "identifier": False}},
        BaseRelation,
        Policy(database=True, schema=True, identifier=True),
    ).quote_policy == Policy(database=False, schema=False, identifier=False)


@pytest.mark.xdist_group("dbt_manifest")
def test_is_incremental(sushi_test_project: Project, assert_exp_eq, mocker):
    model_config = ModelConfig(
        name="model",
        package_name="package",
        schema="sushi",
        alias="some_table",
        sql="""
        SELECT 1 AS one FROM tbl_a
        {% if is_incremental() %}
        WHERE ds > (SELECT MAX(ds) FROM model)
        {% endif %}
        """,
    )
    context = sushi_test_project.context

    assert_exp_eq(
        model_config.to_sqlmesh(context).render_query_or_raise().sql(),
        'SELECT 1 AS "one" FROM "tbl_a" AS "tbl_a"',
    )

    snapshot = mocker.Mock()
    snapshot.intervals = [1]
    snapshot.is_incremental = True

    assert_exp_eq(
        model_config.to_sqlmesh(context).render_query_or_raise(snapshot=snapshot).sql(),
        'SELECT 1 AS "one" FROM "tbl_a" AS "tbl_a" WHERE "ds" > (SELECT MAX("ds") FROM "model" AS "model")',
    )

    # If the snapshot_table_exists flag was set to False, intervals should be ignored
    assert_exp_eq(
        model_config.to_sqlmesh(context)
        .render_query_or_raise(snapshot=snapshot, snapshot_table_exists=False)
        .sql(),
        'SELECT 1 AS "one" FROM "tbl_a" AS "tbl_a"',
    )

    # If the snapshot_table_exists flag was set to True, intervals should be taken into account
    assert_exp_eq(
        model_config.to_sqlmesh(context)
        .render_query_or_raise(snapshot=snapshot, snapshot_table_exists=True)
        .sql(),
        'SELECT 1 AS "one" FROM "tbl_a" AS "tbl_a" WHERE "ds" > (SELECT MAX("ds") FROM "model" AS "model")',
    )
    snapshot.intervals = []
    assert_exp_eq(
        model_config.to_sqlmesh(context)
        .render_query_or_raise(snaspshot=snapshot, snapshot_table_exists=True)
        .sql(),
        'SELECT 1 AS "one" FROM "tbl_a" AS "tbl_a"',
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_is_incremental_non_incremental_model(sushi_test_project: Project, assert_exp_eq, mocker):
    model_config = ModelConfig(
        name="model",
        package_name="package",
        schema="sushi",
        alias="some_table",
        sql="""
        SELECT 1 AS one FROM tbl_a
        {% if is_incremental() %}
        WHERE ds > (SELECT MAX(ds) FROM model)
        {% endif %}
        """,
    )
    context = sushi_test_project.context

    snapshot = mocker.Mock()
    snapshot.intervals = [1]
    snapshot.is_incremental = False

    assert_exp_eq(
        model_config.to_sqlmesh(context).render_query_or_raise(snapshot=snapshot).sql(),
        'SELECT 1 AS "one" FROM "tbl_a" AS "tbl_a"',
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_dbt_max_partition(sushi_test_project: Project, assert_exp_eq, mocker: MockerFixture):
    model_config = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        schema="sushi",
        partition_by={"field": "`ds`", "data_type": "datetime", "granularity": "month"},
        materialized=Materialization.INCREMENTAL,
        sql="""
        SELECT 1 AS one FROM tbl_a
        {% if is_incremental() %}
        WHERE ds > _dbt_max_partition
        {% endif %}
        """,
    )
    context = sushi_test_project.context
    context.target = BigQueryConfig(
        name="test_target", schema="test_schema", database="test-project"
    )

    pre_statement = model_config.to_sqlmesh(context).pre_statements[-1]  # type: ignore

    assert (
        pre_statement.sql().strip()
        == """
JINJA_STATEMENT_BEGIN;
{% if is_incremental() %}
  DECLARE _dbt_max_partition DATETIME DEFAULT (
    COALESCE((SELECT MAX(PARSE_DATETIME('%Y%m', partition_id)) FROM `{{ target.database }}`.`{{ adapter.resolve_schema(this) }}`.`INFORMATION_SCHEMA.PARTITIONS` AS PARTITIONS WHERE table_name = '{{ adapter.resolve_identifier(this) }}' AND NOT partition_id IS NULL AND partition_id <> '__NULL__'), CAST('1970-01-01' AS DATETIME))
  );
{% endif %}
JINJA_END;""".strip()
    )

    assert d.parse_one(pre_statement.sql()) == pre_statement


@pytest.mark.xdist_group("dbt_manifest")
def test_bigquery_physical_properties(sushi_test_project: Project, mocker: MockerFixture):
    context = sushi_test_project.context
    context.target = BigQueryConfig(
        name="test_target", schema="test_schema", database="test-project"
    )

    base_config = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        schema="sushi",
        partition_by={"field": "`ds`", "data_type": "datetime", "granularity": "month"},
        materialized=Materialization.INCREMENTAL,
        sql="SELECT 1 AS one FROM tbl_a",
    )

    assert base_config.to_sqlmesh(context).physical_properties == {}

    assert base_config.copy(
        update={"require_partition_filter": True, "partition_expiration_days": 7}
    ).to_sqlmesh(context).physical_properties == {
        "require_partition_filter": exp.convert(True),
        "partition_expiration_days": exp.convert(7),
    }

    assert base_config.copy(update={"require_partition_filter": True}).to_sqlmesh(
        context
    ).physical_properties == {
        "require_partition_filter": exp.convert(True),
    }

    assert base_config.copy(update={"partition_expiration_days": 7}).to_sqlmesh(
        context
    ).physical_properties == {
        "partition_expiration_days": exp.convert(7),
    }


@pytest.mark.xdist_group("dbt_manifest")
def test_clickhouse_properties(mocker: MockerFixture):
    context = DbtContext(target_name="production")
    context._project_name = "Foo"
    context._target = ClickhouseConfig(name="production")
    model_config = ModelConfig(
        name="model",
        alias="model",
        schema="test",
        package_name="package",
        materialized="incremental",
        incremental_strategy="delete+insert",
        incremental_predicates=["ds > (SELECT MAX(ds) FROM model)"],
        query_settings={"QUERY_SETTING": "value"},
        sharding_key="rand()",
        engine="MergeTree()",
        partition_by=["toMonday(ds)", "partition_col"],
        order_by=["toStartOfWeek(ds)", "order_col"],
        primary_key=["ds", "primary_key_col"],
        ttl="time + INTERVAL 1 WEEK",
        settings={"SETTING": "value"},
        sql="""SELECT 1 AS one, ds FROM foo""",
    )

    with patch.object(get_console(), "log_warning") as mock_logger:
        model_to_sqlmesh = model_config.to_sqlmesh(context)

    assert [call[0][0] for call in mock_logger.call_args_list] == [
        "The 'delete+insert' incremental strategy is not supported - SQLMesh will use the temp table/partition swap strategy.",
        "SQLMesh does not support 'incremental_predicates' - they will not be applied.",
        "SQLMesh does not support the 'query_settings' model configuration parameter. Specify the query settings directly in the model query.",
        "SQLMesh does not support the 'sharding_key' model configuration parameter or distributed materializations.",
    ]

    assert [e.sql("clickhouse") for e in model_to_sqlmesh.partitioned_by] == [
        'toMonday("ds")',
        '"partition_col"',
    ]
    assert model_to_sqlmesh.storage_format == "MergeTree()"

    physical_properties = model_to_sqlmesh.physical_properties
    assert [e.sql("clickhouse", identify=True) for e in physical_properties["order_by"]] == [
        'toStartOfWeek("ds")',
        '"order_col"',
    ]
    assert [e.sql("clickhouse", identify=True) for e in physical_properties["primary_key"]] == [
        '"ds"',
        '"primary_key_col"',
    ]
    assert physical_properties["ttl"].sql("clickhouse") == "time + INTERVAL 1 WEEK"
    assert physical_properties["SETTING"].sql("clickhouse") == "value"


@pytest.mark.xdist_group("dbt_manifest")
def test_snapshot_json_payload():
    sushi_context = Context(paths=["tests/fixtures/dbt/sushi_test"])
    snapshot_json = json.loads(
        _snapshot_to_json(sushi_context.get_snapshot("sushi.top_waiters", raise_if_missing=True))
    )
    assert snapshot_json["node"]["jinja_macros"]["global_objs"]["target"] == {
        "type": "duckdb",
        "name": "in_memory",
        "schema": "sushi",
        "database": "memory",
        "target_name": "in_memory",
    }


@pytest.mark.xdist_group("dbt_manifest")
@time_machine.travel("2023-01-08 00:00:00 UTC")
def test_dbt_package_macros(sushi_test_project: Project):
    context = sushi_test_project.context

    # Make sure external macros are available.
    assert context.render("{{ dbt.current_timestamp() }}") == "now()"
    # Make sure builtins are available too.
    assert context.render("{{ dbt.run_started_at }}") == "2023-01-08 00:00:00+00:00"


@pytest.mark.xdist_group("dbt_manifest")
def test_dbt_vars(sushi_test_project: Project):
    context = sushi_test_project.context
    context.set_and_render_variables(
        sushi_test_project.packages["customers"].variables, "customers"
    )

    assert context.render("{{ var('some_other_var') }}") == "5"
    assert context.render("{{ var('some_other_var', 0) }}") == "5"
    assert context.render("{{ var('missing') }}") == "None"
    assert context.render("{{ var('missing', 0) }}") == "0"

    assert context.render("{{ var.has_var('some_other_var') }}") == "True"
    assert context.render("{{ var.has_var('missing') }}") == "False"


@pytest.mark.xdist_group("dbt_manifest")
def test_snowflake_session_properties(sushi_test_project: Project, mocker: MockerFixture):
    context = sushi_test_project.context
    context.target = SnowflakeConfig(
        name="target", schema="test", database="test", account="foo", user="bar", password="baz"
    )

    base_config = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        schema="sushi",
        partition_by={"field": "`ds`", "data_type": "datetime", "granularity": "month"},
        materialized=Materialization.INCREMENTAL,
        sql="SELECT 1 AS one FROM tbl_a",
    )

    assert base_config.to_sqlmesh(context).session_properties == {}

    model_with_warehouse = base_config.copy(
        update={"snowflake_warehouse": "test_warehouse"}
    ).to_sqlmesh(context)

    assert model_with_warehouse.session_properties_ == exp.Tuple(
        expressions=[exp.Literal.string("warehouse").eq(exp.Literal.string("test_warehouse"))]
    )
    assert model_with_warehouse.session_properties == {"warehouse": "test_warehouse"}


def test_model_cluster_by():
    context = DbtContext()
    context._target = SnowflakeConfig(
        name="target",
        schema="test",
        database="test",
        account="account",
        user="user",
        password="password",
    )

    model = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        target_schema="test",
        cluster_by="Bar",
        sql="SELECT * FROM baz",
        materialized=Materialization.TABLE.value,
    )
    assert model.to_sqlmesh(context).clustered_by == [exp.to_column('"BAR"')]

    model = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        target_schema="test",
        cluster_by=["Bar", "qux"],
        sql="SELECT * FROM baz",
        materialized=Materialization.TABLE.value,
    )
    assert model.to_sqlmesh(context).clustered_by == [
        exp.to_column('"BAR"'),
        exp.to_column('"QUX"'),
    ]


def test_snowflake_dynamic_table():
    context = DbtContext()
    context._target = SnowflakeConfig(
        name="target",
        schema="test",
        database="test",
        account="account",
        user="user",
        password="password",
    )

    model = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        target_schema="test",
        sql="SELECT * FROM baz",
        materialized=Materialization.DYNAMIC_TABLE.value,
        target_lag="1 hour",
        snowflake_warehouse="SMALL",
    )

    as_sqlmesh = model.to_sqlmesh(context)
    assert as_sqlmesh.kind == ManagedKind()
    assert as_sqlmesh.physical_properties == {
        "target_lag": exp.Literal.string("1 hour"),
        "warehouse": exp.Literal.string("SMALL"),
    }

    # both target_lag and snowflake_warehouse are required properties
    # https://docs.getdbt.com/reference/resource-configs/snowflake-configs#dynamic-tables
    for required_property in ["target_lag", "snowflake_warehouse"]:
        with pytest.raises(ConfigError, match=r".*must be set for dynamic tables"):
            model.copy(update={required_property: None}).to_sqlmesh(context)


@pytest.mark.xdist_group("dbt_manifest")
def test_refs_in_jinja_globals(sushi_test_project: Project, mocker: MockerFixture):
    context = sushi_test_project.context

    sqlmesh_model = t.cast(
        SqlModel,
        sushi_test_project.packages["sushi"].models["simple_model_b"].to_sqlmesh(context),
    )
    assert set(sqlmesh_model.jinja_macros.global_objs["refs"].keys()) == {"simple_model_a"}  # type: ignore

    sqlmesh_model = t.cast(
        SqlModel,
        sushi_test_project.packages["sushi"].models["top_waiters"].to_sqlmesh(context),
    )
    assert set(sqlmesh_model.jinja_macros.global_objs["refs"].keys()) == {  # type: ignore
        "waiter_revenue_by_day",
        "sushi.waiter_revenue_by_day",
    }


def test_allow_partials_by_default():
    context = DbtContext()
    context._target = SnowflakeConfig(
        name="target",
        schema="test",
        database="test",
        account="account",
        user="user",
        password="password",
    )

    model = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        target_schema="test",
        sql="SELECT * FROM baz",
        materialized=Materialization.TABLE.value,
    )
    assert model.allow_partials is None
    assert model.to_sqlmesh(context).allow_partials

    model.materialized = Materialization.INCREMENTAL.value
    assert model.allow_partials is None
    assert model.to_sqlmesh(context).allow_partials

    model.allow_partials = True
    assert model.to_sqlmesh(context).allow_partials

    model.allow_partials = False
    assert not model.to_sqlmesh(context).allow_partials


def test_grain():
    context = DbtContext()
    context._target = SnowflakeConfig(
        name="target",
        schema="test",
        database="test",
        account="account",
        user="user",
        password="password",
    )

    model = ModelConfig(
        name="model",
        alias="model",
        package_name="package",
        target_schema="test",
        sql="SELECT * FROM baz",
        materialized=Materialization.TABLE.value,
        grain=["id_a", "id_b"],
    )
    assert model.to_sqlmesh(context).grains == [exp.to_column("id_a"), exp.to_column("id_b")]

    model.grain = "id_a"
    assert model.to_sqlmesh(context).grains == [exp.to_column("id_a")]


@pytest.mark.xdist_group("dbt_manifest")
def test_on_run_start_end():
    sushi_context = Context(paths=["tests/fixtures/dbt/sushi_test"])
    assert len(sushi_context._environment_statements) == 2

    # Root project's on run start / on run end should be first by checking the macros
    root_environment_statements = sushi_context._environment_statements[0]
    assert "create_tables" in root_environment_statements.jinja_macros.root_macros

    # Validate order of execution to be correct
    assert root_environment_statements.before_all == [
        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_last (col VARCHAR);\nJINJA_END;",
        """JINJA_STATEMENT_BEGIN;\nSELECT {{ var("yet_another_var") }} AS var, '{{ source("raw", "items").identifier }}' AS src, '{{ ref("waiters").identifier }}' AS ref;\nJINJA_END;""",
        "JINJA_STATEMENT_BEGIN;\n{{ log_value('on-run-start') }}\nJINJA_END;",
    ]
    assert root_environment_statements.after_all == [
        "JINJA_STATEMENT_BEGIN;\n{{ create_tables(schemas) }}\nJINJA_END;",
        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_last;\nJINJA_END;",
        "JINJA_STATEMENT_BEGIN;\n{{ graph_usage() }}\nJINJA_END;",
    ]

    assert root_environment_statements.jinja_macros.root_package_name == "sushi"

    rendered_before_all = render_statements(
        root_environment_statements.before_all,
        dialect=sushi_context.default_dialect,
        python_env=root_environment_statements.python_env,
        jinja_macros=root_environment_statements.jinja_macros,
        runtime_stage=RuntimeStage.BEFORE_ALL,
    )

    rendered_after_all = render_statements(
        root_environment_statements.after_all,
        dialect=sushi_context.default_dialect,
        python_env=root_environment_statements.python_env,
        jinja_macros=root_environment_statements.jinja_macros,
        snapshots=sushi_context.snapshots,
        runtime_stage=RuntimeStage.AFTER_ALL,
        environment_naming_info=EnvironmentNamingInfo(name="dev"),
        engine_adapter=sushi_context.engine_adapter,
    )

    assert rendered_before_all == [
        "CREATE TABLE IF NOT EXISTS analytic_stats (physical_table TEXT, evaluation_time TEXT)",
        "CREATE TABLE IF NOT EXISTS to_be_executed_last (col TEXT)",
        "SELECT 1 AS var, 'items' AS src, 'waiters' AS ref",
    ]

    # The jinja macro should have resolved the schemas for this environment and generated corresponding statements
    expected_statements = [
        "CREATE OR REPLACE TABLE schema_table_snapshots__dev AS SELECT 'snapshots__dev' AS schema",
        "CREATE OR REPLACE TABLE schema_table_sushi__dev AS SELECT 'sushi__dev' AS schema",
        "DROP TABLE to_be_executed_last",
    ]
    assert sorted(rendered_after_all[:-1]) == sorted(expected_statements)

    # Assert the models with their materialisations are present in the rendered graph_table statement
    graph_table_stmt = rendered_after_all[-1]
    assert "'model.sushi.simple_model_a' AS unique_id, 'table' AS materialized" in graph_table_stmt
    assert "'model.sushi.waiters' AS unique_id, 'ephemeral' AS materialized" in graph_table_stmt
    assert "'model.sushi.simple_model_b' AS unique_id, 'table' AS materialized" in graph_table_stmt
    assert (
        "'model.sushi.waiter_as_customer_by_day' AS unique_id, 'incremental' AS materialized"
        in graph_table_stmt
    )
    assert "'model.sushi.top_waiters' AS unique_id, 'view' AS materialized" in graph_table_stmt
    assert "'model.customers.customers' AS unique_id, 'view' AS materialized" in graph_table_stmt
    assert (
        "'model.customers.customer_revenue_by_day' AS unique_id, 'incremental' AS materialized"
        in graph_table_stmt
    )
    assert (
        "'model.sushi.waiter_revenue_by_day.v1' AS unique_id, 'incremental' AS materialized"
        in graph_table_stmt
    )
    assert (
        "'model.sushi.waiter_revenue_by_day.v2' AS unique_id, 'incremental' AS materialized"
        in graph_table_stmt
    )

    # Nested dbt_packages on run start / on run end
    packaged_environment_statements = sushi_context._environment_statements[1]

    # Validate order of execution to be correct
    assert packaged_environment_statements.before_all == [
        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS to_be_executed_first (col VARCHAR);\nJINJA_END;",
        "JINJA_STATEMENT_BEGIN;\nCREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table VARCHAR, evaluation_time VARCHAR);\nJINJA_END;",
    ]
    assert packaged_environment_statements.after_all == [
        "JINJA_STATEMENT_BEGIN;\nDROP TABLE to_be_executed_first\nJINJA_END;",
        "JINJA_STATEMENT_BEGIN;\n{{ packaged_tables(schemas) }}\nJINJA_END;",
    ]

    assert "packaged_tables" in packaged_environment_statements.jinja_macros.root_macros
    assert packaged_environment_statements.jinja_macros.root_package_name == "sushi"

    rendered_before_all = render_statements(
        packaged_environment_statements.before_all,
        dialect=sushi_context.default_dialect,
        python_env=packaged_environment_statements.python_env,
        jinja_macros=packaged_environment_statements.jinja_macros,
        runtime_stage=RuntimeStage.BEFORE_ALL,
    )

    rendered_after_all = render_statements(
        packaged_environment_statements.after_all,
        dialect=sushi_context.default_dialect,
        python_env=packaged_environment_statements.python_env,
        jinja_macros=packaged_environment_statements.jinja_macros,
        snapshots=sushi_context.snapshots,
        runtime_stage=RuntimeStage.AFTER_ALL,
        environment_naming_info=EnvironmentNamingInfo(name="dev"),
        engine_adapter=sushi_context.engine_adapter,
    )

    # Validate order of execution to match dbt's
    assert rendered_before_all == [
        "CREATE TABLE IF NOT EXISTS to_be_executed_first (col TEXT)",
        "CREATE TABLE IF NOT EXISTS analytic_stats_packaged_project (physical_table TEXT, evaluation_time TEXT)",
    ]

    # This on run end statement should be executed first
    assert rendered_after_all[0] == "DROP TABLE to_be_executed_first"

    # The table names is an indication of the rendering of the dbt_packages statements
    assert sorted(rendered_after_all) == sorted(
        [
            "DROP TABLE to_be_executed_first",
            "CREATE OR REPLACE TABLE schema_table_snapshots__dev_nested_package AS SELECT 'snapshots__dev' AS schema",
            "CREATE OR REPLACE TABLE schema_table_sushi__dev_nested_package AS SELECT 'sushi__dev' AS schema",
        ]
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_dynamic_var_names(sushi_test_project: Project, sushi_test_dbt_context: Context):
    context = sushi_test_project.context
    context.set_and_render_variables(sushi_test_project.packages["sushi"].variables, "sushi")
    context.target = BigQueryConfig(name="production", database="main", schema="sushi")
    model_config = ModelConfig(
        name="model",
        alias="model",
        schema="test",
        package_name="package",
        materialized="table",
        unique_key="ds",
        partition_by={"field": "ds", "granularity": "month"},
        sql="""
        {% set var_name = "yet_" + "another_" + "var" %}
        {% set results = run_query('select 1 as one') %}
        {% if results %}
        SELECT {{ results.columns[0].values()[0] }} AS one {{ var(var_name) }} AS var FROM {{ this.identifier }}
        {% else %}
        SELECT NULL AS one {{ var(var_name) }} AS var FROM {{ this.identifier }}
        {% endif %}
        """,
        dependencies=Dependencies(has_dynamic_var_names=True),
    )
    converted_model = model_config.to_sqlmesh(context)
    assert "yet_another_var" in converted_model.jinja_macros.global_objs["vars"]  # type: ignore

    # Test the existing model in the sushi project
    assert (
        "dynamic_test_var"  # type: ignore
        in sushi_test_dbt_context.get_model(
            "sushi.waiter_revenue_by_day_v2"
        ).jinja_macros.global_objs["vars"]
    )


@pytest.mark.xdist_group("dbt_manifest")
def test_dynamic_var_names_in_macro(sushi_test_project: Project):
    context = sushi_test_project.context
    context.set_and_render_variables(sushi_test_project.packages["sushi"].variables, "sushi")
    context.target = BigQueryConfig(name="production", database="main", schema="sushi")
    model_config = ModelConfig(
        name="model",
        alias="model",
        schema="test",
        package_name="package",
        materialized="table",
        unique_key="ds",
        partition_by={"field": "ds", "granularity": "month"},
        sql="""
        {% set var_name = "dynamic_" + "test_" + "var" %}
        SELECT {{ sushi.dynamic_var_name_dependency(var_name) }} AS var
        """,
        dependencies=Dependencies(
            macros=[MacroReference(package="sushi", name="dynamic_var_name_dependency")],
            has_dynamic_var_names=True,
        ),
    )
    converted_model = model_config.to_sqlmesh(context)
    assert "dynamic_test_var" in converted_model.jinja_macros.global_objs["vars"]  # type: ignore