diff --git a/ibis-server/app/model/__init__.py b/ibis-server/app/model/__init__.py index 4d1617026..47b3a4d1e 100644 --- a/ibis-server/app/model/__init__.py +++ b/ibis-server/app/model/__init__.py @@ -14,6 +14,18 @@ connection_info_field = Field(alias="connectionInfo") +class BaseConnectionInfo(BaseModel): + def to_key_string(self) -> str: + key_parts = [] + + # Get all fields that are SecretStr and get their values + for _, field_value in self: + if isinstance(field_value, SecretStr): + key_parts.append(field_value.get_secret_value()) + + return "|".join(key_parts) + + class QueryDTO(BaseModel): sql: str manifest_str: str = manifest_str_field @@ -72,7 +84,11 @@ class QueryGcsFileDTO(QueryDTO): connection_info: GcsFileConnectionInfo = connection_info_field -class BigQueryConnectionInfo(BaseModel): +class ConnectionUrl(BaseConnectionInfo): + connection_url: SecretStr = Field(alias="connectionUrl") + + +class BigQueryConnectionInfo(BaseConnectionInfo): project_id: SecretStr = Field(description="GCP project id", examples=["my-project"]) dataset_id: SecretStr = Field( description="BigQuery dataset id", examples=["my_dataset"] @@ -82,7 +98,7 @@ class BigQueryConnectionInfo(BaseModel): ) -class CannerConnectionInfo(BaseModel): +class CannerConnectionInfo(BaseConnectionInfo): host: SecretStr = Field( description="the hostname of your database", examples=["localhost"] ) @@ -101,7 +117,7 @@ class CannerConnectionInfo(BaseModel): ) -class ClickHouseConnectionInfo(BaseModel): +class ClickHouseConnectionInfo(BaseConnectionInfo): host: SecretStr = Field( description="the hostname of your database", examples=["localhost"] ) @@ -117,7 +133,7 @@ class ClickHouseConnectionInfo(BaseModel): ) -class MSSqlConnectionInfo(BaseModel): +class MSSqlConnectionInfo(BaseConnectionInfo): host: SecretStr = Field( description="the hostname of your database", examples=["localhost"] ) @@ -142,7 +158,7 @@ class MSSqlConnectionInfo(BaseModel): ) -class MySqlConnectionInfo(BaseModel): +class MySqlConnectionInfo(BaseConnectionInfo): host: SecretStr = Field( description="the hostname of your database", examples=["localhost"] ) @@ -170,11 +186,7 @@ class MySqlConnectionInfo(BaseModel): ) -class ConnectionUrl(BaseModel): - connection_url: SecretStr = Field(alias="connectionUrl") - - -class PostgresConnectionInfo(BaseModel): +class PostgresConnectionInfo(BaseConnectionInfo): host: SecretStr = Field( examples=["localhost"], description="the hostname of your database" ) @@ -190,7 +202,7 @@ class PostgresConnectionInfo(BaseModel): ) -class OracleConnectionInfo(BaseModel): +class OracleConnectionInfo(BaseConnectionInfo): host: SecretStr = Field( examples=["localhost"], description="the hostname of your database" ) @@ -206,7 +218,7 @@ class OracleConnectionInfo(BaseModel): ) -class SnowflakeConnectionInfo(BaseModel): +class SnowflakeConnectionInfo(BaseConnectionInfo): user: SecretStr = Field( description="the username of your database", examples=["admin"] ) @@ -226,7 +238,7 @@ class SnowflakeConnectionInfo(BaseModel): ) # Use `sf_schema` to avoid `schema` shadowing in BaseModel -class TrinoConnectionInfo(BaseModel): +class TrinoConnectionInfo(BaseConnectionInfo): host: SecretStr = Field( description="the hostname of your database", examples=["localhost"] ) @@ -247,7 +259,7 @@ class TrinoConnectionInfo(BaseModel): ) -class LocalFileConnectionInfo(BaseModel): +class LocalFileConnectionInfo(BaseConnectionInfo): url: SecretStr = Field( description="the root path of the local file", default="/", examples=["/data"] ) @@ -256,7 +268,7 @@ class LocalFileConnectionInfo(BaseModel): ) -class S3FileConnectionInfo(BaseModel): +class S3FileConnectionInfo(BaseConnectionInfo): url: SecretStr = Field( description="the root path of the s3 bucket", default="/", examples=["/data"] ) @@ -277,7 +289,7 @@ class S3FileConnectionInfo(BaseModel): ) -class MinioFileConnectionInfo(BaseModel): +class MinioFileConnectionInfo(BaseConnectionInfo): url: SecretStr = Field( description="the root path of the minio bucket", default="/", examples=["/data"] ) @@ -303,7 +315,7 @@ class MinioFileConnectionInfo(BaseModel): ) -class GcsFileConnectionInfo(BaseModel): +class GcsFileConnectionInfo(BaseConnectionInfo): url: SecretStr = Field( description="the root path of the gcs bucket", default="/", examples=["/data"] ) diff --git a/ibis-server/app/query_cache/__init__.py b/ibis-server/app/query_cache/__init__.py index 2e8044525..203035b47 100644 --- a/ibis-server/app/query_cache/__init__.py +++ b/ibis-server/app/query_cache/__init__.py @@ -7,8 +7,6 @@ from loguru import logger from opentelemetry import trace -from app.model import ConnectionUrl - tracer = trace.get_tracer(__name__) @@ -72,19 +70,10 @@ 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 = [] - 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) + connection_key = info.to_key_string() + + # Combine with data source and SQL + key_string = f"{data_source}|{sql}|{connection_key}" return hashlib.sha256(key_string.encode()).hexdigest() diff --git a/ibis-server/tests/routers/v2/connector/test_bigquery.py b/ibis-server/tests/routers/v2/connector/test_bigquery.py index 764fbd5d8..0cbcb9d92 100644 --- a/ibis-server/tests/routers/v2/connector/test_bigquery.py +++ b/ibis-server/tests/routers/v2/connector/test_bigquery.py @@ -112,6 +112,64 @@ async def test_query(client, manifest_str): } +async def test_query_with_cache(client, manifest_str): + response1 = await client.post( + url=f"{base_url}/query?cacheEnable=true", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": 'SELECT * FROM "Orders" ORDER BY orderkey 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": connection_info, + "manifestStr": manifest_str, + "sql": 'SELECT * FROM "Orders" ORDER BY orderkey 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() + + assert result1["dtypes"] == result2["dtypes"] + + +async def test_query_with_cache_override(client, manifest_str): + response1 = await client.post( + url=f"{base_url}/query?cacheEnable=true", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": 'SELECT * FROM "Orders" ORDER BY orderkey LIMIT 1', + }, + ) + assert response1.status_code == 200 + + response2 = await client.post( + url=f"{base_url}/query?cacheEnable=true&overrideCache=true", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": 'SELECT * FROM "Orders" ORDER BY orderkey 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_without_manifest(client): response = await client.post( url=f"{base_url}/query", diff --git a/ibis-server/tests/routers/v3/connector/bigquery/test_query.py b/ibis-server/tests/routers/v3/connector/bigquery/test_query.py index ebdd0fac8..24e005c24 100644 --- a/ibis-server/tests/routers/v3/connector/bigquery/test_query.py +++ b/ibis-server/tests/routers/v3/connector/bigquery/test_query.py @@ -103,6 +103,62 @@ async def test_query(client, manifest_str, connection_info): } +async def test_query_with_cache(client, manifest_str, connection_info): + response1 = await client.post( + url=f"{base_url}/query?cacheEnable=true", + json={ + "connectionInfo": connection_info, + "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() + + response2 = await client.post( + url=f"{base_url}/query?cacheEnable=true", + json={ + "connectionInfo": connection_info, + "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["dtypes"] == result2["dtypes"] + + +async def test_query_with_cache_override(client, manifest_str, connection_info): + response1 = await client.post( + url=f"{base_url}/query?cacheEnable=true", + json={ + "connectionInfo": connection_info, + "manifestStr": manifest_str, + "sql": "SELECT * FROM wren.public.orders LIMIT 1", + }, + ) + assert response1.status_code == 200 + + response2 = await client.post( + url=f"{base_url}/query?cacheEnable=true&overrideCache=true", + json={ + "connectionInfo": connection_info, + "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",