diff --git a/ibis-server/app/query_cache/__init__.py b/ibis-server/app/query_cache/__init__.py index fc04260b4..2e8044525 100644 --- a/ibis-server/app/query_cache/__init__.py +++ b/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() diff --git a/ibis-server/tests/routers/v2/connector/test_postgres.py b/ibis-server/tests/routers/v2/connector/test_postgres.py index 5717f85c3..6e868da9f 100644 --- a/ibis-server/tests/routers/v2/connector/test_postgres.py +++ b/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 = [ diff --git a/ibis-server/tests/routers/v3/connector/postgres/test_fallback_v2.py b/ibis-server/tests/routers/v3/connector/postgres/test_fallback_v2.py index aa68593e5..2ef62e955 100644 --- a/ibis-server/tests/routers/v3/connector/postgres/test_fallback_v2.py +++ b/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", diff --git a/ibis-server/tests/routers/v3/connector/postgres/test_query.py b/ibis-server/tests/routers/v3/connector/postgres/test_query.py index 29d212eb9..d65dacafa 100644 --- a/ibis-server/tests/routers/v3/connector/postgres/test_query.py +++ b/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",