feat(ibis): query cache with bigquery info

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"]
     )

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()
 

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",

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",