codekingpro/portable-devtools

1# Copyright 2026 The HuggingFace Team. All rights reserved.2#3# Licensed under the Apache License, Version 2.0 (the "License");4# you may not use this file except in compliance with the License.5# You may obtain a copy of the License at6#7#     http://www.apache.org/licenses/LICENSE-2.08#9# Unless required by applicable law or agreed to in writing, software10# distributed under the License is distributed on an "AS IS" BASIS,11# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.12# See the License for the specific language governing permissions and13# limitations under the License.14import json15import os16import shutil17import subprocess18from dataclasses import dataclass19from typing import TYPE_CHECKING, Any, Union20 21from . import constants22from .utils import get_token23 24 25if TYPE_CHECKING:26    import duckdb27 28 29@dataclass(frozen=True)30class DatasetParquetEntry:31    """Represents a single parquet file available for a dataset on the Hub."""32 33    config: str34    split: str35    url: str36    size: int37 38 39def execute_raw_sql_query(sql_query: str, *, token: str | bool | None = None) -> list[dict[str, Any]]:40    normalized_query = sql_query.strip().rstrip(";").strip()41    _raise_on_forbidden_query(normalized_query)42 43    connection = None44    try:45        connection = _get_duckdb_connection(token=token)46        relation = connection.sql(normalized_query)47        if relation is None:48            raise ValueError("SQL query must return rows.")49 50        if isinstance(relation, _DuckDBCliRelation):51            # DuckDB binary => run CLI => parse JSON52            return relation.execute()53        else:54            # DuckDB Python API => fetch columns + rows => convert to dicts55            columns = tuple(column[0] for column in relation.description)56            rows = tuple(tuple(row) for row in relation.fetchall())57            return [dict(zip(columns, row)) for row in rows]58    finally:59        if connection is not None:60            connection.close()61 62 63def _raise_on_forbidden_query(query: str) -> None:64    if len(query) == 0:65        raise ValueError("SQL query cannot be empty.")66 67    # DuckDB CLI meta-commands are dot-prefixed words (e.g. `.shell`, `.output`).68    # Let's forbid them for now but allow SQL expressions like `.5` that can legitimately start a line.69    for line in query.splitlines():70        stripped = line.lstrip()71        if stripped.startswith(".") and stripped[1:2].isalpha():72            raise ValueError("DuckDB CLI meta-commands are not allowed in SQL queries.")73 74 75def _get_duckdb_connection(76    token: str | bool | None,77) -> Union["duckdb.DuckDBPyConnection", "_DuckDBCliConnection"]:78    try:79        # If DuckDB is installed as a Python package, use it!80        import duckdb81    except ImportError as error:82        # Otherwise, use the DuckDB CLI binary.83        duckdb_binary = shutil.which("duckdb")84        if duckdb_binary is None:85            raise ImportError(86                "DuckDB is required for `hf datasets sql`. Install the Python package with `pip install duckdb` or "87                "install the DuckDB CLI binary (for example `brew install duckdb`)."88            ) from error89        return _DuckDBCliConnection(binary_path=duckdb_binary, token=token)90 91    # Create a new connection (Python API).92    connection = duckdb.connect()93    try:94        for statement in _build_duckdb_secret_statements(token):95            connection.execute(statement)96        return connection97    except Exception:98        connection.close()99        raise100 101 102@dataclass103class _DuckDBCliConnection:104    """DuckDB connection.105 106    Mimics the DuckDB Python API, but runs the queries via the DuckDB CLI binary.107    """108 109    binary_path: str110    token: str | bool | None111 112    def __post_init__(self) -> None:113        self._setup_statements = _build_duckdb_secret_statements(self.token)114 115    def sql(self, query: str) -> "_DuckDBCliRelation":116        return _DuckDBCliRelation(binary_path=self.binary_path, setup_statements=self._setup_statements, query=query)117 118    def close(self) -> None:119        pass120 121 122@dataclass123class _DuckDBCliRelation:124    """DuckDB relation.125 126    Mimics the DuckDB Python API, but runs the queries via the DuckDB CLI binary.127    """128 129    binary_path: str130    setup_statements: list[str]131    query: str132 133    def execute(self) -> list[dict[str, Any]]:134        # Build the DuckDB CLI input.135        setup = []136        if self.setup_statements:137            setup = [138                f".output {os.devnull}",139                *(f"{stmt};" for stmt in self.setup_statements),140                ".output",141            ]142        full_query = "\n".join(setup + [self.query + ";"])143 144        # Run DuckDB binary145        result = subprocess.run(146            [self.binary_path, "-json"],147            input=full_query,148            capture_output=True,149            text=True,150            check=False,151        )152        if result.returncode != 0:153            error_message = result.stderr.strip() or result.stdout.strip() or "DuckDB CLI command failed."154            raise RuntimeError(error_message)155 156        # Parse JSON output and return157        return json.loads(result.stdout.strip())158 159 160def _build_duckdb_secret_statements(token: str | bool | None) -> list[str]:161    if token is None or token is True:162        token = get_token()163 164    if not token:165        return []166 167    escaped_token = token.replace("'", "''")168    escaped_endpoint = constants.ENDPOINT.replace("'", "''")169    return [170        f"CREATE OR REPLACE SECRET hf_hub_token (TYPE HTTP, BEARER_TOKEN '{escaped_token}', SCOPE '{escaped_endpoint}')",171        f"CREATE OR REPLACE SECRET hf_token (TYPE HUGGINGFACE, TOKEN '{escaped_token}')",172    ]173