feat: add auth_header to TransformContext and support ctx in lighweight (#100)

Author: nicornk

libs/transforms/src/transforms/api/_transform.py

@@ -113,7 +113,7 @@ def _compute_spark(
         """
         kwargs = {name: i.init_input(context).dataframe() for name, i in self.inputs.items()}
         if self._use_context:
-            kwargs["ctx"] = TransformContext()
+            kwargs["ctx"] = TransformContext(context)
 
         output_df = self(**kwargs)
 
@@ -136,7 +136,7 @@ def _compute_pandas(
         """
         kwargs = {name: i.init_input(context).dataframe().toPandas() for name, i in self.inputs.items()}
         if self._use_context:
-            kwargs["ctx"] = TransformContext()
+            kwargs["ctx"] = TransformContext(context)
 
         output_df = self(**kwargs)
         from foundry_dev_tools._optional.pandas import pd
@@ -159,7 +159,7 @@ def _compute_transform(  # noqa: ANN202
         kwargs = {**inputs, **outputs}
 
         if self._use_context:
-            kwargs["ctx"] = TransformContext()
+            kwargs["ctx"] = TransformContext(context)
 
         self(**kwargs)
 
@@ -169,17 +169,16 @@ def _compute_lightweight(  # noqa: ANN202
         self,
         context: FoundryContext,
     ):
-        if self._use_context:
-            msg = "Lightweight transforms do not support the context argument."
-            raise ValueError(msg)
-
         inputs = {argument_name: LightweightTransformInput(i, context) for argument_name, i in self.inputs.items()}
         outputs = {
             argument_name: LightweightTransformOutput(o, argument_name, context)
             for argument_name, o in self.outputs.items()
         }
         kwargs = {**inputs, **outputs}
 
+        if self._use_context:
+            kwargs["ctx"] = TransformContext(context)
+
         self(**kwargs)
 
         return {name: i.df for name, i in outputs.items()}
@@ -188,34 +187,30 @@ def _compute_lightweight_pandas(
         self,
         context: FoundryContext,
     ) -> pd.core.frame.DataFrame:
-        if self._use_context:
-            msg = "Lightweight transforms do not support the context argument."
-            raise ValueError(msg)
-
         inputs = {
             argument_name: LightweightTransformInput(i, context).pandas() for argument_name, i in self.inputs.items()
         }
+        if self._use_context:
+            inputs["ctx"] = TransformContext(context)
         return self(**inputs)
 
     def _compute_lightweight_polars(
         self,
         context: FoundryContext,
     ) -> pl.DataFrame:
-        if self._use_context:
-            msg = "Lightweight transforms do not support the context argument."
-            raise ValueError(msg)
-
         inputs = {
             argument_name: LightweightTransformInput(i, context).polars() for argument_name, i in self.inputs.items()
         }
+        if self._use_context:
+            inputs["ctx"] = TransformContext(context)
         return self(**inputs)
 
 
 class TransformContext:
     """The TransformContext is passed to the transform function if ctx is the first argument."""
 
-    def __init__(self):
-        pass
+    def __init__(self, foundry_ctx: FoundryContext):
+        self._foundry_ctx = foundry_ctx
 
     @property
     def spark_session(self) -> pyspark.sql.SparkSession:
@@ -241,6 +236,10 @@ def is_incremental(self) -> bool:
         warnings.warn("is_incremental functionality not implemented in Foundry DevTools")
         return False
 
+    @property
+    def auth_header(self) -> str:
+        return f"Bearer {self._foundry_ctx.token_provider.token}"
+
 
 class TransformInput:
     """TransformInput class, passed when using @transform decorator."""

tests/unit/test_transforms.py

@@ -275,16 +275,17 @@ def transform_me(input1, input2, out):
 
 
 def test_transform_polars_transform_two_inputs(spark_df_return_data_one, spark_df_return_data_two, transforms_context):
-    from transforms.api import Input, Output, transform_polars
+    from transforms.api import Input, Output, TransformContext, transform_polars
 
     @transform_polars(
         Output("/output/to/dataset"),
         input1=Input("/input1"),
         input2=Input("/input2"),
     )
-    def transform_me(input1, input2) -> DataFrame:
+    def transform_me(ctx, input1, input2) -> DataFrame:
         assert_frame_equal(input1.to_pandas(), spark_df_return_data_one.toPandas())
         assert_frame_equal(input2.to_pandas(), spark_df_return_data_two.toPandas())
+        assert isinstance(ctx, TransformContext)  # ctx is our TransformContext
         return input1.extend(input2)
 
     df = transform_me.compute(transforms_context)
@@ -358,20 +359,34 @@ def transform_me(input1: pd.DataFrame) -> pd.DataFrame:
 
 
 def test_transform_pandas_one_input_with_ctx(spark_df_return_data_one, transforms_context):
-    from transforms.api import Input, Output, TransformContext, transform_pandas
+    from transforms.api import Input, Output, TransformContext, lightweight, transform_pandas
+
+    expected_df = spark_df_return_data_one.toPandas()
 
     @transform_pandas(Output("/output/to/dataset"), input1=Input("/input1"))
     def transform_me(ctx, input1: pd.DataFrame) -> pd.DataFrame:
         assert isinstance(input1, pd.DataFrame)
         assert isinstance(ctx, TransformContext)  # ctx is our TransformContext
         assert isinstance(ctx.spark_session, SparkSession)  # ctx.spark_session is a SparkSession
-        assert_frame_equal(input1, spark_df_return_data_one.toPandas())
+        assert_frame_equal(input1, expected_df)
         return input1
 
     df = transform_me.compute(transforms_context)
     assert isinstance(df, pd.DataFrame)
     assert_frame_equal(df, spark_df_return_data_one.toPandas())
 
+    @lightweight
+    @transform_pandas(Output("/output/to/dataset"), input1=Input("/input1"))
+    def transform_me(ctx, input1: pd.DataFrame) -> pd.DataFrame:
+        assert isinstance(input1, pd.DataFrame)
+        assert isinstance(ctx, TransformContext)  # ctx is our TransformContext
+        assert_frame_equal(input1, expected_df)
+        return input1
+
+    df = transform_me.compute(transforms_context)
+    assert isinstance(df, pd.DataFrame)
+    assert_frame_equal(df, expected_df)
+
 
 def test_transform_pandas_two_inputs(spark_df_return_data_one, spark_df_return_data_two, transforms_context):
     from transforms.api import Input, Output, TransformContext, transform_pandas
@@ -502,6 +517,20 @@ def transform_me(output1, input1):
         assert "will have no effect" in record_message
 
 
+def test_lightweight_transforms_with_context(transforms_context):
+    from transforms.api import Input, Output, TransformContext, lightweight, transform
+
+    @lightweight()
+    @transform(
+        output1=Output("/output/to/dataset"),
+        input1=Input("/input1"),
+    )
+    def transform_me(ctx, output1, input1):
+        assert isinstance(ctx, TransformContext)  # ctx is our TransformContext
+
+    transform_me.compute(transforms_context)
+
+
 def test_transforms_with_incremental():
     from transforms.api import Input, Output, incremental, transform
 
@@ -587,3 +616,17 @@ def transform_me(input1):
         return input1
 
     transform_me.compute(transforms_context)
+
+
+def test_transform_context_auth_header(transforms_context):
+    from transforms.api import Input, Output, transform_df
+
+    @transform_df(
+        Output("output1"),
+        input1=Input("/input1"),
+    )
+    def transform_me(ctx, input1):
+        assert ctx.auth_header == "Bearer test_transform_context_auth_header_token"
+        return input1
+
+    transform_me.compute(transforms_context)