use opendal to unify fs access

Author: douenergy

ibis-server/app/query_cache/__init__.py

@@ -1,28 +1,32 @@
 import hashlib
-import os
 from typing import Any, Optional
 
 import ibis
+import opendal
 from loguru import logger
 from opentelemetry import trace
 
 tracer = trace.get_tracer(__name__)
 
 
 class QueryCacheManager:
-    def __init__(self):
-        pass
+    def __init__(self, root: str = "/tmp/wren-engine/"):
+        self.root = root
 
     @tracer.start_as_current_span("get_cache", kind=trace.SpanKind.INTERNAL)
     def get(self, data_source: str, sql: str, info) -> Optional[Any]:
         cache_key = self._generate_cache_key(data_source, sql, info)
-        cache_path = self._get_cache_path(cache_key)
+        cache_file_name = self._get_cache_file_name(cache_key)
+        op = self._get_dal_operator()
+        full_path = self._get_full_path(cache_file_name)
 
         # Check if cache file exists
-        if os.path.exists(cache_path):
+        if op.exists(cache_file_name):
             try:
-                cache = ibis.read_parquet(cache_path)
+                logger.info(f"\nReading query cache {cache_file_name}\n")
+                cache = ibis.read_parquet(full_path)
                 df = cache.execute()
+                logger.info("\nquery cache to dataframe\n")
                 return df
             except Exception as e:
                 logger.debug(f"Failed to read query cache {e}")
@@ -33,14 +37,17 @@ def get(self, data_source: str, sql: str, info) -> Optional[Any]:
     @tracer.start_as_current_span("set_cache", kind=trace.SpanKind.INTERNAL)
     def set(self, data_source: str, sql: str, result: Any, info) -> None:
         cache_key = self._generate_cache_key(data_source, sql, info)
-        cache_path = self._get_cache_path(cache_key)
+        cache_file_name = self._get_cache_file_name(cache_key)
+        op = self._get_dal_operator()
+        full_path = self._get_full_path(cache_file_name)
 
         try:
             # Create cache directory if it doesn't exist
-            os.makedirs(os.path.dirname(cache_path), exist_ok=True)
-            cache = ibis.memtable(result)
-            logger.info(f"\nWriting query cache to {cache_path}\n")
-            cache.to_parquet(cache_path)
+            with op.open(cache_file_name, mode="wb") as file:
+                cache = ibis.memtable(result)
+                logger.info(f"\nWriting query cache to {cache_file_name}\n")
+                if file.writable():
+                    cache.to_parquet(full_path)
         except Exception as e:
             logger.debug(f"Failed to write query cache: {e}")
             return
@@ -57,5 +64,12 @@ def _generate_cache_key(self, data_source: str, sql: str, info) -> str:
 
         return hashlib.sha256(key_string.encode()).hexdigest()
 
-    def _get_cache_path(self, cache_key: str) -> str:
-        return f"/tmp/wren-engine/{cache_key}.cache"
+    def _get_cache_file_name(self, cache_key: str) -> str:
+        return f"{cache_key}.cache"
+
+    def _get_full_path(self, path: str) -> str:
+        return self.root + path
+
+    def _get_dal_operator(self) -> Any:
+        # Default implementation using local filesystem
+        return opendal.Operator("fs", root=self.root)

ibis-server/tests/routers/v2/connector/test_postgres.py

@@ -1,6 +1,4 @@
 import base64
-import os
-import shutil
 from urllib.parse import quote_plus, urlparse
 
 import orjson
@@ -12,7 +10,6 @@
 from testcontainers.postgres import PostgresContainer
 
 from app.model.validator import rules
-from app.query_cache import QueryCacheManager
 from tests.conftest import file_path
 
 pytestmark = pytest.mark.postgres
@@ -127,15 +124,6 @@ def postgres(request) -> PostgresContainer:
     return pg
 
 
-@pytest.fixture(scope="function")
-def cache_dir():
-    temp_dir = "/tmp/wren-engine-test"
-    os.makedirs(temp_dir, exist_ok=True)
-    yield temp_dir
-    # Clean up after the test
-    shutil.rmtree(temp_dir, ignore_errors=True)
-
-
 async def test_query(client, manifest_str, postgres: PostgresContainer):
     connection_info = _to_connection_info(postgres)
     response = await client.post(
@@ -176,16 +164,7 @@ async def test_query(client, manifest_str, postgres: PostgresContainer):
     }
 
 
-async def test_query_with_cache(
-    client, manifest_str, postgres: PostgresContainer, cache_dir, monkeypatch
-):
-    # Override the cache path to use our test directory
-    monkeypatch.setattr(
-        QueryCacheManager,
-        "_get_cache_path",
-        lambda self, key: f"{cache_dir}/{key}.cache",
-    )
-
+async def test_query_with_cache(client, manifest_str, postgres: PostgresContainer):
     connection_info = _to_connection_info(postgres)
 
     # First request - should miss cache

ibis-server/tests/routers/v3/connector/postgres/test_query.py

@@ -1,11 +1,8 @@
 import base64
-import os
-import shutil
 
 import orjson
 import pytest
 
-from app.query_cache import QueryCacheManager
 from tests.routers.v3.connector.postgres.conftest import base_url
 
 manifest = {
@@ -100,15 +97,6 @@ def manifest_str():
     return base64.b64encode(orjson.dumps(manifest)).decode("utf-8")
 
 
-@pytest.fixture(scope="function")
-def cache_dir():
-    temp_dir = "/tmp/wren-engine-test"
-    os.makedirs(temp_dir, exist_ok=True)
-    yield temp_dir
-    # Clean up after the test
-    shutil.rmtree(temp_dir, ignore_errors=True)
-
-
 async def test_query(client, manifest_str, connection_info):
     response = await client.post(
         url=f"{base_url}/query",
@@ -148,16 +136,7 @@ async def test_query(client, manifest_str, connection_info):
     }
 
 
-async def test_query_with_cache(
-    client, manifest_str, connection_info, cache_dir, monkeypatch
-):
-    # Override the cache path to use our test directory
-    monkeypatch.setattr(
-        QueryCacheManager,
-        "_get_cache_path",
-        lambda self, key: f"{cache_dir}/{key}.cache",
-    )
-
+async def test_query_with_cache(client, manifest_str, connection_info):
     # First request - should miss cache
     response1 = await client.post(
         url=f"{base_url}/query?cacheEnable=true",  # Enable cache