completions.py

from __future__ import annotations

import sqlite3

from harlequin.autocomplete.completion import HarlequinCompletion


def get_completion_data(conn: sqlite3.Connection) -> list[HarlequinCompletion]:
    KEYWORDS = [
        "abort",
        "action",
        "add",
        "after",
        "all",
        "alter",
        "always",
        "analyze",
        "and",
        "as",
        "asc",
        "attach",
        "autoincrement",
        "before",
        "begin",
        "between",
        "by",
        "cascade",
        "case",
        "cast",
        "check",
        "collate",
        "column",
        "commit",
        "conflict",
        "constraint",
        "create",
        "cross",
        "current",
        "current_date",
        "current_time",
        "current_timestamp",
        "database",
        "default",
        "deferrable",
        "deferred",
        "delete",
        "desc",
        "detach",
        "distinct",
        "do",
        "drop",
        "each",
        "else",
        "end",
        "escape",
        "except",
        "exclude",
        "exclusive",
        "exists",
        "explain",
        "fail",
        "filter",
        "first",
        "following",
        "for",
        "foreign",
        "from",
        "full",
        "generated",
        "glob",
        "group",
        "groups",
        "having",
        "if",
        "ignore",
        "immediate",
        "in",
        "index",
        "indexed",
        "initially",
        "inner",
        "insert",
        "instead",
        "intersect",
        "into",
        "is",
        "isnull",
        "join",
        "key",
        "last",
        "left",
        "like",
        "limit",
        "match",
        "materialized",
        "natural",
        "no",
        "not",
        "nothing",
        "notnull",
        "null",
        "nulls",
        "of",
        "offset",
        "on",
        "or",
        "order",
        "others",
        "outer",
        "over",
        "partition",
        "plan",
        "pragma",
        "preceding",
        "primary",
        "query",
        "raise",
        "range",
        "recursive",
        "references",
        "regexp",
        "reindex",
        "release",
        "rename",
        "replace",
        "restrict",
        "returning",
        "right",
        "rollback",
        "row",
        "rows",
        "savepoint",
        "select",
        "set",
        "table",
        "temp",
        "temporary",
        "then",
        "ties",
        "to",
        "transaction",
        "trigger",
        "unbounded",
        "union",
        "unique",
        "update",
        "using",
        "vacuum",
        "values",
        "view",
        "virtual",
        "when",
        "where",
        "window",
        "with",
        "without",
    ]

    PRAGMAS = [
        "analysis_limit",
        "application_id",
        "auto_vacuum",
        "automatic_index",
        "busy_timeout",
        "cache_size",
        "cache_spill",
        "cell_size_check",
        "checkpoint_fullfsync",
        "collation_list",
        "compile_options",
        "data_version",
        "database_list",
        "defer_foreign_keys",
        "encoding",
        "foreign_key_check",
        "foreign_key_list",
        "foreign_keys",
        "freelist_count",
        "fullfsync",
        "function_list",
        "hard_heap_limit",
        "ignore_check_constraints",
        "incremental_vacuum",
        "index_info",
        "index_list",
        "index_xinfo",
        "integrity_check",
        "journal_mode",
        "journal_size_limit",
        "legacy_alter_table",
        "legacy_file_format",
        "locking_mode",
        "max_page_count",
        "mmap_size",
        "module_list",
        "optimize",
        "page_count",
        "page_size",
        "pragma_list",
        "query_only",
        "quick_check",
        "read_uncommitted",
        "recursive_triggers",
        "reverse_unordered_selects",
        "secure_delete",
        "shrink_memory",
        "soft_heap_limit",
        "synchronous",
        "table_info",
        "table_list",
        "table_xinfo",
        "temp_store",
        "threads",
        "trusted_schema",
        "user_version",
        "wal_autocheckpoint",
        "wal_checkpoint",
    ]

    function_data = conn.execute(
        """
        select distinct
            name as label,
            case when type = 'w' then 'agg' else 'fn' end as type_label
        from pragma_function_list
        order by name asc
        """
    ).fetchall()

    keyword_completions = [
        HarlequinCompletion(
            label=kw_name, type_label="kw", value=kw_name, priority=100, context=None
        )
        for kw_name in KEYWORDS
    ]

    pragma_completions = [
        HarlequinCompletion(
            label=label, type_label="pragma", value=label, priority=1000, context=None
        )
        for label in PRAGMAS
    ]

    function_completions = [
        HarlequinCompletion(
            label=label, type_label=type_label, value=label, priority=1000, context=None
        )
        for label, type_label in function_data
    ]

    return [*keyword_completions, *pragma_completions, *function_completions]