from __future__ import annotations

import datetime
import logging
import os
import random
import re
import string
from copy import deepcopy
from textwrap import dedent
from typing import Optional

import pytest
from helpers.mock_soda_cloud import MockResponse, MockSodaCloud
from helpers.test_table import TestColumn, TestTable, TestTableSpecification
from soda_core.common.data_source_impl import DataSourceImpl
from soda_core.common.logs import Logs
from soda_core.common.metadata_types import SodaDataTypeName, SqlDataType
from soda_core.common.soda_cloud import SodaCloud
from soda_core.common.sql_ast import (
    COLUMN,
    CREATE_MATERIALIZED_VIEW,
    CREATE_TABLE,
    CREATE_TABLE_COLUMN,
    CREATE_VIEW,
    DROP_MATERIALIZED_VIEW,
    DROP_TABLE,
    DROP_VIEW,
    FROM,
    INSERT_INTO,
    SELECT,
    STAR,
    VALUES_ROW,
)
from soda_core.common.statements.metadata_tables_query import (
    FullyQualifiedTableName,
    FullyQualifiedViewName,
    MetadataTablesQuery,
    TableType,
)
from soda_core.common.statements.table_types import FullyQualifiedMaterializedViewName
from soda_core.common.yaml import (
    ContractYamlSource,
    DataSourceYamlSource,
    SodaCloudYamlSource,
)
from soda_core.contracts.contract_verification import (
    ContractVerificationResult,
    ContractVerificationSession,
    ContractVerificationSessionResult,
)

logger = logging.getLogger(__name__)


