cat dialect_inferer.py

import re
import sys
from collections import defaultdict
from typing import Dict, List, Tuple

class AmbiguousDialectError(Exception):
    def __init__(self, dialects):
        self.dialects = dialects
        message = f"Ambiguous dialect match: {', '.join(dialects)}"
        super().__init__(message)


class SQLDialectInferrer:
    def __init__(self):
        self.dialect_patterns: Dict[str, List[Tuple[str, int]]] = {
            "redshift": [
                # Redshift-specific keywords and procedural patterns
                (r"\bdiststyle\b", 2),
                (r"\bdistkey\b", 2),
                (r"\bsortkey\b", 2),
                (r"\bcompound sortkey\b", 3),
                (r"\binterleaved sortkey\b", 3),
                (r"\bvacuum\b", 2),
                (r"\bunload\b", 2),
                (r"\bcopy\b", 2),
                (r"\bcreate\s+materialized\s+view\b", 3),
                (r"\blanguage\s+plpgsql\b", 3),  # Redshift supports PL/pgSQL
                (r"\bdo\s*\$\$", 3),  # PL/pgSQL procedural language block
                (r"\breturn\s+query\b", 2),  # PL/pgSQL return query
                (r"\bexecute\b", 3),  # Redshift procedural EXECUTE statements
                (r"\bcreate\s+procedure\b", 3),  # Redshift stored procedures
                (r"\bprocedure\b", 2),  # General procedure keyword
                (r"\bcommit\b", 1),  # Transactions are supported in Redshift SPs
                # Redshift's standard SQL keywords:
                (r"\bcurrent_timestamp\b", 2),
                (r"\bgetdate\(\)", 2),  # Redshift supports GETDATE()
                (r"\bnow\(\)", 2),  # NOW() is common for current timestamp in Redshift
                (r"\bnvarchar\b", 2),  # Redshift uses NVARCHAR
                (r"\bchar\b", 2),  # Redshift supports CHAR
                (r"\bvarchar\b", 2),
                (r"\bboolean\b", 2),  # Redshift BOOLEAN type
                (r"\bdate\b", 2),
                (r"\btimestamp\b", 2),
                (r"\binteger\b", 2),  # Redshift uses INTEGER types
            ],
            "tsql": [
                (r"\bbegin\s+try\b", 2),
                (r"\bend\s+try\b", 2),
                (r"\bbegin\s+catch\b", 2),
                (r"\bend\s+catch\b", 2),
                (r"\bexec\b", 2),
                (r"\bexecute\b", 2),
                (r"\bsp_\w+", 2),  # Stored procedure calls
                (r"\bwith\s*\(\s*nolock\s*\)", 3),  # NOLOCK hint for table access
                (r"\bmerge\b", 2),
                (r"\bdeclare\s+@", 4),  # Variable declarations in T-SQL
                (r"@@", 4),  # System variable declarations in T-SQL
                (
                    r"\bisnumeric\(\)",
                    2,
                ),  # ISNUMERIC() function to check if a value is numeric
                (r"\btransaction\b", 2),
                (r"\bidentity\b", 2),
                (r"\bgetdate\(\)", 2),  # T-SQL specific datetime function
                (r"\bcoalesce\(", 2),  # COALESCE is common but more prominent in T-SQL
                (r"\btry_convert\b", 2),  # T-SQL specific type conversion function
                (r"\brow_number\(\)", 3),  # T-SQL's analytic function
                (r"\btop\s+\d+", 2),  # TOP N syntax in T-SQL for limiting rows
                (
                    r"\bdb_name\(\)",
                    2,
                ),  # DB_NAME() function to get the name of the current database
                (
                    r"\bobject_schema_name\(\)",
                    2,
                ),  # OBJECT_SCHEMA_NAME(@@PROCID) to get the schema of the current object
                (
                    r"\bobject_name\(\)",
                    2,
                ),  # OBJECT_NAME() to get the name of the current object
                (
                    r"\bobject_schema_name\(\s*@@procid\s*\)",
                    2,
                ),  # OBJECT_SCHEMA_NAME() to get the schema of the current object
                (
                    r"\bobject_name\(\s*@@procid\s*\)",
                    2,
                ),  # OBJECT_NAME(@@PROCID) to get the name of the current object
                (
                    r"\bcreate\s+nonclustered\s+index\b",
                    2,
                ),  # CREATE NONCLUSTERED INDEX for performance optimization in T-SQL
                (
                    r"\buse\s*\[\s*\w+\s*\]",
                    2,
                ),  # USE [database_name] for switching database context
            ],
        }

        # Common SQL keywords that should not affect dialect inference
        self.common_sql_keywords = set(
            [
                "select",
                "from",
                "where",
                "and",
                "or",
                "insert",
                "update",
                "delete",
                "create",
                "table",
                "index",
                "drop",
                "alter",
                "join",
                "on",
                "group by",
                "order by",
                "having",
                "union",
                "all",
                "as",
                "distinct",
                "like",
                "in",
                "between",
                "is",
                "null",
                "not",
                "case",
                "when",
                "then",
                "else",
                "end",
            ]
        )

    def remove_comments(self, content: str) -> str:
        """Remove single-line and multi-line comments from the SQL content."""
        # Remove single-line comments
        content = re.sub(r"--.*$", "", content, flags=re.MULTILINE)
        # Remove multi-line comments
        content = re.sub(r"/\*[\s\S]*?\*/", "", content)
        return content

    def calculate_sql_likelihood(self, content: str) -> float:
        """Calculate how likely the content is to be SQL based on common keywords."""
        words = re.findall(r"\b\w+\b", content.lower())
        sql_keywords = sum(1 for word in words if word in self.common_sql_keywords)
        return sql_keywords / len(words) if words else 0

    def score_content_against_dialects(self, content: str) -> list:
        # Score the content against known dialect patterns
        scores = defaultdict(int)
        for dialect, patterns in self.dialect_patterns.items():
            for pattern, weight in patterns:
                if re.search(pattern, content):
                    scores[dialect] += weight

        if not scores:
            return "ansi"  # Assume ANSI SQL if no specific dialect patterns are matched

        max_score = max(scores.values())
        return [d for d, s in scores.items() if s == max_score]

    def infer_dialect(self, file_path: str) -> str:
        """Infer SQL dialect from the provided file."""
        try:
            with open(file_path, "r", encoding="utf-8") as file:
                content = file.read().lower()
        except UnicodeDecodeError:
            try:
                with open(file_path, "r", encoding="latin-1") as file:
                    content = file.read().lower()
            except Exception as e:
                print(f"Error reading file: {e}")
                return "unknown"

        content = self.remove_comments(content)

        # Check if the file is likely to be SQL based on common keywords
        sql_likelihood = self.calculate_sql_likelihood(content)
        if sql_likelihood < 0.1:  # Adjust this threshold as needed
            raise RuntimeError("File isn't SQL or does not contain enough SQL keywords")

        top_dialects = self.score_content_against_dialects(content)

        if len(top_dialects) == 1:
            return top_dialects[0]  # Return the single top dialect
        else:
            raise AmbiguousDialectError(
                top_dialects
            )  # Raise exception if multiple dialects have the same score

    def update_sqlfluff_config(self, config_path: str, inferred_dialect: str) -> None:
        try:
            with open(config_path, "r") as config_file:
                config_lines = config_file.readlines()

            updated = False
            with open(config_path, "w") as config_file:
                for line in config_lines:
                    if line.strip().startswith("dialect ="):
                        config_file.write(f"dialect = {inferred_dialect}\n")
                        updated = True
                    else:
                        config_file.write(line)

                if not updated:
                    config_file.write(f"dialect = {inferred_dialect}\n")

            print(f"Updated .sqlfluff config with dialect: {inferred_dialect}")

        except Exception as e:
            print(f"Error updating .sqlfluff config: {e}")


def main(file_path: str) -> None:
    """Main function to run the SQL dialect inference."""
    inferrer = SQLDialectInferrer()
    try:
        dialect = inferrer.infer_dialect(file_path)
        print(dialect)
    except AmbiguousDialectError as e:
        print(f"Error: {e}")
        sys.exit(1)


if __name__ == "__main__":
    if len(sys.argv) != 2:
        print("Usage: python sql_dialect_inferrer.py <path_to_sql_file>")
        sys.exit(1)
    main(sys.argv[1])