relationalai 0.12.0__py3-none-any.whl → 0.12.2__py3-none-any.whl

relationalai/experimental/solvers.py

@@ -9,11 +9,13 @@ from ..metamodel import Builtins
 from ..tools.cli_controls import Spinner
 from ..tools.constants import DEFAULT_QUERY_TIMEOUT_MINS
 from .. import debugging
+from .. errors import ResponseStatusException
 import uuid
 import relationalai
 import json
 from ..clients.util import poll_with_specified_overhead
 from ..clients.snowflake import Resources as SnowflakeResources
+from ..clients.snowflake import DirectAccessClient, DirectAccessResources
 from ..util.timeout import calc_remaining_timeout_minutes
 
 rel_sv = rel._tagged(Builtins.SingleValued)
@@ -23,7 +25,8 @@ APP_NAME = relationalai.clients.snowflake.APP_NAME
 ENGINE_TYPE_SOLVER = "SOLVER"
 # TODO (dba) The ERP still uses `worker` instead of `engine`. Change
 # this once we fix this in the ERP.
-ENGINE_ERRORS = ["worker is suspended", "create/resume", "worker not found", "no workers found", "worker was deleted"]
+WORKER_ERRORS = ["worker is suspended", "create/resume", "worker not found", "no workers found", "worker was deleted"]
+ENGINE_ERRORS = ["engine is suspended", "create/resume", "engine not found", "no engines found", "engine was deleted"]
 ENGINE_NOT_READY_MSGS = ["worker is in pending", "worker is provisioning", "worker is not ready to accept jobs"]
 
 # --------------------------------------------------
@@ -213,13 +216,6 @@ class SolverModel:
         config_file_path = getattr(rai_config, 'file_path', None)
         start_time = time.monotonic()
         remaining_timeout_minutes = query_timeout_mins
