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

relationalai/semantics/rel/executor.py

@@ -18,7 +18,8 @@ from relationalai.clients.snowflake import APP_NAME
 from relationalai.semantics.metamodel import ir, executor as e, factory as f
 from relationalai.semantics.rel import Compiler
 from relationalai.clients.config import Config
-from relationalai.tools.constants import USE_DIRECT_ACCESS, Generation
+from relationalai.tools.constants import USE_DIRECT_ACCESS, Generation, QUERY_ATTRIBUTES_HEADER
+from relationalai.tools.query_utils import prepare_metadata_for_headers
 
 
 class RelExecutor(e.Executor):
@@ -57,10 +58,10 @@ class RelExecutor(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, "rel")
         return self._resources
 
-    def check_graph_index(self):
+    def check_graph_index(self, headers: dict[str, Any] | None = None):
         # Has to happen first, so self.dry_run is populated.
         resources = self.resources
 
@@ -84,7 +85,16 @@ class RelExecutor(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="rel",
+                program_span_id=program_span_id,
+                headers=headers,
+            )
 
     def report_errors(self, problems: list[dict[str, Any]], abort_on_error=True):
         from relationalai import errors
@@ -143,7 +153,7 @@ class RelExecutor(e.Executor):
         elif len(all_errors) > 1:
             raise errors.RAIExceptionSet(all_errors)
 
-    def _export(self, raw_code: str, dest_fqn: str, actual_cols: list[str], declared_cols: list[str], update:bool):
+    def _export(self, raw_code: str, dest_fqn: str, actual_cols: list[str], declared_cols: list[str], update:bool, headers: dict[str, Any] | None = None):
         _exec = self.resources._exec
         output_table = "out" + str(uuid.uuid4()).replace("-", "_")
         txn_id = None
@@ -153,7 +163,7 @@ class RelExecutor(e.Executor):
         with debugging.span("transaction"):
             try:
                 with debugging.span("exec_format") as span:
-                    res = _exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, ?);", [self.database, self.engine, raw_code, output_table, False, True])
+                    res = _exec(f"call {APP_NAME}.api.exec_into_table(?, ?, ?, ?, ?, NULL, ?, {headers}, ?, ?);", [self.database, self.engine, raw_code, output_table, False, True, False, None])
                     txn_id = json.loads(res[0]["EXEC_INTO_TABLE"])["rai_transaction_id"]
                     span["txn_id"] = txn_id
 
@@ -229,14 +239,18 @@ class RelExecutor(e.Executor):
                 else:
                     raise e
             if txn_id:
-                artifact_info = self.resources._list_exec_async_artifacts(txn_id)
+                artifact_info = self.resources._list_exec_async_artifacts(txn_id, headers=headers)
                 with debugging.span("fetch"):
                     artifacts = self.resources._download_results(artifact_info, txn_id, "ABORTED")
             return artifacts
 
     def execute(self, model: ir.Model, task: ir.Task, format: Literal["pandas", "snowpark"] = "pandas",
-                result_cols: list[str] | None = None, export_to: Optional[str] = None, update: bool = False) -> Any:
-        self.check_graph_index()
+                result_cols: list[str] | None = None, export_to: Optional[str] = None, update: bool = False, meta: dict[str, Any] | None = None) -> Any:
+        # Format meta as headers
+        json_meta = prepare_metadata_for_headers(meta)
+        headers = {QUERY_ATTRIBUTES_HEADER: json_meta} if json_meta else {}
+        
+        self.check_graph_index(headers)
         resources= self.resources
 
         rules_code = ""
@@ -277,12 +291,16 @@ class RelExecutor(e.Executor):
 
         if not export_to:
             if format == "pandas":
-                raw_results = resources.exec_raw(self.database, self.engine, full_code, False, nowait_durable=True)
+                raw_results = resources.exec_raw(self.database, self.engine, full_code, False, nowait_durable=True, headers=headers)
                 df, errs = result_helpers.format_results(raw_results, None, cols, generation=Generation.QB)  # Pass None for task parameter
                 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)]
+
                 return self._postprocess_df(self.config, df, extra_cols)
             elif format == "snowpark":
-                results, raw = resources.exec_format(self.database, self.engine, full_code, cols, format=format, readonly=False, nowait_durable=True)
+                results, raw = resources.exec_format(self.database, self.engine, full_code, cols, format=format, readonly=False, nowait_durable=True, headers=headers)
                 if raw:
                     df, errs = result_helpers.format_results(raw, None, cols, generation=Generation.QB)  # Pass None for task parameter
                     self.report_errors(errs)
@@ -296,7 +314,7 @@ class RelExecutor(e.Executor):
             else:
                 result_cols = [col for col in cols if col not in extra_cols]
             assert result_cols
-            raw = self._export(full_code, export_to, cols, result_cols, update)
+            raw = self._export(full_code, export_to, cols, result_cols, update, headers)
             errors = []
             if raw:
                 dataframe, errors = result_helpers.format_results(raw, None, result_cols, generation=Generation.QB)

relationalai/semantics/sql/executor/snowflake.py

@@ -62,7 +62,7 @@ class SnowflakeExecutor(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) -> Union[pd.DataFrame, Any]:
+                update: bool = False, meta: dict[str, Any] | None = None) -> Union[pd.DataFrame, Any]:
         """ Execute the SQL query directly. """
 
         warehouse = self.resources.config.get("warehouse", None)

relationalai/tools/cli.py

@@ -719,7 +719,10 @@ def config_check(all_profiles:bool=False):
 @cli.command(help="Print version info")
 def version():
     from .. import __version__
-    from railib import __version__ as railib_version
+    try:
+        from railib import __version__ as railib_version
+    except Exception:
+        railib_version = None
 
     table = Table(show_header=False, border_style="dim", header_style="bold", box=rich_box.SIMPLE)
     def print_version(name, version, latest=None):
@@ -730,7 +733,8 @@ def version():
 
     divider()
     print_version("RelationalAI", __version__, latest_version("relationalai"))
-    print_version("Rai-sdk", railib_version, latest_version("rai-sdk"))
+    if railib_version is not None:
+        print_version("Rai-sdk", railib_version, latest_version("rai-sdk"))
     print_version("Python", sys.version.split()[0])
 
     app_version = None