spice-mcp 0.1.1__py3-none-any.whl → 0.1.3__py3-none-any.whl

spice_mcp/adapters/dune/extract.py

@@ -45,7 +45,7 @@ if TYPE_CHECKING:
 
 
 # ---------------------------------------------------------------------------
-# Back-compat helpers expected by tests and adapter code
+# Internal helpers used by adapter code and tests
 
 def _is_sql(query: int | str) -> bool:
     if isinstance(query, int):
@@ -305,11 +305,24 @@ def query(
                 df = get_results(**execute_kwargs, **result_kwargs)
                 if df is not None:
                     return process_result(df, execution, **output_kwargs)
-            execution = execute_query(**execute_kwargs, verbose=verbose)
+            try:
+                execution = execute_query(**execute_kwargs, verbose=verbose)
+            except Exception as e:
+                # Re-raise with more context about the failure
+                if verbose:
+                    print(f'execute_query failed for query_id={query_id}, parameters={parameters}')
+                raise Exception(f'failed to execute query {query_id}: {e}') from e
+        else:
+            # query_id is None or falsy - this shouldn't happen for valid inputs
+            if verbose:
+                print(f'query_id is falsy: {query_id}, query_or_execution={query_or_execution}')
 
-        # await execution completion
+        # check execution status
         if execution is None:
-            raise Exception('could not determine execution')
+            error_detail = f'query_id={query_id}, query_type={type(query_or_execution).__name__}'
+            if isinstance(query_or_execution, str):
+                error_detail += f', query_preview={query_or_execution[:100]}'
+            raise Exception(f'could not determine execution ({error_detail})')
         if poll:
             poll_execution(execution, **poll_kwargs)
             df = get_results(execution, api_key, **result_kwargs)
@@ -322,206 +335,8 @@ def query(
 
 
 @overload
-async def async_query(
-    query_or_execution: Query | Execution,
-    *,
-    verbose: bool = True,
-    refresh: bool = False,
-    max_age: float | None = None,
-    parameters: Mapping[str, Any] | None = None,
-    api_key: str | None = None,
-    performance: Performance = 'medium',
-    poll: Literal[False],
-    poll_interval: float = 1.0,
-    timeout_seconds: float | None = None,
-    limit: int | None = None,
-    offset: int | None = None,
-    sample_count: int | None = None,
-    sort_by: str | None = None,
-    columns: Sequence[str] | None = None,
-    extras: Mapping[str, Any] | None = None,
-    types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    all_types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    cache: bool = True,
-    cache_dir: str | None = None,
-    save_to_cache: bool = True,
-    load_from_cache: bool = True,
-    include_execution: bool = False,
-) -> Execution: ...
-
-
 @overload
-async def async_query(
-    query_or_execution: Query | Execution,
-    *,
-    verbose: bool = True,
-    refresh: bool = False,
-    max_age: float | None = None,
-    parameters: Mapping[str, Any] | None = None,
-    api_key: str | None = None,
-    performance: Performance = 'medium',
-    poll: Literal[True] = True,
-    poll_interval: float = 1.0,
-    limit: int | None = None,
-    offset: int | None = None,
-    sample_count: int | None = None,
-    sort_by: str | None = None,
-    columns: Sequence[str] | None = None,
-    extras: Mapping[str, Any] | None = None,
-    types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    all_types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    cache: bool = True,
-    cache_dir: str | None = None,
-    save_to_cache: bool = True,
-    load_from_cache: bool = True,
-    include_execution: Literal[False] = False,
-) -> pl.DataFrame: ...
-
-
 @overload
-async def async_query(
-    query_or_execution: Query | Execution,
-    *,
-    verbose: bool = True,
-    refresh: bool = False,
-    max_age: float | None = None,
-    parameters: Mapping[str, Any] | None = None,
-    api_key: str | None = None,
-    performance: Performance = 'medium',
-    poll: Literal[True] = True,
-    poll_interval: float = 1.0,
-    limit: int | None = None,
-    offset: int | None = None,
-    sample_count: int | None = None,
-    sort_by: str | None = None,
-    columns: Sequence[str] | None = None,
-    extras: Mapping[str, Any] | None = None,
-    types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    all_types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    cache: bool = True,
-    cache_dir: str | None = None,
-    save_to_cache: bool = True,
-    load_from_cache: bool = True,
-    include_execution: Literal[True],
-) -> tuple[pl.DataFrame, Execution]: ...
-
-
-async def async_query(
-    query_or_execution: Query | Execution,
-    *,
-    verbose: bool = True,
-    refresh: bool = False,
-    max_age: float | None = None,
-    parameters: Mapping[str, Any] | None = None,
-    api_key: str | None = None,
-    performance: Performance = 'medium',
-    poll: bool = True,
-    poll_interval: float = 1.0,
-    timeout_seconds: float | None = None,
-    limit: int | None = None,
-    offset: int | None = None,
-    sample_count: int | None = None,
-    sort_by: str | None = None,
-    columns: Sequence[str] | None = None,
-    extras: Mapping[str, Any] | None = None,
-    types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    all_types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    cache: bool = True,
-    cache_dir: str | None = None,
-    save_to_cache: bool = True,
-    load_from_cache: bool = True,
-    include_execution: bool = False,
-):
-
-    # determine whether target is a query or an execution
-    query_id, execution, parameters = _determine_input_type(
-        query_or_execution,
-        parameters,
-    )
-
-    # gather arguments
-    execute_kwargs: ExecuteKwargs = {
-        'query_id': query_id,
-        'api_key': api_key,
-        'parameters': parameters,
-        'performance': performance,
-    }
-    poll_kwargs: PollKwargs = {
-        'poll_interval': poll_interval,
-        'api_key': api_key,
-        'verbose': verbose,
-        'timeout_seconds': timeout_seconds,
-    }
-    result_kwargs: RetrievalKwargs = {
-        'limit': limit,
-        'offset': offset,
-        'sample_count': sample_count,
-        'sort_by': sort_by,
-        'columns': columns,
-        'extras': extras,
-        'types': types,
-        'all_types': all_types,
-        'verbose': verbose,
-    }
-    output_kwargs: OutputKwargs = {
-        'execute_kwargs': execute_kwargs,
-        'result_kwargs': result_kwargs,
-        'cache': cache,
-        'save_to_cache': save_to_cache,
-        'cache_dir': cache_dir,
-        'include_execution': include_execution,
-    }
-
-    # execute or retrieve query
-    if query_id:
-        if cache and cache_dir is not None and not refresh:
-            cache_result, cache_execution = await _cache.async_load_from_cache(
-                execute_kwargs, result_kwargs, output_kwargs
-            )
-            if cache_result is not None:
-                return cache_result
-            if execution is None and cache_execution is not None:
-                execution = cache_execution
-        if max_age is not None and not refresh:
-            age = await async_get_query_latest_age(**execute_kwargs, verbose=verbose)  # type: ignore
-            if age is None or age > max_age:
-                refresh = True
-        if not refresh:
-            df = await async_get_results(**execute_kwargs, **result_kwargs)
-            if df is not None:
-                return await async_process_result(df, execution, **output_kwargs)
-        execution = await async_execute_query(**execute_kwargs, verbose=verbose)
-
-    # await execution completion
-    if execution is None:
-        raise Exception('could not determine execution')
-    if poll:
-        await async_poll_execution(execution, **poll_kwargs)
-        df = await async_get_results(execution, api_key, **result_kwargs)
-        if df is not None:
-            return await async_process_result(df, execution, **output_kwargs)
-        else:
-            raise Exception('no successful execution for query')
-    else:
-        return execution
-
-
 def _process_result(
     df: pl.DataFrame,
     execution: Execution | None,
@@ -551,35 +366,6 @@ def _process_result(
         return df
 
 
-async def _async_process_result(
-    df: pl.DataFrame,
-    execution: Execution | None,
-    execute_kwargs: ExecuteKwargs,
-    result_kwargs: RetrievalKwargs,
-    cache: bool,
-    save_to_cache: bool,
-    cache_dir: str | None,
-    include_execution: bool,
-) -> pl.DataFrame | tuple[pl.DataFrame, Execution]:
-    if cache and save_to_cache and execute_kwargs['query_id'] is not None:
-        if execution is None:
-            execution = await async_get_latest_execution(execute_kwargs)
-            if execution is None:
-                raise Exception('could not get execution')
-        _cache.save_to_cache(
-            df, execution, execute_kwargs, result_kwargs, cache_dir
-        )
-
-    if include_execution:
-        if execution is None:
-            execution = await async_get_latest_execution(execute_kwargs)
-            if execution is None:
-                raise Exception('could not get execution')
-        return df, execution
-    else:
-        return df
-
-
 def _get_query_latest_age(
     query_id: int,
     *,
@@ -654,83 +440,6 @@ def _parse_timestamp(timestamp: str) -> int:
     return int(timestamp_float)
 
 
-async def _async_get_query_latest_age(
-    query_id: int,
-    *,
-    verbose: bool = True,
-    parameters: Mapping[str, Any] | None = None,
-    performance: Performance = 'medium',
-    api_key: str | None = None,
-) -> float | None:
-    import datetime
-    import json
-
-    import aiohttp
-
-    # process inputs
-    if api_key is None:
-        api_key = _urls.get_api_key()
-    headers = {'X-Dune-API-Key': api_key, 'User-Agent': get_user_agent()}
-    data = {}
-    if parameters is not None:
-        data['query_parameters'] = parameters
-    url = _urls.get_query_results_url(query_id, parameters=data, csv=False)
-
-    # print summary
-    if verbose:
-        print('checking age of last execution, query_id = ' + str(query_id))
-
-    # perform request with retry/backoff for 429/502
-    timeout = aiohttp.ClientTimeout(total=30)
-    async with aiohttp.ClientSession(timeout=timeout) as session:
-        attempts = 0
-        backoff = 0.5
-        while True:
-            async with session.get(url, headers=headers) as response:
-                if response.status in (429, 502):
-                    attempts += 1
-                    if attempts >= 3:
-                        break
-                    import asyncio
-                    import random
-                    await asyncio.sleep(backoff * random.uniform(1.5, 2.5))
-                    backoff = min(5.0, backoff * 2)
-                    continue
-                result: Mapping[str, Any] = await response.json()
-                break
-
-    # check if result is error
-    try:
-        if 'error' in result:
-            if (
-                result['error']
-                == 'not found: No execution found for the latest version of the given query'
-            ):
-                if verbose:
-                    print(
-                        'no age for query, because no previous executions exist'
-                    )
-                return None
-            raise Exception(result['error'])
-    except json.JSONDecodeError:
-        pass
-
-    # process result
-    if 'execution_started_at' in result:
-        now = datetime.datetime.now(datetime.UTC).timestamp()
-        started = _parse_timestamp(result['execution_started_at'])
-        age = now - started
-
-        if verbose:
-            print('latest result age:', age)
-
-        return age
-    else:
-        if verbose:
-            print('no age for query, because no previous executions exist')
-        return None
-
-
 def _execute(
     query_id: int | str,
     *,
@@ -755,50 +464,25 @@ def _execute(
 
     # perform request
     response = _transport_post(url, headers=headers, json=data, timeout=_POST_TIMEOUT)
-    result: Mapping[str, Any] = response.json()
-
-    # check for errors
-    if 'execution_id' not in result:
-        raise Exception(result['error'])
-
-    # process result
-    execution_id = result['execution_id']
-    return {'execution_id': execution_id, 'timestamp': None}
-
-
-async def _async_execute(
-    query_id: int | str,
-    *,
-    parameters: Mapping[str, Any] | None = None,
-    performance: Performance = 'medium',
-    api_key: str | None = None,
-    verbose: bool = True,
-) -> Execution:
-    import aiohttp
-
-    # process inputs
-    url = _urls.get_query_execute_url(query_id)
-    if api_key is None:
-        api_key = _urls.get_api_key()
-    headers = {'X-Dune-API-Key': api_key, 'User-Agent': get_user_agent()}
-    data = {}
-    if parameters is not None:
-        data['query_parameters'] = parameters
-    data['performance'] = performance
-
-    # print summary
-    if verbose:
-        print('executing query, query_id = ' + str(query_id))
-
-    # perform request
-    timeout = aiohttp.ClientTimeout(total=_POST_TIMEOUT)
-    async with aiohttp.ClientSession(timeout=timeout) as session:
-        async with session.post(url, headers=headers, json=data) as response:
-            result: Mapping[str, Any] = await response.json()
+    
+    # Parse response with better error handling
+    try:
+        result: Mapping[str, Any] = response.json()
+    except Exception as e:
+        if verbose:
+            print(f'failed to parse response JSON: {e}')
+            print(f'response status: {response.status_code}')
+            print(f'response text: {response.text[:500]}')
+        raise Exception(f'failed to parse response: {e}') from e
 
     # check for errors
     if 'execution_id' not in result:
-        raise Exception(result['error'])
+        error_msg = result.get('error', f'response missing execution_id: {result}')
+        if verbose:
+            print(f'execution failed: {error_msg}')
+            print(f'response status: {response.status_code}')
+            print(f'full response: {result}')
+        raise Exception(error_msg)
 
     # process result
     execution_id = result['execution_id']
@@ -912,124 +596,6 @@ def _get_results(
     return df
 
 
-async def _async_get_results(
-    execution: Execution | None = None,
-    api_key: str | None = None,
-    *,
-    query_id: int | None = None,
-    parameters: Mapping[str, Any] | None = None,
-    performance: Performance = 'medium',
-    limit: int | None = None,
-    offset: int | None = None,
-    sample_count: int | None = None,
-    sort_by: str | None = None,
-    columns: Sequence[str] | None = None,
-    extras: Mapping[str, Any] | None = None,
-    types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    all_types: Sequence[type[pl.DataType]]
-    | Mapping[str, type[pl.DataType]]
-    | None = None,
-    verbose: bool = True,
-) -> pl.DataFrame | None:
-    import asyncio
-    import random
-
-    import aiohttp
-    import polars as pl
-
-    if api_key is None:
-        api_key = _urls.get_api_key()
-    headers = {'X-Dune-API-Key': api_key, 'User-Agent': get_user_agent()}
-    params: dict[str, Any] = {
-        'limit': limit,
-        'offset': offset,
-        'sample_count': sample_count,
-        'sort_by': sort_by,
-        'columns': columns,
-    }
-    if extras is not None:
-        params.update(extras)
-    if parameters is not None:
-        params['query_parameters'] = parameters
-    if query_id is not None:
-        url = _urls.get_query_results_url(query_id, parameters=params)
-    elif execution is not None:
-        url = _urls.get_execution_results_url(execution['execution_id'], params)
-    else:
-        raise Exception('must specify query_id or execution')
-
-    # print summary
-    if verbose:
-        if query_id is not None:
-            print('getting results, query_id = ' + str(query_id))
-        elif execution is not None:
-            print('getting results, execution_id = ' + str(execution['execution_id']))
-
-    # perform request
-    timeout = aiohttp.ClientTimeout(total=_GET_TIMEOUT)
-    async with aiohttp.ClientSession(timeout=timeout) as session:
-        # GET with simple retry/backoff for 429/502
-        attempts = 0
-        backoff = 0.5
-        while True:
-            async with session.get(url, headers=headers) as response:
-                if response.status in (429, 502):
-                    attempts += 1
-                    if attempts >= 3:
-                        break
-                    await asyncio.sleep(backoff * random.uniform(1.5, 2.5))
-                    backoff = min(5.0, backoff * 2)
-                    continue
-                if response.status == 404:
-                    return None
-                result = await response.text()
-                response_headers = response.headers
-                break
-
-    # process result
-    df = _process_raw_table(result, types=types, all_types=all_types)
-
-    # support pagination when using limit
-    if limit is not None:
-        import polars as pl
-
-        n_rows = len(df)
-        pages = []
-        timeout = aiohttp.ClientTimeout(total=30)
-        async with aiohttp.ClientSession(timeout=timeout) as session:
-            while 'x-dune-next-uri' in response_headers and n_rows < limit:
-                if verbose:
-                    off = response.headers.get('x-dune-next-offset', 'unknown')
-                    print('gathering additional page, offset = ' + str(off))
-                next_url = response_headers['x-dune-next-uri']
-                # Pager GET with retry/backoff
-                attempts = 0
-                backoff = 0.5
-                while True:
-                    async with session.get(next_url, headers=headers) as response:
-                        if response.status in (429, 502):
-                            attempts += 1
-                            if attempts >= 3:
-                                break
-                            await asyncio.sleep(backoff * random.uniform(1.5, 2.5))
-                            backoff = min(5.0, backoff * 2)
-                            continue
-                        result = await response.text()
-                        response_headers = response.headers
-                        break
-                page = _process_raw_table(
-                    result, types=types, all_types=all_types
-                )
-                n_rows += len(page)
-                pages.append(page)
-
-        df = pl.concat([df, *pages]).limit(limit)
-
-    return df
-
-
 def _process_raw_table(
     raw_csv: str,
     types: Sequence[type[pl.DataType] | None]
@@ -1224,97 +790,6 @@ def _poll_execution(
         )
 
 
-async def _async_poll_execution(
-    execution: Execution,
-    *,
-    api_key: str | None,
-    poll_interval: float,
-    verbose: bool,
-    timeout_seconds: float | None,
-) -> None:
-    import asyncio
-    import random
-
-    import aiohttp
-
-    # process inputs
-    url = _urls.get_execution_status_url(execution['execution_id'])
-    execution_id = execution['execution_id']
-    if api_key is None:
-        api_key = _urls.get_api_key()
-    headers = {'X-Dune-API-Key': api_key, 'User-Agent': get_user_agent()}
-
-    # print summary
-    t_start = time.time()
-
-    # poll until completion
-    timeout = aiohttp.ClientTimeout(total=_GET_TIMEOUT)
-    async with aiohttp.ClientSession(timeout=timeout) as session:
-        sleep_amount = poll_interval
-        while True:
-            t_poll = time.time()
-
-            # print summary
-            if verbose:
-                print(
-                    'waiting for results, execution_id = '
-                    + str(execution['execution_id'])
-                    + ', t = '
-                    + str(t_poll - t_start)
-                )
-
-            # poll
-            async with session.get(url, headers=headers) as response:
-                result = await response.json()
-                if (
-                    'is_execution_finished' not in result
-                    and response.status == 429
-                ):
-                    sleep_amount = sleep_amount * random.uniform(1, 2)
-                    await asyncio.sleep(sleep_amount)
-                    continue
-                if result['is_execution_finished']:
-                    if result['state'] == 'QUERY_STATE_FAILED':
-                        err_detail = ''
-                        try:
-                            if 'error' in result and result['error']:
-                                err_detail = f", error={result['error']}"
-                        except Exception:
-                            pass
-                        raise Exception(
-                            f"QUERY FAILED execution_id={execution_id} state={result.get('state')}{err_detail}"
-                        )
-                    execution['timestamp'] = _parse_timestamp(
-                        result['execution_started_at']
-                    )
-                    break
-
-            # timeout check
-            if timeout_seconds is not None and (t_poll - t_start) > timeout_seconds:
-                raise TimeoutError(
-                    f'query polling timed out after {timeout_seconds} seconds'
-                )
-
-            # wait until polling interval
-            t_wait = time.time() - t_poll
-            if t_wait < poll_interval:
-                import asyncio
-
-                await asyncio.sleep(poll_interval - t_wait)
-
-    # check for errors
-    if result['state'] == 'QUERY_STATE_FAILED':
-        err_detail = ''
-        try:
-            if 'error' in result and result['error']:
-                err_detail = f", error={result['error']}"
-        except Exception:
-            pass
-        raise Exception(
-            f"QUERY FAILED execution_id={execution_id} state={result.get('state')}{err_detail}"
-        )
-
-
 def get_latest_execution(
     execute_kwargs: ExecuteKwargs,
     *,
@@ -1378,74 +853,6 @@ def get_latest_execution(
     return execution
 
 
-async def async_get_latest_execution(
-    execute_kwargs: ExecuteKwargs,
-    *,
-    allow_unfinished: bool = False,
-) -> Execution | None:
-    import json
-    import random
-
-    import aiohttp
-
-    query_id = execute_kwargs['query_id']
-    api_key = execute_kwargs['api_key']
-    parameters = execute_kwargs['parameters']
-    if query_id is None:
-        raise Exception('query_id required for async_get_latest_execution')
-
-    # process inputs
-    if api_key is None:
-        api_key = _urls.get_api_key()
-    headers = {'X-Dune-API-Key': api_key, 'User-Agent': get_user_agent()}
-    data: dict[str, Any] = {}
-    if parameters is not None:
-        data['query_parameters'] = parameters
-    data['limit'] = 0
-    url = _urls.get_query_results_url(query_id, parameters=data, csv=False)
-
-    # perform request
-    timeout = aiohttp.ClientTimeout(total=30)
-    async with aiohttp.ClientSession(timeout=timeout) as session:
-        sleep_amount = 1.0
-        while True:
-            async with session.get(url, headers=headers) as response:
-                if response.status in (429, 502):
-                    sleep_amount = sleep_amount * random.uniform(1, 2)
-                    import asyncio
-                    await asyncio.sleep(sleep_amount)
-                    continue
-                result: Mapping[str, Any] = await response.json()
-
-                # check if result is error
-                try:
-                    if 'error' in result:
-                        if (
-                            result['error']
-                            == 'not found: No execution found for the latest version of the given query'
-                        ):
-                            return None
-                        if response.status == 429:
-                            import asyncio
-
-                            sleep_amount = sleep_amount * random.uniform(1, 2)
-                            await asyncio.sleep(sleep_amount)
-                        raise Exception(result['error'])
-                except json.JSONDecodeError:
-                    pass
-            break
-
-    # process result
-    if not result['is_execution_finished'] and not allow_unfinished:
-        return None
-    execution: Execution = {'execution_id': result['execution_id']}
-    if 'execution_started_at' in result:
-        execution['timestamp'] = int(
-            _parse_timestamp(result['execution_started_at'])
-        )
-    return execution
-
-
 def get_user_agent() -> str:
     # Identify as spice-mcp vendored spice client
     return 'spice-mcp/' + ADAPTER_VERSION
@@ -1455,12 +862,7 @@ def get_user_agent() -> str:
 
 determine_input_type = _determine_input_type
 get_query_latest_age = _get_query_latest_age
-async_get_query_latest_age = _async_get_query_latest_age
 execute_query = _execute
-async_execute_query = _async_execute
 get_results = _get_results
-async_get_results = _async_get_results
 process_result = _process_result
-async_process_result = _async_process_result
 poll_execution = _poll_execution
-async_poll_execution = _async_poll_execution