fix: send correct username to consumer trino (#6629)

* fix: send correct username to consumer trino

* Apply suggestion from @ravenac95

Author: ravenac95

warehouse/scheduler/scheduler/dlt_destination.py

@@ -14,7 +14,9 @@ class DLTDestinationResource(abc.ABC):
 
     @abc.abstractmethod
     @asynccontextmanager
-    def get_destination(self, *, dataset_schema: str) -> t.AsyncIterator[Destination]:
+    def get_destination(
+        self, *, dataset_schema: str, user: str | None = None
+    ) -> t.AsyncIterator[Destination]:
         """Get a DLT destination configured for the given dataset schema."""
         ...
 
@@ -27,7 +29,7 @@ def __init__(self, *, database_path: str) -> None:
 
     @asynccontextmanager
     async def get_destination(
-        self, *, dataset_schema: str
+        self, *, dataset_schema: str, user: str | None = None
     ) -> t.AsyncIterator[Destination]:
         from dlt.destinations import duckdb
         from dlt.destinations.impl.duckdb.configuration import DuckDbCredentials
@@ -49,7 +51,7 @@ def __init__(self, *, trino: TrinoResource, catalog: str) -> None:
 
     @asynccontextmanager
     async def get_destination(
-        self, *, dataset_schema: str
+        self, *, dataset_schema: str, user: str | None = None
     ) -> t.AsyncIterator[Destination]:
         from dlt.destinations import sqlalchemy
 
@@ -58,7 +60,7 @@ async def get_destination(
             await cursor.execute(
                 f'CREATE SCHEMA IF NOT EXISTS "{self._catalog}"."{dataset_schema}"'
             )
-            user = conn.user or "scheduler"
+            user = user or conn.user or "scheduler"
             credentials = f"trinoso://{user}@{conn.host}:{conn.port}/{self._catalog}/{dataset_schema}"
             yield sqlalchemy(
                 credentials=credentials,

warehouse/scheduler/scheduler/mq/handlers/data_ingestion.py

@@ -16,6 +16,7 @@
     HandlerResponse,
     RunContext,
 )
+from scheduler.utils import get_warehouse_user
 
 
 class DataIngestionRunRequestHandler(RunHandler[DataIngestionRunRequest]):
@@ -106,6 +107,9 @@ async def handle_run_message(
                 message=f"Unsupported type: {config.factory_type}",
             )
 
+        warehouse_user = get_warehouse_user(
+            user_type="rw", org_id=org_id, org_name=context.organization.name
+        )
         async with context.step_context(
             name="execute_data_ingestion_pipeline",
             display_name="Execute Data Ingestion Pipeline",
@@ -116,6 +120,7 @@ async def handle_run_message(
                 config=config.config,
                 dataset_id=dataset_id,
                 org_id=org_id,
+                destination_user=warehouse_user,
                 dlt_destination=dlt_destination,
                 common_settings=common_settings,
             )

warehouse/scheduler/scheduler/mq/handlers/data_model.py

@@ -25,7 +25,7 @@
     StepContext,
     SuccessResponse,
 )
-from scheduler.utils import OSOClientTableResolver, ctas_query, get_trino_user
+from scheduler.utils import OSOClientTableResolver, ctas_query, get_warehouse_user
 from sqlglot import exp
 from sqlmesh import EngineAdapter
 
@@ -125,7 +125,7 @@ async def handle_run_message(
         dataset_name = dataset.node.name
         org_name = dataset.node.organization.name
 
-        user = get_trino_user("rw", context.organization.id, org_name)
+        user = get_warehouse_user("rw", context.organization.id, org_name)
 
         assert isinstance(
             data_model_def,

warehouse/scheduler/scheduler/mq/handlers/ingestion/archive.py

@@ -53,6 +53,7 @@ async def execute(
         config: dict[str, object],
         dataset_id: str,
         org_id: str,
+        destination_user: str,
         dlt_destination: DLTDestinationResource,
         common_settings: CommonSettings,
     ) -> HandlerResponse:
@@ -129,7 +130,8 @@ async def execute(
 
         try:
             async with dlt_destination.get_destination(
-                dataset_schema=dataset_schema
+                dataset_schema=dataset_schema,
+                user=destination_user,
             ) as destination:
                 dlt_pipeline = pipeline(
                     pipeline_name=pipeline_name,

warehouse/scheduler/scheduler/mq/handlers/ingestion/base.py

@@ -19,6 +19,7 @@ async def execute(
         config: dict[str, object],
         dataset_id: str,
         org_id: str,
+        destination_user: str,
         dlt_destination: DLTDestinationResource,
         common_settings: CommonSettings,
     ) -> HandlerResponse:

warehouse/scheduler/scheduler/mq/handlers/ingestion/rest.py

@@ -33,6 +33,7 @@ async def execute(
         config: dict[str, object],
         dataset_id: str,
         org_id: str,
+        destination_user: str,
         dlt_destination: DLTDestinationResource,
         common_settings: CommonSettings,
     ) -> HandlerResponse:
@@ -83,7 +84,8 @@ async def execute(
             dataset_schema = placeholder_target_table.db
             pipeline_name = f"{org_id}_{dataset_id}".replace("-", "")[:50]
             async with dlt_destination.get_destination(
-                dataset_schema=dataset_schema
+                dataset_schema=dataset_schema,
+                user=destination_user,
             ) as destination:
                 p = pipeline(
                     pipeline_name=pipeline_name,

warehouse/scheduler/scheduler/mq/handlers/static_model.py

@@ -16,7 +16,7 @@
     SuccessResponse,
     TableReference,
 )
-from scheduler.utils import dlt_to_oso_schema
+from scheduler.utils import dlt_to_oso_schema, get_warehouse_user
 
 if t.TYPE_CHECKING:
     from scheduler.config import CommonSettings
@@ -99,8 +99,15 @@ async def handle_run_message(
             },
         )
 
+        warehouse_user = get_warehouse_user(
+            user_type="rw",
+            org_id=org_id,
+            org_name=context.organization.name,
+        )
+
         async with dlt_destination.get_destination(
             dataset_schema=schema_name,
+            user=warehouse_user,
         ) as dlt_destination_instance:
             for model_id in message.model_ids:
                 async with context.step_context(

warehouse/scheduler/scheduler/mq/handlers/sync_connection.py

@@ -18,7 +18,7 @@
     RunContext,
     SuccessResponse,
 )
-from scheduler.utils import get_trino_user
+from scheduler.utils import get_warehouse_user
 
 EXCLUDED_SCHEMAS = {"information_schema"}
 MAX_CONCURRENT_QUERIES = 5  # Number of concurrent queries to run
@@ -139,7 +139,7 @@ async def handle_run_message(
             extra={"catalog_name": catalog_name, "org_id": organization.id},
         )
 
-        user = get_trino_user("ro", organization.id, organization.name)
+        user = get_warehouse_user("ro", organization.id, organization.name)
 
         # Use a single Trino client for all queries
         async with consumer_trino.async_get_client(user=user) as client:

warehouse/scheduler/scheduler/utils.py

@@ -18,7 +18,9 @@ def convert_uuid_bytes_to_str(uuid_bytes: bytes) -> str:
     return str(uuid.UUID(bytes=uuid_bytes))
 
 
-def get_trino_user(user_type: Literal["rw", "ro"], org_id: str, org_name: str) -> str:
+def get_warehouse_user(
+    user_type: Literal["rw", "ro"], org_id: str, org_name: str
+) -> str:
     return f"{user_type}-{org_name.strip().lower()}-{org_id.replace('-', '').lower()}"