feat(ibis): generic-default index operations (create/drop/index_exists)

src/mountainash_data/backends/ibis/_index.py

@@ -0,0 +1,249 @@
+"""Generic-default index DDL: pure builders + dispatchers (spec §5).
+
+Pure builders take pre-computed, already-validated parts so registry golden
+tests render every dialect without a live connection.
+"""
+
+from __future__ import annotations
+
+import typing as t
+
+from mountainash_data.backends.ibis._render import (
+    compile_index_predicate,
+    dialect_of,
+    qualified_name,
+    quote_identifier,
+)
+from mountainash_data.backends.ibis.dialects._registry import DropScope, IndexCapability
+from mountainash_data.backends.ibis.operations import (
+    _generate_index_name,
+    _normalize_columns,
+    _validate_simple_identifier,
+)
+
+# USING <method> position differs across dialects (verified against official docs):
+#   - Postgres:    CREATE INDEX i ON tbl USING gin (cols)   -> after ON, before columns
+#   - MySQL/MariaDB: CREATE INDEX i USING btree ON tbl (cols) -> after index name, before ON
+#   - SingleStore: CREATE INDEX i ON tbl (cols) USING hash   -> after columns (the default)
+# sqlite/duckdb/motherduck/mssql/oracle have empty index_types -> no USING emitted.
+_USING_BEFORE_ON: frozenset[str] = frozenset({"mysql"})
+_USING_BEFORE_COLUMNS: frozenset[str] = frozenset({"postgres"})
+
+
+def build_create_index_sql(
+    *,
+    dialect: t.Any,
+    target: str,
+    index_name: str,
+    cols: list[str],
+    unique: bool,
+    index_type: t.Optional[str],
+    guard: str,
+    where_sql: t.Optional[str],
+) -> str:
+    """Render a CREATE INDEX statement from pre-validated parts.
+
+    Args:
+        dialect: sqlglot dialect string (e.g. ``dialect_of(ibis_conn)``).
+        target: already-qualified, already-quoted table reference.
+        index_name: unquoted index name.
+        cols: unquoted column names.
+        unique: emit CREATE UNIQUE INDEX.
+        index_type: USING <type>, or None for no USING clause.
+        guard: ``"IF NOT EXISTS "`` or ``""`` (emulation supplies idempotency).
+        where_sql: rendered partial-index WHERE body, or None.
+    """
+    unique_sql = "UNIQUE " if unique else ""
+    cols_sql = ", ".join(quote_identifier(c, dialect) for c in cols)
+    name_sql = quote_identifier(index_name, dialect)
+    where = f" WHERE {where_sql}" if where_sql else ""
+    name_part = f"{guard}{name_sql}"
+    # dialect may be a sqlglot Dialect class (live path) or a plain string (tests/golden).
+    # Normalise to the lowercase name so membership checks work in both cases.
+    if isinstance(dialect, type):
+        d = dialect.__name__.lower()
+    elif isinstance(dialect, str):
+        d = dialect.lower()
+    else:
+        d = str(dialect).lower()
+    using = f"USING {index_type}" if index_type else None
+
+    if using and d in _USING_BEFORE_ON:
+        # MySQL/MariaDB: USING sits between the index name and ON.
+        name_part = f"{name_part} {using}"
+        tail = f"ON {target} ({cols_sql})"
+    elif using and d in _USING_BEFORE_COLUMNS:
+        # Postgres: USING sits after ON, before the column list.
+        tail = f"ON {target} {using} ({cols_sql})"
+    elif using:
+        # SingleStore (and the general default): USING after the column list.
+        tail = f"ON {target} ({cols_sql}) {using}"
+    else:
+        tail = f"ON {target} ({cols_sql})"
+
+    return f"CREATE {unique_sql}INDEX {name_part} {tail}{where}"
+
+
+def build_drop_index_sql(
+    *,
+    dialect: t.Any,
+    drop_scope: DropScope,
+    index_name: str,
+    target: t.Optional[str],
+    guard: str,
+) -> str:
+    """Render a DROP INDEX statement. `target` is required (already quoted) when
+    `drop_scope` is TABLE_SCOPED."""
+    name_sql = quote_identifier(index_name, dialect)
+    if drop_scope is DropScope.TABLE_SCOPED:
+        return f"DROP INDEX {guard}{name_sql} ON {target}"
+    return f"DROP INDEX {guard}{name_sql}"
+
+
+# ---------------------------------------------------------------------------
+# Generic dispatchers (spec §5-§8)
+# ---------------------------------------------------------------------------
+
+
+def _generic_index_exists(
+    ibis_conn: t.Any,
+    index_name: str,
+    *,
+    table_name: t.Optional[str] = None,
+    database: t.Optional[str] = None,
+    exists_sql_fn: t.Any,
+) -> bool:
+    """Run the dialect's introspection SQL and return whether the index exists."""
+    if exists_sql_fn is None:
+        raise NotImplementedError("dialect has no get_index_exists_sql")
+    _validate_simple_identifier(index_name, kind="index_name")
+    if table_name is not None:
+        _validate_simple_identifier(table_name, kind="table_name")
+    if database is not None:
+        _validate_simple_identifier(database, kind="database")
+    result = ibis_conn.sql(exists_sql_fn(index_name, table_name, database))
+    if result is None:
+        return False
+    import mountainash as ma
+
+    # Read the single returned column BY POSITION, not by the alias name:
+    # Oracle upper-cases the unquoted `count` alias ("count" -> "COUNT"), so
+    # keying by "count" would KeyError. Every introspection query returns
+    # exactly one column.
+    data = ma.relation(result).to_dict()
+    first_col = next(iter(data.values()))
+    return first_col[0] > 0
+
+
+def _target_ref(ibis_conn: t.Any, table_name: str, database: t.Optional[str]) -> str:
+    dialect = dialect_of(ibis_conn)
+    parts = [database, table_name] if database else [table_name]
+    return qualified_name(parts, dialect)
+
+
+def _generic_create_index(
+    ibis_conn: t.Any,
+    table_name: str,
+    columns: t.Union[list[str], str],
+    *,
+    index_name: t.Optional[str] = None,
+    unique: bool = False,
+    index_type: t.Optional[str] = None,
+    where: t.Any = None,
+    database: t.Optional[str] = None,
+    if_not_exists: bool = True,
+    caps: IndexCapability,
+    exists_sql_fn: t.Any,
+) -> None:
+    """Render and execute a CREATE INDEX via the generic path (spec §5-§8).
+
+    Emulation failure modes (TOCTOU / privilege / catalog-isolation /
+    auto-commit DDL) are documented-and-accepted per spec §6: the engine's
+    error is surfaced, never swallowed.
+    """
+    _validate_simple_identifier(table_name, kind="table_name")
+    if database is not None:
+        _validate_simple_identifier(database, kind="database")
+    cols = _normalize_columns(columns)
+    for c in cols:
+        _validate_simple_identifier(c, kind="column")
+
+    if index_type is not None and index_type not in caps.index_types:
+        raise ValueError(
+            f"index_type {index_type!r} not supported by this dialect; "
+            f"valid: {sorted(caps.index_types) or 'none'}"
+        )
+    if where is not None and not caps.partial:
+        raise ValueError("this dialect does not support partial indexes (where=)")
+
+    if index_name is None:
+        index_name = _generate_index_name(table_name, cols, unique=unique)
+    _validate_simple_identifier(index_name, kind="index_name")
+
+    # Idempotency: native guard, or emulate via precheck.
+    guard = ""
+    if if_not_exists:
+        if caps.native_if_not_exists:
+            guard = "IF NOT EXISTS "
+        elif _generic_index_exists(
+            ibis_conn, index_name, table_name=table_name, database=database,
+            exists_sql_fn=exists_sql_fn,
+        ):
+            return  # emulated: already present
+
+    where_sql = None
+    if where is not None:
+        schema = ibis_conn.table(table_name, database=database).schema()
+        where_sql = compile_index_predicate(ibis_conn, schema, table_name, where)
+
+    sql = build_create_index_sql(
+        dialect=dialect_of(ibis_conn),
+        target=_target_ref(ibis_conn, table_name, database),
+        index_name=index_name, cols=cols, unique=unique,
+        index_type=index_type, guard=guard, where_sql=where_sql,
+    )
+    ibis_conn.raw_sql(sql)
+
+
+def _generic_drop_index(
+    ibis_conn: t.Any,
+    index_name: str,
+    *,
+    table_name: t.Optional[str] = None,
+    database: t.Optional[str] = None,
+    if_exists: bool = True,
+    caps: IndexCapability,
+    exists_sql_fn: t.Any,
+) -> None:
+    """Render and execute a DROP INDEX via the generic path (spec §5-§8).
+
+    Emulation failure modes (TOCTOU / privilege / catalog-isolation /
+    auto-commit DDL) are documented-and-accepted per spec §6: the engine's
+    error is surfaced, never swallowed.
+    """
+    _validate_simple_identifier(index_name, kind="index_name")
+    if caps.drop_scope is DropScope.TABLE_SCOPED and table_name is None:
+        raise ValueError(
+            "drop_index requires table_name for this dialect (DROP INDEX ... ON tbl)"
+        )
+    if table_name is not None:
+        _validate_simple_identifier(table_name, kind="table_name")
+    if database is not None:
+        _validate_simple_identifier(database, kind="database")
+
+    guard = ""
+    if if_exists:
+        if caps.native_if_exists:
+            guard = "IF EXISTS "
+        elif not _generic_index_exists(
+            ibis_conn, index_name, table_name=table_name, database=database,
+            exists_sql_fn=exists_sql_fn,
+        ):
+            return  # emulated: already absent
+
+    target = _target_ref(ibis_conn, table_name, database) if table_name else None
+    sql = build_drop_index_sql(
+        dialect=dialect_of(ibis_conn), drop_scope=caps.drop_scope,
+        index_name=index_name, target=target, guard=guard,
+    )
+    ibis_conn.raw_sql(sql)

