We recently also ran into this limitation. I put my +1 to Palantir supporting this in the LightweightInput.
Here is code that you can use to achieve it in the meantime:
from transforms.api import LightweightContext
# import polars as pl
from transforms.api import Input, Output, transform, LightweightInput, LightweightOutput
from foundry.transforms import (
LightweightFoundryConnector,
)
import requests
from requests.adapters import HTTPAdapter
from urllib3.util.retry import Retry
from urllib.parse import quote_plus
def _get_session(retries=3, backoff_factor=0.3):
session = requests.Session()
retry = Retry(
total=retries,
backoff_factor=backoff_factor,
)
adapter = HTTPAdapter(max_retries=retry)
session.mount("http://", adapter)
session.mount("https://", adapter)
return session
SESSION = _get_session()
def _request(
method: str,
auth_header: str,
url: str,
params: dict,
expected_status_codes: list[int] = [200, 204],
json: dict = None,
) -> requests.Response:
response = SESSION.request(
method=method,
url=url,
headers={"Authorization": auth_header},
json=json,
params=params,
timeout=30,
)
response.raise_for_status()
if response.status_code not in expected_status_codes:
raise ValueError(
f"Url: {url}\n"
+ f"Fields: {params}\n"
+ f"Body: {json}\n"
+ f"Response Code: {response.status_code}\n"
+ f"Expected Status Code: {expected_status_codes}\n"
+ f"Response Body: {response.text}"
)
return response
def _get_dataset_column_description(
auth_header: str,
foundry_url: str,
dataset_rid: str,
transaction_rid: str,
branch: str,
) -> dict:
response = _request(
method="GET",
auth_header=auth_header,
url=f"{foundry_url}/metadata/datasets/"
f"{dataset_rid}/branches/{quote_plus(branch)}/view/namespace/foundryColumnMetadata",
params={"endTransactionRid": transaction_rid},
expected_status_codes=[200, 204],
)
if response.status_code == 204:
return {}
maybe_descriptions = response.json()
return {
column_name: description_pack["description"]
for column_name, description_pack in maybe_descriptions.items()
if "description" in description_pack
}
@transform.using(
output_dataset1=Output(
"/.../column-test"
),
output_dataset2=Output(
"/.../column-test2"
),
lineitem=Input("ri.foundry.main.dataset.d137c770-137c-441c-9029-420db0078c33"),
no_desc=Input("ri.foundry.main.dataset.b73a4565-f7db-4ea1-bbc8-2ff4444de878"),
)
def compute(
ctx: LightweightContext,
lineitem: LightweightInput,
no_desc: LightweightInput,
output_dataset1: LightweightOutput,
output_dataset2: LightweightOutput,
) -> None:
service_url = LightweightFoundryConnector()._service_provider.services_config[
"foundry-metadata-api"
][0]
cd = _get_dataset_column_description(
ctx.auth_header,
foundry_url=service_url,
dataset_rid=lineitem.rid,
transaction_rid=lineitem.end_transaction_rid,
branch=lineitem.branch,
)
print(cd)
output_dataset1.write_table(
lineitem.polars(lazy=True).limit(1), column_descriptions=cd
)
cd2 = _get_dataset_column_description(
ctx.auth_header,
foundry_url=service_url,
dataset_rid=no_desc.rid,
transaction_rid=no_desc.end_transaction_rid,
branch=no_desc.branch,
)
print(cd2)
output_dataset2.write_table(
no_desc.polars(lazy=True).limit(1), column_descriptions=cd2
)