feat(substitute): add default header to model substitute

ibis-server/app/mdl/substitute.py

@@ -10,10 +10,11 @@
 
 
 class ModelSubstitute:
-    def __init__(self, data_source: DataSource, manifest_str: str):
+    def __init__(self, data_source: DataSource, manifest_str: str, headers=None):
         self.data_source = data_source
         self.manifest = base64_to_dict(manifest_str)
         self.model_dict = self._build_model_dict(self.manifest["models"])
+        self.headers = dict(headers) if headers else None
 
     @tracer.start_as_current_span("substitute", kind=trace.SpanKind.INTERNAL)
     def substitute(self, sql: str, write: str | None = None) -> str:
@@ -44,8 +45,15 @@ def key(model):
         return {key(model): model for model in models if "tableReference" in model}
 
     def _find_model(self, source: exp.Table) -> dict | None:
-        catalog = source.catalog or ""
-        schema = source.db or ""
+        if source.catalog and source.db:
+            catalog = source.catalog
+            schema = source.db
+        elif self.headers is not None:
+            catalog = self.headers.get("x-user-catalog")
+            schema = self.headers.get("x-user-schema")
+        else:
+            catalog = ""
+            schema = ""
         table = source.name
         return self.model_dict.get(f"{catalog}.{schema}.{table}", None)
 

ibis-server/app/routers/v2/connector.py