src/mountainash_data/backends/ibis/_render.py

@@ -209,3 +209,60 @@ def _remap(n: exp.Expression) -> exp.Expression:
         return n
 
     return on.transform(_remap)
+
+
+# ---------------------------------------------------------------------------
+# Index-predicate compiler (§5.2)
+# ---------------------------------------------------------------------------
+
+INDEX_SENTINEL = "__ma_index_tbl__"
+
+IndexPredicate = t.Callable[[ir.Table], ir.BooleanValue]
+
+
+def compile_index_predicate(
+    ibis_conn: t.Any,
+    schema: t.Any,
+    table_name: str,
+    predicate: IndexPredicate,
+) -> str:
+    """Compile a single-table ``(table) -> bool`` predicate to an UNQUALIFIED
+    boolean SQL string for the connection's dialect (partial-index WHERE).
+
+    Mechanism (spec §5.2): bind one sentinel-named ibis table at `schema`,
+    filter it by the predicate, compile to sqlglot, extract the WHERE, then
+    strip every column's table/db/catalog qualifier at the AST level (NOT by
+    string replacement). The predicate may reference any column of the table,
+    not only the indexed columns, so the full `schema` is bound. The
+    `table_name` parameter is validated only (sentinel-collision check) and
+    does NOT appear in the returned SQL.
+
+    Raises:
+        ValueError: if `table_name` collides with the reserved sentinel, or the
+            predicate contains a forbidden op (aggregation/window/subquery).
+    """
+    if table_name == INDEX_SENTINEL:
+        raise ValueError(
+            f"target table name {table_name!r} collides with a reserved sentinel."
+        )
+    tbl = ibis.table(schema, name=INDEX_SENTINEL)
+    pred = predicate(tbl)
+    validate_predicate(pred)
+
+    filtered = tbl.filter(pred)
+    ast = ibis_conn.compiler.to_sqlglot(filtered)
+    ast = ast if isinstance(ast, exp.Expression) else ast[0]
+
+    where = next(ast.find_all(exp.Where), None)
+    if where is None or where.this is None:
+        raise ValueError("could not extract WHERE predicate from compiled AST")
+    cond = where.this.copy()
+
+    def _strip(n: exp.Expression) -> exp.Expression:
+        if isinstance(n, exp.Column):
+            n.set("table", None)
+            n.set("db", None)
+            n.set("catalog", None)
+        return n
+
+    return cond.transform(_strip).sql(dialect=dialect_of(ibis_conn))