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

relationalai/environments/snowbook.py

@@ -20,7 +20,16 @@ class SnowbookEnvironment(NotebookRuntimeEnvironment, SessionEnvironment):
 
     def __init__(self):
         super().__init__()
-        self.runner = "container" if "snowflake.connector.auth" in sys.modules else "warehouse"
+        # Detect runner type based on module presence:
+        # - Warehouse runtime has '_snowflake' module
+        # - Container runtime has 'snowflake._legacy' module
+        if "_snowflake" in sys.modules:
+            self.runner = "warehouse"
+        elif "snowflake._legacy" in sys.modules:
+            self.runner = "container"
+        else:
+            # Fallback to original check
+            self.runner = "container" if "snowflake.connector.auth" in sys.modules else "warehouse"
 
     @classmethod
     def detect(cls):

relationalai/semantics/internal/internal.py

@@ -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."""
@@ -267,7 +267,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 +344,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 +366,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 +378,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 +390,7 @@ class LQPExecutor(e.Executor):
             # transactions are serialized.
             readonly=False,
             nowait_durable=True,
+            headers=headers,
         )
         assert isinstance(raw_results, TransactionAsyncResponse)
 

relationalai/semantics/lqp/model2lqp.py

@@ -102,6 +102,7 @@ def _get_export_reads(export_ids: list[tuple[lqp.RelationId, int, lqp.Type]]) ->
         data_columns=csv_columns,
         compression="gzip",
         partition_size=200,
+        syntax_escapechar='"',  # To follow Snowflake's expected format
         meta=None,
     )
     reads.append(lqp.Read(read_type=lqp.Export(config=export_csv_config, meta=None), meta=None))

relationalai/semantics/metamodel/executor.py

@@ -56,5 +56,6 @@ class Executor():
             rich.print(f"[black]{df}[/black]")
         if not df.empty:
             for col in extra_cols:
-                df = df.drop(col, axis=1)
+                if col in df.columns:
+                    df = df.drop(col, axis=1)
         return df

relationalai/semantics/metamodel/rewrite/flatten.py

@@ -558,22 +558,23 @@ class Flatten(Pass):
     def rewrite_wide_output(self, output: ir.Output):
         assert(output.keys)
 
-        # only prefix keys that are not already in the output
-        prefix_keys = []
+        # only append keys that are not already in the output
+        suffix_keys = []
         for key in output.keys:
             if all([val is not key for _, val in output.aliases]):
-                prefix_keys.append(key)
+                suffix_keys.append(key)
 
         aliases: OrderedSet[Tuple[str, ir.Value]] = ordered_set()
-        # add the keys to the output
-        for key in prefix_keys:
-            aliases.add((key.name, key))
 
         # add the remaining args, unless it is already a key
         for name, val in output.aliases:
-            if not isinstance(val, ir.Var) or val not in prefix_keys:
+            if not isinstance(val, ir.Var) or val not in suffix_keys:
                 aliases.add((name, val))
 
+        # add the keys to the output
+        for key in suffix_keys:
+            aliases.add((key.name, key))
+
         # TODO - we are assuming that the Rel compiler will translate nullable lookups
         # properly, returning a `Missing` if necessary, like this:
         # (nested_192(_adult, _adult_name) or (not nested_192(_adult, _) and _adult_name = Missing)) and