feat: update ai_list_files schema and upgrade opendal

Author: BohuTANG

README.md

@@ -23,7 +23,7 @@ uv run databend-aiserver --port 8815
 
 ```sql
 CREATE OR REPLACE FUNCTION ai_list_files(stage_location STAGE_LOCATION, max_files INT)
-RETURNS TABLE (stage_name VARCHAR, relative_path VARCHAR, path VARCHAR, is_dir BOOLEAN, size BIGINT, mode VARCHAR, content_type VARCHAR, etag VARCHAR, truncated BOOLEAN)
+RETURNS TABLE (stage_name VARCHAR, path VARCHAR, fullpath VARCHAR, size UINT64, last_modified VARCHAR, etag VARCHAR, content_type VARCHAR)
 LANGUAGE PYTHON HANDLER = 'ai_list_files' ADDRESS = '<your-ai-server-address>';
 
 CREATE OR REPLACE FUNCTION ai_embed_1024(text VARCHAR)

databend_aiserver/stages/operator.py

@@ -134,7 +134,19 @@ def _build_memory_options(_: Mapping[str, Any]) -> Dict[str, Any]:
     return {}
 
 
-_STORAGE_BUILDERS: Dict[str, Any] = {"s3": _build_s3_options, "memory": _build_memory_options}
+def _build_fs_options(storage: Mapping[str, Any]) -> Dict[str, Any]:
+    # Useful for local testing with real files.
+    root = _first_present(storage, "root")
+    if not root:
+        raise StageConfigurationError("FS stage is missing root configuration")
+    return {"root": root}
+
+
+_STORAGE_BUILDERS: Dict[str, Any] = {
+    "s3": _build_s3_options,
+    "memory": _build_memory_options,
+    "fs": _build_fs_options,
+}
 
 
 def _cache_key(stage: StageLocation) -> str:
@@ -251,23 +263,30 @@ def stage_file_suffix(path: str) -> str:
     return Path(path).suffix or ".bin"
 
 
-def resolve_full_path(stage_location: StageLocation, path: str) -> str:
+def resolve_storage_uri(stage_location: StageLocation, path: str) -> str:
     """
-    Resolve the full path (URI) of a file in the stage.
-    Useful for returning the full S3 path in metadata.
+    Resolve the full URI of a path relative to the storage root.
     """
-    resolved_path = resolve_stage_subpath(stage_location, path)
     storage = stage_location.storage or {}
     storage_root = str(storage.get("root", "") or "")
     bucket = storage.get("bucket") or storage.get("name")
 
     if storage_root.startswith("s3://"):
         base = storage_root.rstrip("/")
-        return f"{base}/{resolved_path}"
+        return f"{base}/{path}"
     elif bucket:
         base = f"s3://{bucket}"
         if storage_root:
             base = f"{base}/{storage_root.strip('/')}"
-        return f"{base}/{resolved_path}"
+        return f"{base}/{path}"
 
-    return resolved_path or path
+    return path
+
+
+def resolve_full_path(stage_location: StageLocation, path: str) -> str:
+    """
+    Resolve the full path (URI) of a file in the stage.
+    Useful for returning the full S3 path in metadata.
+    """
+    resolved_path = resolve_stage_subpath(stage_location, path)
+    return resolve_storage_uri(stage_location, resolved_path)

databend_aiserver/udfs/stage.py

@@ -21,13 +21,13 @@
 from typing import Any, Dict, Iterable, List, Optional
 
 from databend_udf import StageLocation, udf
-from opendal import exceptions as opendal_exceptions
-
 from databend_aiserver.stages.operator import (
     StageConfigurationError,
     get_operator,
     resolve_stage_subpath,
+    resolve_storage_uri,
 )
+from opendal import exceptions as opendal_exceptions
 
 
 def _format_last_modified(value: Any) -> Optional[str]:
@@ -119,14 +119,12 @@ def _collect_stage_files(
     input_types=["INT"],
     result_type=[
         ("stage_name", "VARCHAR"),
-        ("relative_path", "VARCHAR"),
         ("path", "VARCHAR"),
-        ("is_dir", "BOOLEAN"),
-        ("size", "BIGINT"),
-        ("mode", "VARCHAR"),
-        ("content_type", "VARCHAR"),
+        ("fullpath", "VARCHAR"),
+        ("size", "UINT64"),
+        ("last_modified", "VARCHAR"),
         ("etag", "VARCHAR"),
-        ("last_modified", "TIMESTAMP"),
+        ("content_type", "VARCHAR"),
     ],
     name="ai_list_files",
 )
@@ -170,22 +168,19 @@ def ai_list_files(
             metadata = op.stat(entry.path)
             # Check if directory using mode (opendal.Metadata doesn't have is_dir())
             # Mode for directories typically has specific bits set, or path ends with /
-            is_dir = entry.path.endswith('/')
 
             # Convert mode to string if it exists, otherwise None
-            mode_str = str(metadata.mode) if metadata.mode is not None else None
-            last_modified = _format_last_modified(getattr(metadata, "last_modified", None))
 
             yield {
                 "stage_name": stage_location.stage_name,
-                "relative_path": stage_location.relative_path,
                 "path": entry.path,
-                "is_dir": is_dir,
+                "fullpath": resolve_storage_uri(stage_location, entry.path),
                 "size": metadata.content_length,
-                "mode": mode_str,
-                "content_type": metadata.content_type,
+                "last_modified": _format_last_modified(
+                    getattr(metadata, "last_modified", None)
+                ),
                 "etag": metadata.etag,
-                "last_modified": last_modified,
+                "content_type": metadata.content_type,
             }
 
     except Exception as e:

pyproject.toml

@@ -11,7 +11,7 @@ requires-python = ">=3.10"
 authors = [{ name = "Databend Labs" }]
 dependencies = [
     "databend-udf==0.2.18",
-    "opendal==0.45.1",
+    "opendal>=0.46.0",
     "pypdf==4.0.0",
     "python-docx==1.2.0",
     "uvicorn>=0.29.0",

tests/integration/test_listing_integration.py

@@ -12,6 +12,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import pytest
 from databend_udf.client import UDFClient
 
 from tests.integration.conftest import build_stage_mapping
@@ -42,18 +43,31 @@ def test_list_stage_files_content(running_server, memory_stage):
 def test_list_stage_files_metadata(running_server, memory_stage):
     rows = _get_listing(running_server, memory_stage)
     assert {row["stage_name"] for row in rows} == {memory_stage.stage_name}
-    assert {row["relative_path"] for row in rows} == {memory_stage.relative_path}
+    # Check for fullpath instead of relative_path
+    # Memory stage fullpath might be just the path if no bucket/root
+    assert all("fullpath" in row for row in rows)
+    assert all(row["fullpath"].endswith(row["path"]) for row in rows)
+    # Check that last_modified key exists (value might be None for memory backend)
+    assert all("last_modified" in row for row in rows)
 
 
 def test_list_stage_files_schema(running_server, memory_stage):
     rows = _get_listing(running_server, memory_stage)
     for row in rows:
-        assert "is_dir" in row
+        assert "path" in row
+        assert "fullpath" in row
         assert "size" in row
-        assert "mode" in row
-        assert "content_type" in row  # May be None
-        assert "etag" in row  # May be None
         assert "last_modified" in row
+        assert "etag" in row  # May be None
+        assert "content_type" in row  # May be None
+        
+        # Verify order implicitly by checking keys list if needed, 
+        # but for now just existence is enough as dicts are ordered in Python 3.7+
+        keys = list(row.keys())
+        # Expected keys: stage_name, path, fullpath, size, last_modified, etag, content_type
+        # Note: stage_name is added by _get_listing or the UDF logic, let's check the core ones
+        assert keys.index("path") < keys.index("fullpath")
+        assert keys.index("last_modified") < keys.index("etag")
 
 
 def test_list_stage_files_truncation(running_server, memory_stage):