-        # 1. Materialize the model and store it.
-        # TODO(coey) Currently we must run a dummy query to install the pyrel rules in a separate txn
-        # to the solve_output updates. Ideally pyrel would offer an option to flush the rules separately.
-        self.graph.exec_raw("", query_timeout_mins=remaining_timeout_minutes)
-        remaining_timeout_minutes = calc_remaining_timeout_minutes(
-            start_time, query_timeout_mins, config_file_path=config_file_path,
-        )
         response = self.graph.exec_raw(
             textwrap.dedent(f"""
             @inline
@@ -269,7 +265,9 @@ class SolverModel:
             job_id = solver._exec_job(payload, log_to_console=log_to_console, query_timeout_mins=remaining_timeout_minutes)
         except Exception as e:
             err_message = str(e).lower()
-            if any(kw in err_message.lower() for kw in ENGINE_ERRORS + ENGINE_NOT_READY_MSGS):
+            if isinstance(e, ResponseStatusException):
+                err_message = e.response.json().get("message", "")
+            if any(kw in err_message.lower() for kw in ENGINE_ERRORS + WORKER_ERRORS + ENGINE_NOT_READY_MSGS):
                 solver._auto_create_solver_async()
                 remaining_timeout_minutes = calc_remaining_timeout_minutes(
                     start_time, query_timeout_mins, config_file_path=config_file_path
@@ -553,7 +551,11 @@ class Solver:
         # may configure each individual solver.
         self.engine_settings = settings
 
-        return self._auto_create_solver_async()
+        # Optimistically set the engine object to a `READY` engine to
+        # avoid checking the engine status on each execution.
+        self.engine:Optional[dict[str,Any]] = {"name": engine_name, "state": "READY"}
+
+        return None
 
     # --------------------------------------------------
     # Helper
@@ -572,6 +574,7 @@ class Solver:
             assert len(engines) == 1 or len(engines) == 0
             if len(engines) != 0:
                 engine = engines[0]
+
             if engine:
                 # TODO (dba) Logic engines support altering the
                 # auto_suspend_mins setting. Currently, we don't have
@@ -653,31 +656,20 @@ class Solver:
 
             self.engine = engine
 
-    def _exec_job_async(self, payload, query_timeout_mins: Optional[int]=None):
-        payload_json = json.dumps(payload)
-        engine_name = self.engine["name"]
-        if query_timeout_mins is None and (timeout_value := self.rai_config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
-            query_timeout_mins = int(timeout_value)
-        if query_timeout_mins is not None:
-            sql_string = textwrap.dedent(f"""
-            CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
-            """)
-        else:
-            sql_string = textwrap.dedent(f"""
-            CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
-            """)
-        res = self.provider.resources._exec(sql_string)
-        return res[0]["ID"]
-
     def _exec_job(self, payload, log_to_console=True, query_timeout_mins: Optional[int]=None):
+        if self.engine is None:
+            raise Exception("Engine not initialized.")
+
         # Make sure the engine is ready.
         if self.engine["state"] != "READY":
             poll_with_specified_overhead(lambda: self._is_solver_ready(), 0.1)
 
         with debugging.span("job") as job_span:
-            job_id = self._exec_job_async(payload, query_timeout_mins=query_timeout_mins)
+            job_id = self.provider.create_job_async(self.engine["name"], payload, query_timeout_mins=query_timeout_mins)
             job_span["job_id"] = job_id
             debugging.event("job_created", job_span, job_id=job_id, engine_name=self.engine["name"], job_type=ENGINE_TYPE_SOLVER)
+            if not isinstance(job_id, str):
+                job_id = ""
             polling_state = PollingState(job_id, "", False, log_to_console)
 
             try:
@@ -693,7 +685,14 @@ class Solver:
             return job_id
 
     def _is_solver_ready(self):
+        if self.engine is None:
+            raise Exception("Engine not initialized.")
+
         result = self.provider.get_solver(self.engine["name"])
+
+        if result is None:
+            raise Exception("No engine available.")
+            
         self.engine = result
         state = result["state"]
         if state != "READY" and state != "PENDING":
@@ -711,20 +710,11 @@ class Solver:
 
         return status == "COMPLETED" or status == "FAILED" or status == "CANCELED"
 
-    def _get_job_events(self, job_id: str, continuation_token: str = ""):
-        results = self.provider.resources._exec(
-            f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');"
-        )
-        if not results:
-            return {"events": [], "continuation_token": None}
-        row = results[0][0]
-        return json.loads(row)
-
     def _print_solver_logs(self, state: PollingState):
         if state.is_done:
             return
 
-        resp = self._get_job_events(state.job_id, state.continuation_token)
+        resp = self.provider.get_job_events(state.job_id, state.continuation_token)
 
         # Print solver logs to stdout.
         for event in resp["events"]:
@@ -754,7 +744,12 @@ class Provider:
             resources = relationalai.Resources()
         if not isinstance(resources, relationalai.clients.snowflake.Resources):
             raise Exception("Solvers are only supported on SPCS.")
+
         self.resources = resources
+        self.direct_access_client: Optional[DirectAccessClient] = None
+
+        if isinstance(self.resources, DirectAccessResources):
+            self.direct_access_client = self.resources.direct_access_client
 
     def create_solver(
         self,
@@ -770,75 +765,285 @@ class Provider:
         engine_config: dict[str, Any] = {"settings": settings}
         if auto_suspend_mins is not None:
             engine_config["auto_suspend_mins"] = auto_suspend_mins
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});"
+        self.resources._exec_sql(
+            f"CALL {APP_NAME}.experimental.create_engine('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});", None
         )
 
     def create_solver_async(
         self,
         name: str,
         size: str | None = None,
-        settings: dict = {},
+        settings: dict | None = None,
         auto_suspend_mins: int | None = None,
     ):
         if size is None:
             size = "HIGHMEM_X64_S"
-        if settings is None:
-            settings = ""
-        engine_config: dict[str, Any] = {"settings": settings}
-        if auto_suspend_mins is not None:
-            engine_config["auto_suspend_mins"] = auto_suspend_mins
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});"
-        )
+
+        if self.direct_access_client is not None:
+            payload:dict[str, Any] = {
+                "name": name,
+                "settings": settings,
+               }
+            if auto_suspend_mins is not None:
+                payload["auto_suspend_mins"] = auto_suspend_mins
+            if size is not None:
+                payload["size"] = size
+            response = self.direct_access_client.request(
+                "create_engine",
+                payload=payload,
+                path_params={"engine_type": "solver"},
+               )
+            if response.status_code != 200:
+                raise ResponseStatusException(
+                    f"Failed to create engine {name} with size {size}.", response
+                )
+        else:
+            engine_config: dict[str, Any] = {}
+            if settings is not None:
+                engine_config["settings"] = settings
+            if auto_suspend_mins is not None:
+                engine_config["auto_suspend_mins"] = auto_suspend_mins
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.create_engine_async('{ENGINE_TYPE_SOLVER}', '{name}', '{size}', {engine_config});",
+                None
+               )
 
     def delete_solver(self, name: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "delete_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to delete engine.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.delete_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+            )
 
     def resume_solver_async(self, name: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "resume_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to resume engine.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.resume_engine_async('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+            )
+            return None
 
     def get_solver(self, name: str):
-        results = self.resources._exec(
-            f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');"
-        )
-        return solver_list_to_dicts(results)[0]
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_engine", path_params = {"engine_type": ENGINE_TYPE_SOLVER, "engine_name": name}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get engine.", response)
+            solver = response.json()
+            if not solver :
+                return None
+            solver_state = {
+                "name": solver["name"],
+                "id": solver["id"],
+                "size": solver["size"],
+                "state": solver["status"], # callers are expecting 'state'
+                "created_by": solver["created_by"],
+                "created_on": solver["created_on"],
+                "updated_on": solver["updated_on"],
+                "version": solver["version"],
+                "auto_suspend": solver["auto_suspend_mins"],
+                "suspends_at": solver["suspends_at"],
+                "solvers": []
+                if solver["settings"] == ""
+                else [
+                        k
+                        for (k,v) in json.loads(solver["settings"]).items()
+                        if isinstance(v, dict) and v.get("enabled", False)
+                ],
+                }
+            return solver_state
+        else:
+            results = self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.get_engine('{ENGINE_TYPE_SOLVER}', '{name}');",
+                None
+               )
+            return solver_list_to_dicts(results)[0]
 
     def list_solvers(self, state: str | None = None):
-        where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
-        where_clause = (
-            f"{where_clause} AND STATUS = '{state.upper()}'" if state else where_clause
-        )
-        statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
-        results = self.resources._exec(statement)
-        return solver_list_to_dicts(results)
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "list_engines"
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to list engines.", response)
+            response_content = response.json()
+            if not response_content:
+                return []
+            engines = [
+                {
+                    "name": engine["name"],
+                    "id": engine["id"],
+                    "size": engine["size"],
+                    "state": engine["status"], # callers are expecting 'state'
+                    "created_by": engine["created_by"],
+                    "created_on": engine["created_on"],
+                    "updated_on": engine["updated_on"],
+                    "auto_suspend_mins": engine["auto_suspend_mins"],
+                    "solvers": []
+                    if engine["settings"] == ""
+                    else [
+                        k
+                        for (k, v) in json.loads(engine["settings"]).items()
+                        if isinstance(v, dict) and v.get("enabled", False)
+                           ],
+                }
+                for engine in response_content.get("engines", [])
+                if (state is None or engine.get("status") == state) and (engine.get("type") == ENGINE_TYPE_SOLVER)
+               ]
+            return sorted(engines, key=lambda x: x["name"])
+        else:
+            where_clause = f"WHERE TYPE='{ENGINE_TYPE_SOLVER}'"
+            where_clause = (
+                f"{where_clause} AND STATUS = '{state.upper()}'" if state else where_clause
+               )
+            statement = f"SELECT NAME,ID,SIZE,STATUS,CREATED_BY,CREATED_ON,UPDATED_ON,AUTO_SUSPEND_MINS,SETTINGS FROM {APP_NAME}.experimental.engines {where_clause};"
+            results = self.resources._exec_sql(statement, None)
+            return solver_list_to_dicts(results)
 
     # --------------------------------------------------
     # Job API
     # --------------------------------------------------
 
+    def create_job_async(self, engine_name, payload, query_timeout_mins: Optional[int]=None):
+        payload_json = json.dumps(payload)
+
+        if query_timeout_mins is None and (timeout_value := self.resources.config.get("query_timeout_mins", DEFAULT_QUERY_TIMEOUT_MINS)) is not None:
+            query_timeout_mins = int(timeout_value)
+            
+        if self.direct_access_client is not None:
+            job = {
+                "job_type":ENGINE_TYPE_SOLVER,
+                "worker_name": engine_name,
+                "timeout_mins": query_timeout_mins,
+                "payload": payload_json,
+            }
+            response = self.direct_access_client.request(
+                "create_job",
+                payload=job,
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to create job.", response)
+            response_content = response.json()
+            return response_content["id"]
+        else:
+            if query_timeout_mins is not None:
+                sql_string = textwrap.dedent(f"""
+                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}', null, {query_timeout_mins})
+                """)
+            else:
+                sql_string = textwrap.dedent(f"""
+                CALL {APP_NAME}.experimental.exec_job_async('{ENGINE_TYPE_SOLVER}', '{engine_name}', '{payload_json}')
+                """)
+            res = self.resources._exec_sql(sql_string, None)
+            return res[0]["ID"]
+
     def list_jobs(self, state=None, limit=None):
-        state_clause = f"AND STATE = '{state.upper()}'" if state else ""
-        limit_clause = f"LIMIT {limit}" if limit else ""
-        results = self.resources._exec(
-            f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};"
-        )
-        return job_list_to_dicts(results)
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "list_jobs"
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to list jobs.", response)
+            response_content = response.json()
+            if not response_content:
+                return []
+            jobs = [
+                {
+                    "id": job["id"],
+                    "state": job["state"],
+                    "created_by": job["created_by"],
+                    "created_on": job["created_on"],
+                    "finished_at": job.get("finished_at", None),
+                    "duration": job["duration"] if "duration" in job else 0,
+                    "solver": json.loads(job["payload"]).get("solver", ""),
+                    "engine": job.get("engine_name", job["worker_name"]),
+                }
+                for job in response_content.get("jobs", [])
+                if state is None or job.get("state") == state
+               ]
+            return sorted(jobs, key=lambda x: x["created_on"], reverse=True)
+        else:
+            state_clause = f"AND STATE = '{state.upper()}'" if state else ""
+            limit_clause = f"LIMIT {limit}" if limit else ""
+            results = self.resources._exec_sql(
+                f"SELECT ID,STATE,CREATED_BY,CREATED_ON,FINISHED_AT,DURATION,PAYLOAD,ENGINE_NAME FROM {APP_NAME}.experimental.jobs where type='{ENGINE_TYPE_SOLVER}' {state_clause} ORDER BY created_on DESC {limit_clause};",
+                None
+               )
+            return job_list_to_dicts(results)
 
     def get_job(self, id: str):
-        results = self.resources._exec(
-            f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');"
-        )
-        return job_list_to_dicts(results)[0]
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get job.", response)
+            response_content = response.json()
+            return response_content["job"]
+        else:
+            results = self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.get_job('{ENGINE_TYPE_SOLVER}', '{id}');",
+                None
+            )
+            return job_list_to_dicts(results)[0]
+
+    def get_job_events(self, job_id: str, continuation_token: str = ""):
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "get_job_events",
+                path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": job_id, "stream_name": "progress"},
+                query_params={"continuation_token": continuation_token},
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to get job events.", response)
+            response_content = response.json()
+            if not response_content:
+                return {
+                    "events": [],
+                    "continuation_token": None
+                }
+            return response_content
+        else:
+            results = self.resources._exec_sql(
+                f"SELECT {APP_NAME}.experimental.get_job_events('{ENGINE_TYPE_SOLVER}', '{job_id}', '{continuation_token}');",
+                None
+               )
+            if not results:
+                return {"events": [], "continuation_token": None}
+            row = results[0][0]
+            if not isinstance(row, str):
+                row = ""
+            return json.loads(row)
 
     def cancel_job(self, id: str):
-        self.resources._exec(
-            f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');"
-        )
+        if self.direct_access_client is not None:
+            response = self.direct_access_client.request(
+                "cancel_job", path_params = {"job_type": ENGINE_TYPE_SOLVER, "job_id": id}
+            )
+            if response.status_code != 200:
+                raise ResponseStatusException("Failed to cancel job.", response)
+            return None
+        else:
+            self.resources._exec_sql(
+                f"CALL {APP_NAME}.experimental.cancel_job('{ENGINE_TYPE_SOLVER}', '{id}');",
+                None
+            )
+            return None
 
 
 def solver_list_to_dicts(results):
@@ -865,7 +1070,6 @@ def solver_list_to_dicts(results):
         for row in results
     ]
 
-
 def job_list_to_dicts(results):
     if not results:
         return []

relationalai/semantics/internal/internal.py

@@ -346,7 +346,7 @@ def find_select_keys(item: Any, keys:OrderedSet[Key]|None = None, enable_primiti
 
     if isinstance(item, (list, tuple)):
         for it in item:
-            find_select_keys(it, keys)
+            find_select_keys(it, keys, enable_primitive_key=enable_primitive_key)
 
     elif isinstance(item, (Relationship, RelationshipReading)) and item._parent:
         find_select_keys(item._parent, keys)
@@ -390,7 +390,7 @@ def find_select_keys(item: Any, keys:OrderedSet[Key]|None = None, enable_primiti
         find_select_keys(item._arg, keys)
 
     elif isinstance(item, Alias):
-        find_select_keys(item._thing, keys)
+        find_select_keys(item._thing, keys, enable_primitive_key=enable_primitive_key)
 
     elif isinstance(item, Aggregate):
         keys.update( Key(k, True) for k in item._group )
@@ -2338,6 +2338,7 @@ class Fragment():
             self._define.extend(parent._define)
             self._order_by.extend(parent._order_by)
             self._limit = parent._limit
+            self._meta.update(parent._meta)
 
     def _add_items(self, items:PySequence[Any], to_attr:list[Any]):
         # TODO: ensure that you are _either_ a select, require, or then
@@ -2416,9 +2417,26 @@ class Fragment():
         return f
 
     def meta(self, **kwargs: Any) -> Fragment:
+        """Add metadata to the query.
+
+        Metadata can be used for debugging and observability purposes.
+
+        Args:
+            **kwargs: Metadata key-value pairs
+
+        Returns:
+            Fragment: Returns self for method chaining
+
+        Example:
+            select(Person.name).meta(workload_name="test", priority=1, enabled=True)
+        """
+        if not kwargs:
+            raise ValueError("meta() requires at least one argument")
+
         self._meta.update(kwargs)
         return self
 
+
     def annotate(self, *annos:Expression|Relationship|ir.Annotation) -> Fragment:
         self._annotations.extend(annos)
         return self
@@ -2497,7 +2515,7 @@ class Fragment():
         # @TODO for now we set tag to None but we need to work out how to properly propagate user-provided tag here
         with debugging.span("query", tag=None, dsl=str(self), **with_source(self), meta=self._meta) as query_span:
             query_task = qb_model._compiler.fragment(self)
-            results = qb_model._to_executor().execute(ir_model, query_task)
+            results = qb_model._to_executor().execute(ir_model, query_task, meta=self._meta)
             query_span["results"] = results
             # For local debugging mostly
             dry_run = qb_model._dry_run or bool(qb_model._config.get("compiler.dry_run", False))
@@ -2524,7 +2542,7 @@ class Fragment():
         # @TODO for now we set tag to None but we need to work out how to properly propagate user-provided tag here
         with debugging.span("query", tag=None, dsl=str(clone), **with_source(clone), meta=clone._meta) as query_span:
             query_task = qb_model._compiler.fragment(clone)
-            results = qb_model._to_executor().execute(ir_model, query_task, format="snowpark")
+            results = qb_model._to_executor().execute(ir_model, query_task, format="snowpark", meta=clone._meta)
             query_span["alt_format_results"] = results
             return results
 
@@ -2541,7 +2559,8 @@ class Fragment():
         clone._source = runtime_env.get_source_pos()
         with debugging.span("query", dsl=str(clone), **with_source(clone), meta=clone._meta):
             query_task = qb_model._compiler.fragment(clone)
-            qb_model._to_executor().execute(ir_model, query_task, result_cols=result_cols, export_to=table._fqn, update=update)
+            qb_model._to_executor().execute(ir_model, query_task, result_cols=result_cols, export_to=table._fqn, update=update, meta=clone._meta)
+
 
 #--------------------------------------------------
 # Select / Where

relationalai/semantics/lqp/executor.py

@@ -21,8 +21,8 @@ from relationalai.clients.config import Config
 from relationalai.clients.snowflake import APP_NAME
 from relationalai.clients.types import TransactionAsyncResponse
 from relationalai.clients.util import IdentityParser
-from relationalai.tools.constants import USE_DIRECT_ACCESS
-
+from relationalai.tools.constants import USE_DIRECT_ACCESS, QUERY_ATTRIBUTES_HEADER
+from relationalai.tools.query_utils import prepare_metadata_for_headers
 
 class LQPExecutor(e.Executor):
     """Executes LQP using the RAI client."""
@@ -60,7 +60,7 @@ class LQPExecutor(e.Executor):
                 if not self.dry_run:
                     self.engine = self._resources.get_default_engine_name()
                     if not self.keep_model:
-                        atexit.register(self._resources.delete_graph, self.database, True)
+                        atexit.register(self._resources.delete_graph, self.database, True, "lqp")
         return self._resources
 
     # Checks the graph index and updates it if necessary
@@ -88,7 +88,15 @@ class LQPExecutor(e.Executor):
         assert self.engine is not None
 
         with debugging.span("poll_use_index", sources=sources, model=model, engine=engine_name):
-            resources.poll_use_index(app_name, sources, model, self.engine, engine_size, program_span_id)
+            resources.poll_use_index(
+                app_name=app_name,
+                sources=sources,
+                model=model,
+                engine_name=self.engine,
+                engine_size=engine_size,
+                language="lqp",
+                program_span_id=program_span_id,
+            )
 
     def report_errors(self, problems: list[dict[str, Any]], abort_on_error=True):
         from relationalai import errors
@@ -267,7 +275,7 @@ class LQPExecutor(e.Executor):
         if ivm_flag:
             config_dict['ivm.maintenance_level'] = lqp_ir.Value(value=ivm_flag, meta=None)
         return construct_configure(config_dict, None)
-    
+
     def _compile_intrinsics(self) -> lqp_ir.Epoch:
         """Construct an epoch that defines a number of built-in definitions used by the
         emitter."""
@@ -344,6 +352,10 @@ class LQPExecutor(e.Executor):
         df, errs = result_helpers.format_results(raw_results, cols)
         self.report_errors(errs)
 
+        # Rename columns if wide outputs is enabled
+        if self.wide_outputs and len(cols) - len(extra_cols) == len(df.columns):
+            df.columns = cols[: len(df.columns)]
+
         # Process exports
         if export_to and not self.dry_run:
             assert cols, "No columns found in the output"
@@ -362,7 +374,7 @@ class LQPExecutor(e.Executor):
 
     def execute(self, model: ir.Model, task: ir.Task, format: Literal["pandas", "snowpark"] = "pandas",
                 result_cols: Optional[list[str]] = None, export_to: Optional[str] = None,
-                update: bool = False) -> DataFrame:
+                update: bool = False, meta: dict[str, Any] | None = None) -> DataFrame:
         self.prepare_data()
         previous_model = self._last_model
 
@@ -374,6 +386,9 @@ class LQPExecutor(e.Executor):
         if format != "pandas":
             raise ValueError(f"Unsupported format: {format}")
 
+        # Format meta as headers
+        json_meta = prepare_metadata_for_headers(meta)
+        headers = {QUERY_ATTRIBUTES_HEADER: json_meta} if json_meta else {}
         raw_results = self.resources.exec_lqp(
             self.database,
             self.engine,
@@ -383,6 +398,7 @@ class LQPExecutor(e.Executor):
             # transactions are serialized.
             readonly=False,
             nowait_durable=True,
+            headers=headers,
         )
         assert isinstance(raw_results, TransactionAsyncResponse)