"""pymssql data dictionary."""
from typing import TYPE_CHECKING, Any, ClassVar, cast
from mypy_extensions import mypyc_attr
from sqlspec.data_dictionary import (
ColumnMetadata,
DDLResult,
ForeignKeyMetadata,
IndexMetadata,
MetadataSupport,
SystemMetadataCapability,
SystemMetadataRequest,
SystemMetadataResult,
TableMetadata,
VersionInfo,
ensure_system_metadata_request,
get_data_dictionary_loader,
get_dialect_config,
system_metadata_gated_result,
)
from sqlspec.data_dictionary.dialects.mssql import (
build_mssql_metadata_capability_profile,
build_mssql_system_metadata_capability,
build_mssql_system_metadata_result,
build_mssql_table_ddl_result,
extract_mssql_version_value,
get_mssql_data_dictionary_options,
is_mssql_azure_sql,
list_mssql_available_features,
merge_mssql_table_lists,
mssql_supports_native_json,
parse_mssql_engine_edition,
parse_mssql_version_components,
resolve_mssql_feature_flag,
validate_mssql_system_metadata_options,
)
from sqlspec.driver import SyncDataDictionaryBase
from sqlspec.utils.logging import get_logger
if TYPE_CHECKING:
from collections.abc import Sequence
from sqlspec.adapters.pymssql.driver import PymssqlDriver
from sqlspec.core import SQL
from sqlspec.data_dictionary._types import DialectConfig, MetadataCapabilityProfile
__all__ = ("MssqlVersionInfo", "PymssqlSyncDataDictionary")
logger = get_logger("sqlspec.adapters.pymssql.data_dictionary")
class MssqlVersionInfo(VersionInfo):
"""MSSQL database version info with build, revision, and Azure SQL detection."""
def __init__(
self,
major: int,
minor: int = 0,
build: int = 0,
revision: int = 0,
edition: str | None = None,
engine_edition: int | None = None,
) -> None:
super().__init__(major, minor, 0)
self.build = build
self.revision = revision
self.edition = edition
self.engine_edition = engine_edition
self.is_azure_sql = is_mssql_azure_sql(engine_edition)
def supports_native_json(self) -> bool:
"""Return whether this server supports the native JSON type."""
return mssql_supports_native_json(self.major, is_azure_sql=self.is_azure_sql)
@property
def version_tuple(self) -> "tuple[int, int, int]":
"""Get version tuple using the MSSQL build number as the third component."""
return (self.major, self.minor, self.build)
def __str__(self) -> str:
"""String representation of version info."""
version_str = f"{self.major}.{self.minor}.{self.build}.{self.revision}"
if self.edition:
version_str += f" ({self.edition})"
if self.is_azure_sql:
version_str += " [Azure]"
return version_str
class _MssqlDataDictionaryMixin:
"""Shared helpers for MSSQL data dictionaries."""
dialect: ClassVar[str] = "mssql"
def get_dialect_config(self) -> "DialectConfig":
"""Return the dialect configuration for this data dictionary."""
return get_dialect_config(type(self).dialect)
def resolve_schema(self, schema: str | None) -> str | None:
"""Return a schema name using dialect defaults when missing."""
if schema is not None:
return schema
return self.get_dialect_config().default_schema
def list_available_features(self) -> list[str]:
"""List available feature flags for this dialect."""
return list_mssql_available_features(self.get_dialect_config())
def get_domain_query(self, domain: str, name: str) -> "SQL":
"""Return a SQL Server domain query."""
query = get_data_dictionary_loader().get_domain_query(type(self).dialect, domain, name)
return cast("SQL", query.sql)
def _build_version_info(
self, version_value: str | None, edition: str | None, engine_edition_value: Any
) -> MssqlVersionInfo | None:
if not version_value:
return None
major, minor, build, revision = parse_mssql_version_components(version_value)
return MssqlVersionInfo(
major,
minor,
build,
revision,
edition=edition,
engine_edition=parse_mssql_engine_edition(engine_edition_value),
)
def _get_optimal_type_from_version(self, version_info: MssqlVersionInfo | None, type_category: str) -> str:
if type_category in {"json", "jsonb"} and version_info is not None and version_info.supports_native_json():
return "JSON"
return self.get_dialect_config().get_optimal_type(type_category)
@mypyc_attr(allow_interpreted_subclasses=True, native_class=False)
class PymssqlSyncDataDictionary(_MssqlDataDictionaryMixin, SyncDataDictionaryBase):
"""MSSQL sync data dictionary."""
dialect: ClassVar[str] = "mssql"
def __init__(self) -> None:
super().__init__()