diff --git a/lumen/ai/controls/ingest/constants.py b/lumen/ai/controls/ingest/constants.py index 723d63aec..735dbf244 100644 --- a/lumen/ai/controls/ingest/constants.py +++ b/lumen/ai/controls/ingest/constants.py @@ -4,7 +4,12 @@ # File type constants # ───────────────────────────────────────────────────────────────────────────── -TABLE_EXTENSIONS = ("csv", "parquet", "parq", "json", "xlsx", "geojson", "wkt", "zip") +XARRAY_EXTENSIONS = ( + "nc", "nc4", "netcdf", "h5", "hdf5", "he5", "zarr", + "grib", "grib2", "grb", "grb2", +) + +TABLE_EXTENSIONS = ("csv", "parquet", "parq", "json", "xlsx", "geojson", "wkt", "zip", *XARRAY_EXTENSIONS) METADATA_EXTENSIONS = ("md", "txt", "yaml", "yml", "json", "pdf", "docx", "doc", "pptx", "ppt") METADATA_FILENAME_PATTERNS = ("_metadata", "metadata_", "readme", "schema") diff --git a/lumen/ai/ui.py b/lumen/ai/ui.py index 0629e03a1..426caedef 100644 --- a/lumen/ai/ui.py +++ b/lumen/ai/ui.py @@ -1,6 +1,9 @@ from __future__ import annotations import asyncio +import atexit +import os +import tempfile from contextlib import contextmanager from functools import partial @@ -33,7 +36,8 @@ from ..pipeline import Pipeline from ..sources import Source from ..sources.duckdb import DuckDBSource -from ..util import log +from ..sources.xarray_sql import XArraySQLSource +from ..util import check_xarray_available, log from .agents import ( AnalysisAgent, BaseCodeAgent, ChatAgent, DocumentListAgent, DocumentSummarizerAgent, SQLAgent, TableListAgent, ValidationAgent, @@ -48,6 +52,7 @@ BaseSourceControls, DownloadSourceControls, FileSourceControls, SourceCatalog, TableExplorer, UploadSourceControls, ) +from .controls.ingest.constants import XARRAY_EXTENSIONS from .coordinator import Coordinator, Plan, Planner from .editors import AnalysisOutput, LumenEditor, SQLEditor from .export import export_notebook @@ -508,6 +513,36 @@ def __init__( self._configure_session() self._render_page() + @staticmethod + def _get_xarray_upload_handlers(): + """Build upload handlers for xarray file formats. + + Returns empty dict if xarray-sql is not installed. + """ + def _cleanup_temp_files(): + for path in _temp_files: + try: + os.unlink(path) + except OSError: + pass + + def _handle_xarray_upload(context, file_obj, alias, filename): + ext = filename.rsplit('.', 1)[-1] if '.' in filename else 'nc' + tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f'.{ext}') + file_obj.seek(0) + tmp.write(file_obj.read()) + tmp.flush() + tmp.close() + _temp_files.append(tmp.name) + return XArraySQLSource(uri=tmp.name, name=alias) + + if not check_xarray_available(): + return {} + + _temp_files: list[str] = [] + atexit.register(_cleanup_temp_files) + return {ext: _handle_xarray_upload for ext in XARRAY_EXTENSIONS} + @classmethod def _resolve_data( cls, data: DataT | list[DataT] | dict[DataT] | None @@ -597,6 +632,12 @@ def _resolve_data( sources.append(source) continue + # Handle xarray files (NetCDF, Zarr, HDF5, GRIB) + if Path(src).suffix.lstrip('.').lower() in XARRAY_EXTENSIONS: + source = XArraySQLSource(uri=str(Path(src).absolute())) + sources.append(source) + continue + if src.startswith('http'): remote = True if src.endswith(('.parq', '.parquet', '.csv', '.json', '.tsv', '.jsonl', '.ndjson')): @@ -1272,6 +1313,10 @@ async def _do_sync(): param.parameterized.async_executor(_do_sync) self._source_catalog.param.watch(_schedule_visibility_change, 'visibility_changed') + # Register xarray upload handlers (if xarray-sql is available) + xarray_handlers = self._get_xarray_upload_handlers() + merged_upload_handlers = {**xarray_handlers, **self.upload_handlers} + # Initialize source controls self._source_controls = [] control_tabs = [] @@ -1288,7 +1333,7 @@ async def _do_sync(): 'source_catalog': self._source_catalog, } if issubclass(control, FileSourceControls): - control_kwargs['upload_handlers'] = self.upload_handlers + control_kwargs['upload_handlers'] = merged_upload_handlers if control is UploadSourceControls and self.filedropper_kwargs: control_kwargs['filedropper_kwargs'] = self.filedropper_kwargs control_inst = control(**control_kwargs) diff --git a/lumen/command/ai.py b/lumen/command/ai.py index e4f673070..39a1db25d 100644 --- a/lumen/command/ai.py +++ b/lumen/command/ai.py @@ -241,10 +241,15 @@ def __init__( continue path = Path(table_path) - if not path.is_file(): - if path.exists(): - raise ValueError(f"Table path is not a file: {table_path}") + if not path.exists(): raise FileNotFoundError(f"Table file not found: {table_path}") + if path.is_dir() and not table_path.endswith('.zarr'): + raise ValueError( + f"Table path is a directory: {table_path}. " + "Only .zarr directories are supported." + ) + if not path.is_file() and not path.is_dir(): + raise ValueError(f"Table path is not a file: {table_path}") source = self._build_source_code( tables=tables, diff --git a/lumen/sources/xarray_sql.py b/lumen/sources/xarray_sql.py index 1b51d7886..fe152a923 100644 --- a/lumen/sources/xarray_sql.py +++ b/lumen/sources/xarray_sql.py @@ -14,18 +14,10 @@ import param from ..transforms.sql import SQLFilter +from ..util import check_xarray_available from .base import BaseSQLSource, cached -def _check_xarray_available(): - """Check xarray and xarray-sql are installed, import lazily.""" - try: - import xarray # noqa - import xarray_sql # noqa - return True - except ImportError: - return False - class XArraySQLSource(BaseSQLSource): """ SQL-queryable xarray data source for Lumen. @@ -101,7 +93,7 @@ class XArraySQLSource(BaseSQLSource): The SQL expression template for table queries.""") def __init__(self, _dataset=None, _ctx=None, **params): - if not _check_xarray_available(): + if not check_xarray_available(): raise ImportError( "xarray and xarray-sql are required for XArraySQLSource. " "Install them with: pip install lumen[xarray]" diff --git a/lumen/tests/ai/test_ui.py b/lumen/tests/ai/test_ui.py index f07d818bf..23fdef43e 100644 --- a/lumen/tests/ai/test_ui.py +++ b/lumen/tests/ai/test_ui.py @@ -1,13 +1,15 @@ import asyncio +import io import sqlite3 import sys import tempfile from pathlib import Path from types import SimpleNamespace -from unittest.mock import AsyncMock, patch +from unittest.mock import AsyncMock, MagicMock, patch import duckdb +import numpy as np import pytest try: @@ -15,6 +17,14 @@ except ModuleNotFoundError: pytest.skip("lumen.ai could not be imported, skipping tests.", allow_module_level=True) +try: + import xarray as xr + + from lumen.sources.xarray_sql import XArraySQLSource + HAS_XARRAY = True +except ImportError: + HAS_XARRAY = False + from panel.layout import Column, Row from panel.tests.util import async_wait_until from panel.util import edit_readonly @@ -1126,6 +1136,106 @@ def test_resolve_data_pipeline_and_files(self): # Table key is sanitized assert 'data_csv' in duckdb_source.tables + @pytest.mark.skipif(not HAS_XARRAY, reason="xarray not installed") + def test_resolve_data_netcdf_file(self): + """Test resolving a NetCDF file creates XArraySQLSource.""" + with tempfile.TemporaryDirectory() as tmpdir: + nc_path = Path(tmpdir) / 'test.nc' + ds = xr.Dataset({'temperature': (('lat', 'lon'), np.random.rand(3, 4))}) + try: + ds.to_netcdf(nc_path) + except Exception: + pytest.skip("NetCDF write backend unavailable") + result = UI._resolve_data(str(nc_path)) + assert len(result) == 1 + assert isinstance(result[0], XArraySQLSource) + assert 'temperature' in result[0].get_tables() + + def test_resolve_data_xarray_import_error(self): + """Test graceful error when xarray-sql not installed.""" + with patch('lumen.sources.xarray_sql.check_xarray_available', return_value=False): + with pytest.raises(ImportError, match="xarray"): + UI._resolve_data('data.nc') + + @pytest.mark.skipif(not HAS_XARRAY, reason="xarray not installed") + def test_resolve_data_mixed_xarray_and_csv(self): + """Test that .nc creates separate source while .csv goes to DuckDB.""" + with tempfile.TemporaryDirectory() as tmpdir: + nc_path = Path(tmpdir) / 'test.nc' + ds = xr.Dataset({'temp': (('x',), np.array([1.0, 2.0]))}) + try: + ds.to_netcdf(nc_path) + except Exception: + pytest.skip("NetCDF write backend unavailable") + result = UI._resolve_data([str(nc_path), 'data.csv']) + assert len(result) == 2 + xarray_sources = [s for s in result if isinstance(s, XArraySQLSource)] + duckdb_sources = [s for s in result if isinstance(s, DuckDBSource)] + assert len(xarray_sources) == 1 + assert len(duckdb_sources) == 1 + + +class TestCLIPathValidation: + """Tests for CLI path validation in command/ai.py.""" + + def test_zarr_directory_accepted(self): + """Test that .zarr directories pass validation.""" + with tempfile.TemporaryDirectory() as tmpdir: + zarr_path = Path(tmpdir) / 'test.zarr' + zarr_path.mkdir() + # Validation should not raise for .zarr directories + path = Path(str(zarr_path)) + assert path.exists() + assert path.is_dir() + assert str(zarr_path).endswith('.zarr') + + def test_random_directory_rejected(self): + """Test that non-zarr directories are rejected by validation logic.""" + with tempfile.TemporaryDirectory() as tmpdir: + dir_path = Path(tmpdir) / 'not_zarr' + dir_path.mkdir() + path = Path(str(dir_path)) + # Should be rejected: is a directory but not .zarr + assert path.is_dir() + assert not str(dir_path).endswith('.zarr') + + +@pytest.mark.skipif(not HAS_XARRAY, reason="xarray not installed") +class TestXarrayUploadHandler: + """Tests for the xarray upload handler.""" + + def test_upload_handler_returns_dict(self): + """Upload handlers should be returned for all xarray extensions.""" + handlers = UI._get_xarray_upload_handlers() + assert isinstance(handlers, dict) + for ext in ('nc', 'nc4', 'netcdf', 'h5', 'hdf5', 'he5', 'zarr', + 'grib', 'grib2', 'grb', 'grb2'): + assert ext in handlers + + def test_upload_handler_creates_source_from_nc(self): + """Upload handler should create XArraySQLSource from NetCDF bytes.""" + # Create a NetCDF file in memory + ds = xr.Dataset({'temp': (('x',), np.array([1.0, 2.0, 3.0]))}) + buf = io.BytesIO() + try: + ds.to_netcdf(buf) + except Exception: + pytest.skip("NetCDF write backend unavailable") + buf.seek(0) + + handlers = UI._get_xarray_upload_handlers() + source = handlers['nc'](context={}, file_obj=buf, alias='test_data', filename='test') + + assert isinstance(source, XArraySQLSource) + assert 'temp' in source.get_tables() + assert source.name == 'test_data' + + def test_upload_handler_returns_empty_without_xarray(self): + """Upload handlers should be empty if xarray-sql not installed.""" + with patch('lumen.ai.ui.check_xarray_available', return_value=False): + handlers = UI._get_xarray_upload_handlers() + assert handlers == {} + # --- Tests for on_edit callback --- @@ -1229,7 +1339,6 @@ async def test_edit_updates_logs(explorer_ui, tmp_path): ui = explorer_ui # Set up a mock for _logs - from unittest.mock import MagicMock mock_logs = MagicMock() ui._logs = mock_logs diff --git a/lumen/util.py b/lumen/util.py index fcac103a0..6130ed9dc 100644 --- a/lumen/util.py +++ b/lumen/util.py @@ -491,3 +491,13 @@ def normalize_table_name(name: str) -> str: 'table_name' """ return re.sub(r'\W+', '_', name).strip('_').lower() + + +def check_xarray_available(): + """Check if xarray and xarray-sql are installed.""" + try: + import xarray # noqa + import xarray_sql # noqa + return True + except ImportError: + return False