@@ -271,16 +271,17 @@ async def dry_plan_for_data_source(
 async def model_substitute(
     data_source: DataSource,
     dto: TranspileDTO,
+    request: Request,
     java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
     headers: Annotated[str | None, Header()] = None,
 ) -> str:
     span_name = f"v2_model_substitute_{data_source}"
     with tracer.start_as_current_span(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ):
-        sql = ModelSubstitute(data_source, dto.manifest_str).substitute(
-            dto.sql, write="trino"
-        )
+        sql = ModelSubstitute(
+            data_source, dto.manifest_str, dict(request.headers)
+        ).substitute(dto.sql, write="trino")
         Connector(data_source, dto.connection_info).dry_run(
             await Rewriter(
                 dto.manifest_str,

ibis-server/app/routers/v3/connector.py

@@ -1,6 +1,6 @@
 from typing import Annotated
 
-from fastapi import APIRouter, Depends, Header, Query, Response
+from fastapi import APIRouter, Depends, Header, Query, Request, Response
 from fastapi.responses import ORJSONResponse
 from loguru import logger
 from opentelemetry import trace
@@ -258,6 +258,7 @@ def functions(
 async def model_substitute(
     data_source: DataSource,
     dto: TranspileDTO,
+    request: Request,
     headers: Annotated[str | None, Header()] = None,
     java_engine_connector: JavaEngineConnector = Depends(get_java_engine_connector),
 ) -> str:
@@ -266,7 +267,9 @@ async def model_substitute(
         name=span_name, kind=trace.SpanKind.SERVER, context=build_context(headers)
     ):
         try:
-            sql = ModelSubstitute(data_source, dto.manifest_str).substitute(dto.sql)
+            sql = ModelSubstitute(
+                data_source, dto.manifest_str, dict(request.headers)
+            ).substitute(dto.sql)
             Connector(data_source, dto.connection_info).dry_run(
                 await Rewriter(
                     dto.manifest_str,
@@ -280,5 +283,5 @@ async def model_substitute(
                 "Failed to execute v3 model-substitute, fallback to v2: {}", str(e)
             )
             return await v2.connector.model_substitute(
-                data_source, dto, java_engine_connector, headers
+                data_source, dto, request, java_engine_connector, headers
             )

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

@@ -23,6 +23,7 @@
         {
             "name": "Orders",
             "tableReference": {
+                "catalog": "test",
                 "schema": "public",
                 "table": "orders",
             },
@@ -710,12 +711,29 @@ async def test_dry_plan(client, manifest_str):
 
 async def test_model_substitute(client, manifest_str, postgres: PostgresContainer):
     connection_info = _to_connection_info(postgres)
+    # Test with catalog and schema in SQL
     response = await client.post(
         url=f"{base_url}/model-substitute",
         json={
             "connectionInfo": connection_info,
             "manifestStr": manifest_str,
-            "sql": 'SELECT * FROM "public"."orders"',
+            "sql": 'SELECT * FROM "test"."public"."orders"',
+        },
+    )
+    assert response.status_code == 200
+    assert (
+        response.text
+        == '"SELECT * FROM \\"my_catalog\\".\\"my_schema\\".\\"Orders\\" AS \\"orders\\""'
+    )
+
+    # Test without catalog and schema in SQL but in headers(x-user-xxx)
+    response = await client.post(
+        url=f"{base_url}/model-substitute",
+        headers={"x-user-catalog": "test", "x-user-schema": "public"},
+        json={
+            "connectionInfo": connection_info,
+            "manifestStr": manifest_str,
+            "sql": 'SELECT * FROM "orders"',
         },
     )
     assert response.status_code == 200
@@ -729,14 +747,36 @@ async def test_model_substitute_with_cte(
     client, manifest_str, postgres: PostgresContainer
 ):
     connection_info = _to_connection_info(postgres)
+    # Test with catalog and schema in SQL
     response = await client.post(
         url=f"{base_url}/model-substitute",
         json={
             "connectionInfo": connection_info,
             "manifestStr": manifest_str,
             "sql": """
                 WITH orders_cte AS (
-                    SELECT * FROM "public"."orders"
+                    SELECT * FROM "test"."public"."orders"
+                )
+                SELECT * FROM orders_cte;
+            """,
+        },
+    )
+    assert response.status_code == 200
+    assert (
+        response.text
+        == '"WITH orders_cte AS (SELECT * FROM \\"my_catalog\\".\\"my_schema\\".\\"Orders\\" AS \\"orders\\") SELECT * FROM orders_cte"'
+    )
+
+    # Test without catalog and schema in SQL but in headers(x-user-xxx)
+    response = await client.post(
+        url=f"{base_url}/model-substitute",
+        headers={"x-user-catalog": "test", "x-user-schema": "public"},
+        json={
+            "connectionInfo": connection_info,
+            "manifestStr": manifest_str,
+            "sql": """
+                WITH orders_cte AS (
+                    SELECT * FROM "orders"
                 )
                 SELECT * FROM orders_cte;
             """,
@@ -753,14 +793,35 @@ async def test_model_substitute_with_subquery(
     client, manifest_str, postgres: PostgresContainer
 ):
     connection_info = _to_connection_info(postgres)
+    # Test with catalog and schema in SQL
     response = await client.post(
         url=f"{base_url}/model-substitute",
         json={
             "connectionInfo": connection_info,
             "manifestStr": manifest_str,
             "sql": """
                 SELECT * FROM (
-                    SELECT * FROM "public"."orders"
+                    SELECT * FROM "test"."public"."orders"
+                ) AS orders_subquery;
+            """,
+        },
+    )
+    assert response.status_code == 200
+    assert (
+        response.text
+        == '"SELECT * FROM (SELECT * FROM \\"my_catalog\\".\\"my_schema\\".\\"Orders\\" AS \\"orders\\") AS orders_subquery"'
+    )
+
+    # Test without catalog and schema in SQL but in headers(x-user-xxx)
+    response = await client.post(
+        url=f"{base_url}/model-substitute",
+        headers={"x-user-catalog": "test", "x-user-schema": "public"},
+        json={
+            "connectionInfo": connection_info,
+            "manifestStr": manifest_str,
+            "sql": """
+                SELECT * FROM (
+                    SELECT * FROM "orders"
                 ) AS orders_subquery;
             """,
         },
@@ -776,8 +837,22 @@ async def test_model_substitute_out_of_scope(
     client, manifest_str, postgres: PostgresContainer
 ):
     connection_info = _to_connection_info(postgres)
+    # Test with catalog and schema in SQL
+    response = await client.post(
+        url=f"{base_url}/model-substitute",
+        json={
+            "connectionInfo": connection_info,
+            "manifestStr": manifest_str,
+            "sql": 'SELECT * FROM "Nation" LIMIT 1',
+        },
+    )
+    assert response.status_code == 422
+    assert response.text == 'Model not found: "Nation"'
+
+    # Test without catalog and schema in SQL but in headers(x-user-xxx)
     response = await client.post(
         url=f"{base_url}/model-substitute",
+        headers={"x-user-catalog": "test", "x-user-schema": "public"},
         json={
             "connectionInfo": connection_info,
             "manifestStr": manifest_str,
@@ -792,12 +867,26 @@ async def test_model_substitute_non_existent_column(
     client, manifest_str, postgres: PostgresContainer
 ):
     connection_info = _to_connection_info(postgres)
+    # Test with catalog and schema in SQL
+    response = await client.post(
+        url=f"{base_url}/model-substitute",
+        json={
+            "connectionInfo": connection_info,
+            "manifestStr": manifest_str,
+            "sql": 'SELECT x FROM "test"."public"."orders" LIMIT 1',
+        },
+    )
+    assert response.status_code == 422
+    assert 'column "x" does not exist' in response.text
+
+    # Test without catalog and schema in SQL but in headers(x-user-xxx)
     response = await client.post(
         url=f"{base_url}/model-substitute",
+        headers={"x-user-catalog": "test", "x-user-schema": "public"},
         json={
             "connectionInfo": connection_info,
             "manifestStr": manifest_str,
-            "sql": 'SELECT x FROM "public"."orders" LIMIT 1',
+            "sql": 'SELECT x FROM "orders" LIMIT 1',
         },
     )
     assert response.status_code == 422