Wire XArraySQLSource into Lumen AI for natural language queries on scientific data

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")

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
@@ -64,6 +69,8 @@
 
 DataT = str | Path | Source | Pipeline
 
+XARRAY_DOTTED_EXTENSIONS = tuple(f'.{ext}' for ext in XARRAY_EXTENSIONS)
+
 PAGE_SX = {
     ".sidebar": {"transition": "width 0.2s ease-in-out"},
     ".sidebar:hover": {"width": "140px", "transitionDelay": "0.5s"},
@@ -508,6 +515,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 +634,12 @@ def _resolve_data(
                         sources.append(source)
                     continue
 
+                # Handle xarray files (NetCDF, Zarr, HDF5, GRIB)
+                if src.endswith(XARRAY_DOTTED_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 +1315,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 +1335,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)

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,

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]"

lumen/tests/ai/test_ui.py

@@ -1,20 +1,30 @@
 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:
     import lumen.ai  # noqa
 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
 

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