class DataSourceTestHelper:
    @classmethod
    def create(cls, test_datasource: str, name: str) -> DataSourceTestHelper:
        if test_datasource == "postgres":
            from soda_postgres.test_helpers.postgres_data_source_test_helper import (
                PostgresDataSourceTestHelper,
            )

            return PostgresDataSourceTestHelper(name)
        elif test_datasource == "snowflake":
            from soda_snowflake.test_helpers.snowflake_data_source_test_helper import (
                SnowflakeDataSourceTestHelper,
            )

            return SnowflakeDataSourceTestHelper(name)
        elif test_datasource == "databricks":
            from soda_databricks.test_helpers.databricks_data_source_test_helper import (
                DatabricksDataSourceTestHelper,
            )

            return DatabricksDataSourceTestHelper(name)
        elif test_datasource == "duckdb":
            from soda_duckdb.test_helpers.duckdb_data_source_test_helper import (
                DuckdbDataSourceTestHelper,
            )

            return DuckdbDataSourceTestHelper(name)
        elif test_datasource == "bigquery":
            from soda_bigquery.test_helpers.bigquery_data_source_test_helper import (
                BigQueryDataSourceTestHelper,
            )

            return BigQueryDataSourceTestHelper(name)

        elif test_datasource == "oracle":
            from soda_oracle.test_helpers.oracle_data_source_test_helper import (
                OracleDataSourceTestHelper,
            )

            return OracleDataSourceTestHelper(name)

        elif test_datasource == "sqlserver":
            from soda_sqlserver.test_helpers.sqlserver_data_source_test_helper import (
                SqlServerDataSourceTestHelper,
            )

            return SqlServerDataSourceTestHelper(name)
        elif test_datasource == "synapse":
            from soda_synapse.test_helpers.synapse_data_source_test_helper import (
                SynapseDataSourceTestHelper,
            )

            return SynapseDataSourceTestHelper(name)
        elif test_datasource == "redshift":
            from soda_redshift.test_helpers.redshift_data_source_test_helper import (
                RedshiftDataSourceTestHelper,
            )

            return RedshiftDataSourceTestHelper(name)
        elif test_datasource == "fabric":
            from soda_fabric.test_helpers.fabric_data_source_test_helper import (
                FabricDataSourceTestHelper,
            )

            return FabricDataSourceTestHelper(name)
        elif test_datasource == "athena":
            from soda_athena.test_helpers.athena_data_source_test_helper import (
                AthenaDataSourceTestHelper,
            )

            return AthenaDataSourceTestHelper(name)
        elif test_datasource == "sparkdf":
            from soda_sparkdf.test_helpers.sparkdf_data_source_test_helper import (
                SparkDataFrameDataSourceTestHelper,
            )

            return SparkDataFrameDataSourceTestHelper(name)

        elif test_datasource == "dremio":
            from soda_dremio.test_helpers.dremio_data_source_test_helper import (
                DremioDataSourceTestHelper,
            )

            return DremioDataSourceTestHelper(name)
        elif test_datasource == "trino":
            from soda_trino.test_helpers.trino_data_source_test_helper import (
                TrinoDataSourceTestHelper,
            )

            return TrinoDataSourceTestHelper(name)
        else:
            raise AssertionError(f"Unknown test data source {test_datasource}")

    def __init__(self, name: str):
        self.name = name
        self.dataset_prefix: list[str] = self._create_dataset_prefix()
        logs: Logs = Logs()
        self.data_source_impl: "DataSourceImpl" = self._create_data_source_impl()
        logs.remove_from_root_logger()
        if logs.has_errors:
            raise RuntimeError(f"Couldn't create DataSource: {logs}")
        self.is_cicd = os.getenv("GITHUB_ACTIONS") is not None

        # Test table names that are present in the data source.
        # None means the data source is not queried
        self.existing_test_table_names: Optional[list[str]] = None

        # Maps TestTable to their unique_name property
        # (that is the the full table name composed of "SODATEST_" prefix, table purpose & test table hash)
        self.test_tables: dict[str, TestTable] = {}

        self.soda_cloud: Optional[SodaCloud] = None
        self.use_agent: bool = False

        if os.environ.get("SEND_RESULTS_TO_SODA_CLOUD") == "on":
            self.enable_soda_cloud()

    def enable_soda_cloud(self):
        logs: Logs = Logs()
        soda_cloud_yaml_str: str = """
            soda_cloud:
              host: ${env.SODA_CLOUD_HOST}
              api_key_id: ${env.SODA_CLOUD_API_KEY_ID}
              api_key_secret: ${env.SODA_CLOUD_API_KEY_SECRET}
        """
        soda_cloud_yaml_source: SodaCloudYamlSource = SodaCloudYamlSource.from_str(yaml_str=soda_cloud_yaml_str)
        self.soda_cloud = SodaCloud.from_yaml_source(
            soda_cloud_yaml_source=soda_cloud_yaml_source, provided_variable_values={}
        )
        if logs.has_errors:
            raise AssertionError(str(logs))

    def enable_soda_cloud_mock(self, responses: Optional[list[MockResponse]] = None):
        self.soda_cloud = MockSodaCloud(responses)

    def _create_data_source_impl(self) -> "DataSourceImpl":
        """
        Called in constructor to initialized self.data_source
        """
        logs: Logs = Logs()
        data_source_yaml_source: DataSourceYamlSource = self._create_data_source_yaml_source()
        from soda_core.common.data_source_impl import DataSourceImpl

        data_source_impl: DataSourceImpl = DataSourceImpl.from_yaml_source(data_source_yaml_source)
        assert not logs.has_errors
        return data_source_impl

    def _create_data_source_yaml_str(self) -> str:
        """
        Called in _create_data_source_impl to initialized self.data_source_impl
        self.database_name and self.schema_name are available if appropriate for the data source type
        """
        return ""

    def _create_data_source_yaml_source(self) -> DataSourceYamlSource:
        test_data_source_yaml_str: str = self._create_data_source_yaml_str()
        test_data_source_yaml_str = dedent(test_data_source_yaml_str).strip()
        return DataSourceYamlSource.from_str(yaml_str=test_data_source_yaml_str)

    def _create_dataset_prefix(self) -> list[str]:
        database_name: str = self._create_database_name()
        schema_name: str = self._create_schema_name()
        return [database_name, schema_name]

    def _create_database_name(self) -> Optional[str]:
        """
        Called in constructor to initialized self.database_name
        """
        return "soda_test"

    def _create_schema_name(self) -> Optional[str]:
        """
        Called in constructor to initialized self.schema_name
        """

        schema_name_parts = []

        github_ref_name = os.getenv("GITHUB_REF_NAME")
        github_head_ref = os.getenv("GITHUB_HEAD_REF")

        if not github_ref_name and not github_head_ref:
            user = os.getenv("USER", "anonymous")
            schema_name_parts.append("dev")
            schema_name_parts.append(user)

        else:
            python_version = os.getenv("PYTHON_VERSION")
            python_version_short = f'P{python_version.replace(".", "")}' if python_version else ""
            timestamp = datetime.datetime.now().strftime("%Y_%m_%dT%H_%M_%S_%f")

            def generate_random_alpha_num_str(length: int) -> str:
                return "".join(random.choice(string.ascii_lowercase + string.digits) for _ in range(length))

            if github_head_ref:
                github_head_ref_short = (
                    github_head_ref[:15] if github_head_ref and len(github_head_ref) > 15 else github_head_ref
                )
                schema_name_parts.append("ci")
                schema_name_parts.append(github_head_ref_short)
                schema_name_parts.append(python_version_short)
                schema_name_parts.append(timestamp)

            else:
                schema_name_parts.append("ci_main")
                schema_name_parts.append(python_version_short)
                schema_name_parts.append(timestamp)

        schema_name_raw = "_".join(schema_name_parts)
        schema_name = re.sub("[^0-9a-zA-Z]+", "_", schema_name_raw).lower()
        return schema_name

    def _adjust_schema_name(self, schema_name: str) -> str:
        return schema_name

    def start_test_session(self) -> None:
        self.start_test_session_open_connection()
        self.start_test_session_ensure_schema()

    def start_test_session_open_connection(self) -> None:
        logs: Logs = Logs()
        self.data_source_impl.open_connection()
        logs.remove_from_root_logger()
        if logs.has_errors:
            raise AssertionError(f"Connection creation has errors. See logs.")

    def start_test_session_ensure_schema(self) -> None:
        if self.is_cicd:
            self.drop_test_schema_if_exists()
        self.create_test_schema_if_not_exists()

    def end_test_session(self, exception: Optional[Exception]) -> None:
        self.end_test_session_drop_schema()
        self.end_test_session_close_connection()

    def end_test_session_close_connection(self) -> None:
        self.data_source_impl.close_connection()

    def end_test_session_drop_schema(self) -> None:
        if self.is_cicd:
            self.drop_test_schema_if_exists()

    def query_existing_test_tables(self) -> list[FullyQualifiedTableName]:
        database: Optional[str] = self.extract_database_from_prefix()
        schema: Optional[str] = self.extract_schema_from_prefix()

        metadata_tables_query: MetadataTablesQuery = self.data_source_impl.create_metadata_tables_query()
        fully_qualified_table_names: list[FullyQualifiedTableName] = metadata_tables_query.execute(
            database_name=database,
            schema_name=schema,
            include_table_name_like_filters=["SODATEST_%"],
        )
        return fully_qualified_table_names

    def query_existing_test_table_names(self) -> list[str]:
        fully_qualified_table_names = self.query_existing_test_tables()
        return [
            fully_qualified_test_table_name.table_name
            for fully_qualified_test_table_name in fully_qualified_table_names
            if isinstance(fully_qualified_test_table_name, FullyQualifiedTableName)
        ]

    def create_test_schema_if_not_exists(self) -> None:
        sql: str = self.create_test_schema_if_not_exists_sql()
        self.data_source_impl.execute_update(sql)
        post_create_sql = self.post_test_schema_create_sql()
        if post_create_sql:
            [self.data_source_impl.execute_update(cmd) for cmd in post_create_sql]

    def create_test_schema_if_not_exists_sql(self) -> str:
        return self.data_source_impl.sql_dialect.create_schema_if_not_exists_sql(self.dataset_prefix)

    def post_test_schema_create_sql(self) -> str:
        return self.data_source_impl.sql_dialect.post_schema_create_sql(self.dataset_prefix)

    def extract_database_from_prefix(self) -> Optional[str]:
        return self.data_source_impl.extract_database_from_prefix(self.dataset_prefix)

    def extract_schema_from_prefix(self) -> Optional[str]:
        return self.data_source_impl.extract_schema_from_prefix(self.dataset_prefix)

    def drop_test_schema_if_exists(self) -> None:
        schema = self.extract_schema_from_prefix()
        if not schema:
            raise AssertionError("Data source does not support schemas")

        self.drop_schema_if_exists(schema)

    def drop_schema_if_exists(self, schema: str) -> None:
        try:
            sql: str = self.drop_schema_if_exists_sql(schema=schema)
            self.data_source_impl.execute_update(sql)
        except Exception as e:
            logger.warning(f"Error dropping test schema: {e}")

    def drop_schema_if_exists_sql(self, schema: str) -> str:
        return f"DROP SCHEMA IF EXISTS {self.data_source_impl.sql_dialect.quote_default(schema)} CASCADE;"

    def drop_schema_sql(self, schema: str) -> str:
        return f"DROP SCHEMA {self.data_source_impl.sql_dialect.quote_default(schema)} CASCADE;"

    def ensure_test_table(
        self, test_table_specification: TestTableSpecification, force_recreate: bool = False
    ) -> TestTable:
        """
        Returns a test table with the given table data
        """
        if self.existing_test_table_names is None:
            self.existing_test_table_names = self.query_existing_test_table_names()

        # Specifically for BigQuery; when developing locally, the metadata might lag behind from the actual state of the tables.
        # Uncomment this to force the test table to be recreated every time.
        # self.existing_test_table_names = []

        test_table: TestTable = self.test_tables.get(test_table_specification.unique_name)
        if not test_table:
            test_table = self._create_test_table_python_object(test_table_specification)

            existing_test_table_names_lower: list[str] = [
                existing_test_table_name.lower() for existing_test_table_name in self.existing_test_table_names
            ]
            if (
                test_table_specification.unique_name.lower() not in existing_test_table_names_lower
                or not self.verify_test_table_row_count(test_table_specification)
                or force_recreate
            ):
                obsolete_table_names = [
                    existing_test_table
                    for existing_test_table in self.existing_test_table_names
                    if existing_test_table.lower().startswith(f"sodatest_{test_table_specification.name.lower()}_")
                ]
                if obsolete_table_names:
                    for obsolete_table_name in obsolete_table_names:
                        logger.debug(f"Test table {obsolete_table_name} has changed and will be recreated")
                        self._drop_test_table(table_name=obsolete_table_name)
                        self.existing_test_table_names.remove(obsolete_table_name)

                logger.debug(f"Test table {test_table_specification.unique_name} will be created")
                self._create_and_insert_test_table(test_table=test_table)
                self.existing_test_table_names.append(test_table.unique_name)

                self.data_source_impl.data_source_connection.commit()
        else:
            logger.debug(f"Test table {test_table.unique_name} already exists")

        return test_table

    def verify_test_table_row_count(self, test_table_specification: TestTableSpecification) -> bool:
        expected_row_values = test_table_specification.row_values
        if expected_row_values is None:
            expected_row_values = (
                []
            )  # This is a table that is not expected to have any rows. We should check that as well!
        row_count_sql = f"SELECT COUNT(*) FROM {self.data_source_impl.sql_dialect.qualify_dataset_name(self.dataset_prefix, test_table_specification.unique_name)}"
        row_count = self.data_source_impl.execute_query(row_count_sql)
        try:
            row_count_int = int(row_count.rows[0][0])
        except ValueError:
            row_count_int = None

        if row_count_int is not None and row_count_int != len(expected_row_values):
            logger.warning(
                f"Test table {test_table_specification.unique_name} has {row_count_int} rows, expected {len(expected_row_values)}"
            )
            logger.warning(f"Attempting to drop and recreate table {test_table_specification.unique_name}")
            return False
        return True

    def _create_test_table_python_object(self, test_table_specification: TestTableSpecification) -> TestTable:
        columns: list[TestColumn] = []
        for test_column_specification in test_table_specification.columns:
            data_source_sql_data_type: SqlDataType = (
                self.data_source_impl.sql_dialect.map_test_sql_data_type_to_data_source(
                    source_data_type=test_column_specification.sql_data_type
                )
            )
            test_column: TestColumn = TestColumn(
                name=test_column_specification.column_name,
                sql_data_type=data_source_sql_data_type,
            )
            columns.append(test_column)

        sql_dialect = self.data_source_impl.sql_dialect

        return TestTable(
            data_source_name=self.data_source_impl.name,
            dataset_prefix=self.dataset_prefix,
            code_name=test_table_specification.unique_name,
            unique_name=test_table_specification.unique_name,
            qualified_name=sql_dialect.qualify_dataset_name(
                dataset_prefix=self.dataset_prefix,
                dataset_name=test_table_specification.unique_name,
            ),
            columns=columns,
            row_values=test_table_specification.row_values,
            table_type=test_table_specification.table_type,
        )

    def _create_and_insert_test_table(self, test_table: TestTable) -> None:
        self._create_test_table(test_table)
        self._insert_test_table_rows(test_table)

    def _create_test_table(self, test_table: TestTable) -> None:
        my_create_table = CREATE_TABLE(
            fully_qualified_table_name=test_table.qualified_name,
            columns=[
                CREATE_TABLE_COLUMN(name=column.name, type=column.sql_data_type)
                for column in test_table.columns.values()
            ],
        )
        sql: str = self.data_source_impl.sql_dialect.build_create_table_sql(my_create_table)
        self.data_source_impl.execute_update(sql)

    def _insert_test_table_rows(self, test_table: TestTable) -> None:
        sql: str = self._insert_test_table_rows_sql(test_table)
        if sql:
            self.data_source_impl.execute_update(sql)

    def _insert_test_table_rows_sql(self, test_table: TestTable) -> str:
        if test_table.row_values:
            insert_into_sql = self.data_source_impl.sql_dialect.build_insert_into_sql(
                INSERT_INTO(
                    fully_qualified_table_name=test_table.qualified_name,
                    values=[VALUES_ROW(row) for row in test_table.row_values],
                    columns=[COLUMN(column.name) for column in test_table.columns.values()],
                )
            )
            return insert_into_sql

    def _drop_test_table(self, table_name: str) -> None:
        # Note, this is not a fully qualified table name, it's just the table name.
        # So we need to qualify it ourselves.
        fully_qualified_table_name = self.data_source_impl.sql_dialect.qualify_dataset_name(
            dataset_prefix=self.dataset_prefix,
            dataset_name=table_name,
        )
        my_drop_table = DROP_TABLE(
            fully_qualified_table_name=fully_qualified_table_name, cascade=self._cascade_drop_table()
        )
        sql: str = self.data_source_impl.sql_dialect.build_drop_table_sql(my_drop_table)
        self.data_source_impl.execute_update(sql)

    def _cascade_drop_table(self) -> bool:
        return True

    def query_existing_test_views(self) -> list[FullyQualifiedViewName]:
        metadata_tables_query: MetadataTablesQuery = self.data_source_impl.create_metadata_tables_query()
        fully_qualified_view_names: list[FullyQualifiedViewName] = metadata_tables_query.execute(
            database_name=self.extract_database_from_prefix(),
            schema_name=self.extract_schema_from_prefix(),
            include_table_name_like_filters=["SODATEST_%"],
            types_to_return=[TableType.VIEW],
        )
        return fully_qualified_view_names

    def query_existing_materialized_test_views(self) -> list[FullyQualifiedMaterializedViewName]:
        metadata_tables_query: MetadataTablesQuery = self.data_source_impl.create_metadata_tables_query()
        fully_qualified_view_names: list[FullyQualifiedMaterializedViewName] = metadata_tables_query.execute(
            database_name=self.extract_database_from_prefix(),
            schema_name=self.extract_schema_from_prefix(),
            include_table_name_like_filters=["SODATEST_%"],
            types_to_return=[TableType.MATERIALIZED_VIEW],
        )
        return fully_qualified_view_names

    def query_existing_test_view_names(self) -> list[str]:
        fully_qualified_view_names = self.query_existing_test_views()
        return [
            fully_qualified_view_name.view_name
            for fully_qualified_view_name in fully_qualified_view_names
            if isinstance(fully_qualified_view_name, FullyQualifiedViewName)
        ]

    def query_existing_test_materialized_view_names(self) -> list[str]:
        fully_qualified_view_names = self.query_existing_materialized_test_views()
        return [
            fully_qualified_view_name.materialized_view_name
            for fully_qualified_view_name in fully_qualified_view_names
            if isinstance(fully_qualified_view_name, FullyQualifiedMaterializedViewName)
        ]

    def create_view_from_test_table(self, test_table: TestTable) -> TestTable:
        # First gate: skip early if the dialect declares no view support (e.g. Dremio)
        if not self.data_source_impl.can_create_view:
            pytest.skip("View creation is not supported by this data source")

        view_name = f"{test_table.unique_name}_view"
        view_name = self.data_source_impl.sql_dialect.metadata_casify(view_name)
        existing_view_names = self.query_existing_test_view_names()
        if view_name in existing_view_names:
            self._drop_test_view(view_name)
        my_create_view = CREATE_VIEW(
            fully_qualified_view_name=self.data_source_impl.sql_dialect.qualify_dataset_name(
                dataset_prefix=self.dataset_prefix,
                dataset_name=view_name,
            ),
            select_elements=[
                SELECT(STAR()),
                FROM(self.data_source_impl.sql_dialect.get_from_name_from_qualified_name(test_table.qualified_name)),
            ],
        )
        sql: str = self.data_source_impl.sql_dialect.build_create_view_sql(my_create_view)
        # Second gate: actually attempt the DDL — if it fails, the result is cached so
        # future tests skip immediately without retrying (see DataSourceImpl.try_create_view)
        if not self.data_source_impl.try_create_view(sql):
            pytest.skip("View creation is not supported by this data source")

        view_object = deepcopy(test_table)
        view_object.unique_name = view_name
        view_object.qualified_name = self.data_source_impl.sql_dialect.qualify_dataset_name(
            dataset_prefix=self.dataset_prefix,
            dataset_name=view_name,
        )

        return view_object

    def _drop_test_view(self, view_name: str) -> None:
        fully_qualified_view_name = self.data_source_impl.sql_dialect.qualify_dataset_name(
            dataset_prefix=self.dataset_prefix,
            dataset_name=view_name,
        )
        my_drop_view = DROP_VIEW(
            fully_qualified_view_name=fully_qualified_view_name,
        )
        sql: str = self.data_source_impl.sql_dialect.build_drop_view_sql(my_drop_view)
        self.data_source_impl.execute_update(sql)

    def create_materialized_view_from_test_table(self, test_table: TestTable) -> TestTable:
        # Same two-gate pattern as create_view_from_test_table above.
        # Most data sources default to supports_materialized_views()=False (see SqlDialect);
        # only those that explicitly opt in (Postgres, Redshift, Trino) reach the DDL attempt.
        if not self.data_source_impl.can_create_materialized_view:
            pytest.skip("Materialized view creation is not supported by this data source")

        view_name = f"{test_table.unique_name}_mv"
        view_name = self.data_source_impl.sql_dialect.metadata_casify(view_name)
        existing_view_names = self.query_existing_test_materialized_view_names()
        if view_name in existing_view_names:
            self._drop_test_materialized_view(view_name)
        my_create_view = CREATE_MATERIALIZED_VIEW(
            fully_qualified_view_name=self.data_source_impl.sql_dialect.qualify_dataset_name(
                dataset_prefix=self.dataset_prefix,
                dataset_name=view_name,
            ),
            select_elements=[
                SELECT(STAR()),
                FROM(self.data_source_impl.sql_dialect.get_from_name_from_qualified_name(test_table.qualified_name)),
            ],
        )
        sql: str = self.data_source_impl.sql_dialect.build_create_materialized_view_sql(my_create_view)
        if not self.data_source_impl.try_create_materialized_view(sql):
            pytest.skip("Materialized view creation is not supported by this data source")

        view_object = deepcopy(test_table)
        view_object.unique_name = view_name
        view_object.qualified_name = self.data_source_impl.sql_dialect.qualify_dataset_name(
            dataset_prefix=self.dataset_prefix,
            dataset_name=view_name,
        )
        return view_object

    def _drop_test_materialized_view(self, view_name: str) -> None:
        fully_qualified_view_name = self.data_source_impl.sql_dialect.qualify_dataset_name(
            dataset_prefix=self.dataset_prefix,
            dataset_name=view_name,
        )
        my_drop_view = DROP_MATERIALIZED_VIEW(
            fully_qualified_view_name=fully_qualified_view_name,
        )
        sql: str = self.data_source_impl.sql_dialect.build_drop_materialized_view_sql(my_drop_view)
        self.data_source_impl.execute_update(sql)

    def get_parse_errors_str(self, contract_yaml_str: str) -> str:
        contract_yaml_str: str = dedent(contract_yaml_str).strip()
        contract_verification_session_result = ContractVerificationSession.execute(
            contract_yaml_sources=[ContractYamlSource.from_str(contract_yaml_str)], only_validate_without_execute=True
        )
        return contract_verification_session_result.get_errors_str()

    def assert_contract_error(self, contract_yaml_str: str, variables: Optional[dict[str, str]] = None) -> str:
        contract_yaml_str: str = dedent(contract_yaml_str).strip()

        contract_verification_session_result: ContractVerificationSessionResult = ContractVerificationSession.execute(
            contract_yaml_sources=[
                ContractYamlSource.from_str(yaml_str=contract_yaml_str, file_path="yaml_string.yml")
            ],
            only_validate_without_execute=True,
            variables=variables,
            data_source_impls=[self.data_source_impl],
            soda_cloud_impl=self.soda_cloud,
            soda_cloud_use_agent=self.use_agent,
            soda_cloud_publish_results=True,
        )

        errors_str: str = contract_verification_session_result.get_errors_str()
        if not errors_str:
            raise AssertionError(f"Expected contract execution errors, but got none")
        return errors_str

    def assert_contract_pass(
        self,
        test_table: TestTable,
        contract_yaml_str: str,
        variables: Optional[dict[str, str]] = None,
        dwh_data_source_file_path: Optional[str] = None,
        extra_data_source_impls: list[DataSourceImpl] = [],
        check_paths: Optional[list[str]] = None,
        publish_results: bool = True,
    ) -> ContractVerificationResult:
        contract_verification_session_result: ContractVerificationSessionResult = self.verify_contract(
            contract_yaml_str=contract_yaml_str,
            test_table=test_table,
            variables=variables,
            dwh_data_source_file_path=dwh_data_source_file_path,
            extra_data_source_impls=extra_data_source_impls,
            check_paths=check_paths,
            publish_results=publish_results,
        )
        if not isinstance(contract_verification_session_result, ContractVerificationSessionResult):
            raise AssertionError(f"No contract verification result session")
        if not contract_verification_session_result.is_ok:
            raise AssertionError(f"Expected contract verification passed")
        if len(contract_verification_session_result.contract_verification_results) == 0:
            raise AssertionError(f"No contract verification results")

        # Final assert, is_ok does not guarantee all checks passed (could be skipped or non evaluated)
        assert contract_verification_session_result.is_passed, "Contract verification session result should be passed"

        return contract_verification_session_result.contract_verification_results[0]

    def assert_contract_ok(
        self,
        test_table: TestTable,
        contract_yaml_str: str,
        variables: Optional[dict[str, str]] = None,
        dwh_data_source_file_path: Optional[str] = None,
        extra_data_source_impls: list[DataSourceImpl] = [],
        check_paths: Optional[list[str]] = None,
        publish_results: bool = True,
    ) -> ContractVerificationResult:
        contract_verification_session_result: ContractVerificationSessionResult = self.verify_contract(
            contract_yaml_str=contract_yaml_str,
            test_table=test_table,
            variables=variables,
            dwh_data_source_file_path=dwh_data_source_file_path,
            extra_data_source_impls=extra_data_source_impls,
            check_paths=check_paths,
            publish_results=publish_results,
        )
        if not isinstance(contract_verification_session_result, ContractVerificationSessionResult):
            raise AssertionError(f"No contract verification result session")
        if not contract_verification_session_result.is_ok:
            raise AssertionError(f"Expected contract verification passed")
        if len(contract_verification_session_result.contract_verification_results) == 0:
            raise AssertionError(f"No contract verification results")

        return contract_verification_session_result.contract_verification_results[0]

    def assert_contract_fail(
        self,
        test_table: TestTable,
        contract_yaml_str: str,
        variables: Optional[dict[str, str]] = None,
        dwh_data_source_file_path: Optional[str] = None,
        extra_data_source_impls: list[DataSourceImpl] = [],
        publish_results: bool = True,
    ) -> ContractVerificationResult:
        contract_verification_session_result: ContractVerificationSessionResult = self.verify_contract(
            contract_yaml_str=contract_yaml_str,
            test_table=test_table,
            variables=variables,
            dwh_data_source_file_path=dwh_data_source_file_path,
            extra_data_source_impls=extra_data_source_impls,
            publish_results=publish_results,
        )
        if contract_verification_session_result.is_ok:
            raise AssertionError(f"Expected contract verification failed")
        return contract_verification_session_result.contract_verification_results[0]

    def assert_contract_warn(
        self,
        test_table: TestTable,
        contract_yaml_str: str,
        variables: Optional[dict[str, str]] = None,
        dwh_data_source_file_path: Optional[str] = None,
        extra_data_source_impls: list[DataSourceImpl] = [],
        publish_results: bool = True,
    ) -> ContractVerificationResult:
        contract_verification_session_result: ContractVerificationSessionResult = self.verify_contract(
            contract_yaml_str=contract_yaml_str,
            test_table=test_table,
            variables=variables,
            dwh_data_source_file_path=dwh_data_source_file_path,
            extra_data_source_impls=extra_data_source_impls,
            publish_results=publish_results,
        )
        if not contract_verification_session_result.is_warned:
            raise AssertionError(f"Expected contract verification warned")
        return contract_verification_session_result.contract_verification_results[0]

    def verify_contract(
        self,
        contract_yaml_str: str,
        test_table: Optional[TestTable] = None,
        variables: Optional[dict] = None,
        dwh_data_source_file_path: Optional[str] = None,
        extra_data_source_impls: list[DataSourceImpl] = [],
        check_paths: Optional[list[str]] = None,
        publish_results: bool = True,
    ) -> ContractVerificationSessionResult:
        contract_yaml_str = self._dedent_strip_and_prepend_dataset(contract_yaml_str, test_table)
        logger.debug(f"Contract:\n{contract_yaml_str}")

        return ContractVerificationSession.execute(
            contract_yaml_sources=[ContractYamlSource.from_str(contract_yaml_str)],
            variables=variables,
            data_source_impls=[self.data_source_impl, *extra_data_source_impls],
            soda_cloud_impl=self.soda_cloud,
            soda_cloud_use_agent=self.use_agent,
            soda_cloud_publish_results=publish_results,
            dwh_data_source_file_path=dwh_data_source_file_path,
            check_paths=check_paths,
        )

    def _dedent_strip_and_prepend_dataset(self, contract_yaml_str: str, test_table: Optional[TestTable]):
        checks_contract_yaml_str = dedent(contract_yaml_str).strip()
        if test_table:
            header_contract_yaml_str: str = f"dataset: {self.build_dqn(test_table)}\n"
            # This asserts that any "columns" statement in a check is single line.
            columns_contract_yaml_str: str = "columns: []\n" if not "columns:\n" in checks_contract_yaml_str else ""
            checks_contract_yaml_str = header_contract_yaml_str + columns_contract_yaml_str + checks_contract_yaml_str
        return checks_contract_yaml_str

    def build_dqn(self, test_table: TestTable) -> str:
        dqn_parts: list[str] = [self.data_source_impl.name] + self.dataset_prefix + [test_table.unique_name]
        dqn: str = "/".join(dqn_parts)
        return dqn

    def test_method_ended(self) -> None:
        # self.data_source_impl.data_source_connection.rollback() #TODO: this was originally done to theoretically speed up tests, but needs some datasource-specific work.
        self.soda_cloud = None
        self.use_agent = False

    def quote_column(self, column_name: str) -> str:
        """For shorter notation in the tests, we can just point it to the dialect."""
        return self.data_source_impl.sql_dialect.quote_column(column_name)

    def get_qualified_name_from_test_table(self, test_table: TestTable) -> str:
        return self.data_source_impl.sql_dialect.qualify_dataset_name(
            dataset_prefix=self.dataset_prefix,
            dataset_name=test_table.unique_name,
        )

    @classmethod
    def create_test_table_with_all_datatypes(cls, table_purpose: str, number_of_rows: int) -> TestTableSpecification:
        """
        Creates a test table with all data types.
        The smallint_default column will have a unique value for each row.
        """
        test_rows = [
            (
                # Chars
                "a",
                "abc",
                # Varchars
                "a",
                "abc",
                # Texts
                "a",
                # Integers
                i,
                100,
                100000,
                # Decimals
                1.0,
                100.0,
                100000.1234567890,
                # Numerics
                1.0,
                100.0,
                100000.1234567890,
                # Floats
                1.1234,
                # Doubles
                1.1234567890,
                # Timestamps
                datetime.datetime(2021, 1, 1, 1, 1, 1),
                datetime.datetime(2022, 1, 1, 1, 1, 1, 12),
                # Timestamp Zones
                datetime.datetime(2023, 1, 1, 1, 1, 1),
                datetime.datetime(2024, 1, 1, 1, 1, 1, 12),
                # Dates
                datetime.date(2025, 1, 1),
                # Times
                datetime.time(1, 2, 3),
                # Booleans
                True,
            )
            for i in range(number_of_rows)
        ]
        test_table_specification = (
            TestTableSpecification.builder()
            .table_purpose(table_purpose)
            .column_char("char_default")
            .column_char("char_w_length", 100)
            .column_varchar("varchar_default")
            .column_varchar("varchar_w_length", 100)
            .column_text("text_default")
            .column_smallint("smallint_default")
            .column_integer("integer_default")
            .column_bigint("bigint_default")
            .column_decimal("decimal_default")
            .column_decimal("decimal_w_precision", 10)
            .column_decimal("decimal_w_precision_and_scale", 10, 2)
            .column_numeric("numeric_default")
            .column_numeric("numeric_w_precision", 10)
            .column_numeric("numeric_w_precision_and_scale", 10, 2)
            .column_float("float_default")
            .column_double("double_default")
            .column_timestamp("timestamp_default")
            .column_timestamp("timestamp_w_precision", 2)
            .column_timestamp_tz("timestamp_tz_default")
            .column_timestamp_tz("timestamp_tz_w_precision", 4)
            .column_date("date_default")
            .column_time("time_default")
            .column_boolean("boolean_default")
            .rows(rows=test_rows)
            .build()
        )
        return test_table_specification

    def get_column_mappings(self) -> dict[str, SodaDataTypeName]:
        return {
            "char_default": SodaDataTypeName.CHAR,
            "char_w_length": SodaDataTypeName.CHAR,
            "varchar_default": SodaDataTypeName.VARCHAR,
            "varchar_w_length": SodaDataTypeName.VARCHAR,
            "text_default": SodaDataTypeName.TEXT,
            "smallint_default": SodaDataTypeName.SMALLINT,
            "integer_default": SodaDataTypeName.INTEGER,
            "bigint_default": SodaDataTypeName.BIGINT,
            "decimal_default": SodaDataTypeName.DECIMAL,
            "decimal_w_precision": SodaDataTypeName.DECIMAL,
            "decimal_w_precision_and_scale": SodaDataTypeName.DECIMAL,
            "numeric_default": SodaDataTypeName.NUMERIC,
            "numeric_w_precision": SodaDataTypeName.NUMERIC,
            "numeric_w_precision_and_scale": SodaDataTypeName.NUMERIC,
            "float_default": SodaDataTypeName.FLOAT,
            "double_default": SodaDataTypeName.DOUBLE,
            "timestamp_default": SodaDataTypeName.TIMESTAMP,
            "timestamp_w_precision": SodaDataTypeName.TIMESTAMP,
            "timestamp_tz_default": SodaDataTypeName.TIMESTAMP_TZ,
            "timestamp_tz_w_precision": SodaDataTypeName.TIMESTAMP_TZ,
            "date_default": SodaDataTypeName.DATE,
            "time_default": SodaDataTypeName.TIME,
            "boolean_default": SodaDataTypeName.BOOLEAN,
        }

    def map_table_type_to_short_string(self, table_type: TableType) -> str:
        return {
            TableType.TABLE: "t",
            TableType.MATERIALIZED_VIEW: "mv",
            TableType.VIEW: "v",
        }[table_type]
