feat(ibis): query cache with connection url

ibis-server/app/query_cache/__init__.py

@@ -7,6 +7,8 @@
 from loguru import logger
 from opentelemetry import trace
 
+from app.model import ConnectionUrl
+
 tracer = trace.get_tracer(__name__)
 
 
@@ -70,13 +72,18 @@ def get_cache_file_timestamp(self, data_source: str, sql: str, info) -> int | No
         return None
 
     def _generate_cache_key(self, data_source: str, sql: str, info) -> str:
-        key_parts = [
-            data_source,
-            sql,
-            info.host.get_secret_value(),
-            info.port.get_secret_value(),
-            info.user.get_secret_value(),
-        ]
+        key_parts = []
+        if isinstance(info, ConnectionUrl):
+            key_parts = [data_source, sql, info.connection_url.get_secret_value()]
+        else:
+            key_parts = [
+                data_source,
+                sql,
+                info.host.get_secret_value(),
+                info.port.get_secret_value(),
+                info.user.get_secret_value(),
+            ]
+        logger.debug("Hash key components:  ", key_parts)
         key_string = "|".join(key_parts)
 
         return hashlib.sha256(key_string.encode()).hexdigest()

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

@@ -277,6 +277,72 @@ async def test_query_with_connection_url(
     assert result["dtypes"] is not None
 
 
+async def test_query_with_connection_url_and_cache_enable(
+    client, manifest_str, postgres: PostgresContainer
+):
+    connection_url = _to_connection_url(postgres)
+    # First request - should miss cache then create cache
+    response1 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": 'SELECT * FROM "Orders" LIMIT 1',
+        },
+    )
+
+    assert response1.status_code == 200
+    assert response1.headers["X-Cache-Hit"] == "false"
+    result1 = response1.json()
+
+    response2 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": 'SELECT * FROM "Orders" LIMIT 1',
+        },
+    )
+    assert response2.status_code == 200
+    assert response2.headers["X-Cache-Hit"] == "true"
+    assert int(response2.headers["X-Cache-Create-At"]) > 1743984000  # 2025.04.07
+    result2 = response2.json()
+
+    # Verify results are identical
+    assert result1["data"] == result2["data"]
+
+
+async def test_query_with_connection_url_and_cache_override(
+    client, manifest_str, postgres: PostgresContainer
+):
+    connection_url = _to_connection_url(postgres)
+    # First request - should miss cache then create cache
+    response1 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",  # Enable cache
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": 'SELECT * FROM "Orders" LIMIT 1',
+        },
+    )
+    assert response1.status_code == 200
+
+    # Second request with same SQL - should hit cache and override it
+    response2 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true&overrideCache=true",  # Enable cache
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": 'SELECT * FROM "Orders" LIMIT 1',
+        },
+    )
+    assert response2.status_code == 200
+    assert response2.headers["X-Cache-Override"] == "true"
+    assert int(response2.headers["X-Cache-Override-At"]) > int(
+        response2.headers["X-Cache-Create-At"]
+    )
+
+
 async def test_query_with_dot_all(client, manifest_str, postgres: PostgresContainer):
     connection_info = _to_connection_info(postgres)
     test_sqls = [

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

@@ -108,6 +108,82 @@ async def test_query_with_cache_override(client, manifest_str, connection_info):
     )
 
 
+async def test_query_with_connection_url(client, manifest_str, connection_url):
+    response = await client.post(
+        url=f"{base_url}/query",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT orderkey FROM orders LIMIT 1",
+        },
+    )
+    assert response.status_code == 200
+
+
+async def test_query_with_connection_url_and_cache_enable(
+    client, manifest_str, connection_url
+):
+    response1 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT orderkey FROM orders LIMIT 1",
+        },
+    )
+    assert response1.status_code == 200
+    assert response1.headers["X-Cache-Hit"] == "false"
+    result1 = response1.json()
+
+    response2 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT orderkey FROM orders LIMIT 1",
+        },
+    )
+
+    assert response2.status_code == 200
+    assert response2.headers["X-Cache-Hit"] == "true"
+    result2 = response2.json()
+
+    # Verify results are identical
+    assert result1["data"] == result2["data"]
+    assert result1["columns"] == result2["columns"]
+    assert result1["dtypes"] == result2["dtypes"]
+
+
+async def test_query_with_connection_url_and_cache_override(
+    client, manifest_str, connection_url
+):
+    # First request - should miss cache then create cache
+    response1 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT orderkey FROM orders LIMIT 1",
+        },
+    )
+    assert response1.status_code == 200
+
+    # Second request with same SQL - should hit cache and override it
+    response2 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true&overrideCache=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT orderkey FROM orders LIMIT 1",
+        },
+    )
+    assert response2.status_code == 200
+    assert response2.headers["X-Cache-Override"] == "true"
+    assert int(response2.headers["X-Cache-Override-At"]) > int(
+        response2.headers["X-Cache-Create-At"]
+    )
+
+
 async def test_dry_run(client, manifest_str, connection_info):
     response = await client.post(
         url=f"{base_url}/query",

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

@@ -216,6 +216,69 @@ async def test_query_with_connection_url(client, manifest_str, connection_url):
     assert result["dtypes"] is not None
 
 
+async def test_query_with_connection_url_and_cache_enable(
+    client, manifest_str, connection_url
+):
+    response1 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT * FROM wren.public.orders LIMIT 1",
+        },
+    )
+    assert response1.status_code == 200
+    assert response1.headers["X-Cache-Hit"] == "false"
+    result1 = response1.json()
+
+    # Second request with same SQL - should hit cache
+    response2 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT * FROM wren.public.orders LIMIT 1",
+        },
+    )
+    assert response2.status_code == 200
+    assert response2.headers["X-Cache-Hit"] == "true"
+    result2 = response2.json()
+
+    assert result1["data"] == result2["data"]
+    assert result1["columns"] == result2["columns"]
+    assert result1["dtypes"] == result2["dtypes"]
+
+
+async def test_query_with_connection_url_and_cache_override(
+    client, manifest_str, connection_url
+):
+    # First request - should miss cache then create cache
+    response1 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT * FROM wren.public.orders LIMIT 1",
+        },
+    )
+    assert response1.status_code == 200
+
+    # Second request with same SQL - should hit cache and override it
+    response2 = await client.post(
+        url=f"{base_url}/query?cacheEnable=true&overrideCache=true",
+        json={
+            "connectionInfo": {"connectionUrl": connection_url},
+            "manifestStr": manifest_str,
+            "sql": "SELECT * FROM wren.public.orders LIMIT 1",
+        },
+    )
+    assert response2.status_code == 200
+    assert response2.headers["X-Cache-Override"] == "true"
+    assert int(response2.headers["X-Cache-Override-At"]) > int(
+        response2.headers["X-Cache-Create-At"]
+    )
+
+
 async def test_query_with_limit(client, manifest_str, connection_info):
     response = await client.post(
         url=f"{base_url}/query",