Palantir Foundry to Databricks Migration

 # pip install foundry-sdk
from foundry_sdk.v1 import FoundryClient
import foundry_sdk, pyarrow.parquet as pq, boto3, io

client = FoundryClient(
    auth=foundry_sdk.UserTokenAuth(token=''),
    hostname='yourorg.palantirfoundry.com'
)

# Stream dataset as Arrow (most efficient — no CSV overhead)
dataset_rid = 'ri.foundry.main.dataset.c26f11c8-...'
arrow_stream = client.datasets.Dataset.read(
    dataset_rid,
    format='ARROW',          # Arrow is fastest; also supports CSV
    branch_id='master'
)

# Write directly to S3 as Parquet (ready for Delta conversion)
buf = io.BytesIO(arrow_stream)
table = pq.read_table(buf)
pq.write_to_dataset(table, root_path='s3://my-bucket/bronze/my_dataset/')

# For large datasets: paginate using start_transaction_rid / end_transaction_rid
# For schema only: client.datasets.Dataset.get_schema(dataset_rid)

 -- Run this in a Databricks SQL notebook or as a Workflow task

-- Option A: CONVERT TO DELTA (fastest, in-place, no data copy)
CONVERT TO DELTA parquet.`s3://my-bucket/bronze/my_dataset/`;

-- Option B: Write as Delta directly (when you want to control schema)
-- (Python notebook)
df = spark.read.parquet('s3://my-bucket/bronze/my_dataset/')
df.write.format('delta') \
        .mode('overwrite') \
        .option('overwriteSchema', 'true') \
        .saveAsTable('main.bronze.my_dataset')

-- Register in Unity Catalog after CONVERT TO DELTA
CREATE TABLE IF NOT EXISTS main.bronze.my_dataset
USING DELTA LOCATION 's3://my-bucket/bronze/my_dataset/';

 # pip install foundry-dev-tools
from foundry_dev_tools import FoundryContext

ctx = FoundryContext()  # reads credentials from ~/.foundry/config.toml

# List all datasets in a project (for inventory)
datasets = ctx.foundry_sql_server.query_foundry_sql(
    'SELECT * FROM `/path/to/project` WHERE type = \'DATASET\''
)

# Pull a specific dataset as a Spark DataFrame
df = ctx.foundry_sql_server.query_foundry_sql(
    'SELECT * FROM `/ri/foundry.main.dataset/your-dataset-rid`'
)

# Export transform source file from a Code Repository branch
# (Use Foundry REST API Code Repositories endpoint for bulk export)
import requests
headers = {'Authorization': f'Bearer {token}'}
resp = requests.get(
    f'https://yourorg.palantirfoundry.com/code-repositories/api/repositories/{repo_rid}/contents?branch=master',
    headers=headers
)
# Returns list of files — iterate and download each .py transform file

 # pip install sqlglot
import sqlglot
import os, pathlib

# Single query example
foundry_sql = """
    SELECT date_format(event_ts, 'yyyy-MM') AS month,
           approx_count_distinct(user_id)   AS unique_users
    FROM   `/ri/foundry.main.dataset/events-rid`
    WHERE  event_type = 'CLICK'
    GROUP  BY 1
"""

databricks_sql = sqlglot.transpile(
    foundry_sql,
    read='spark',        # Foundry uses Spark SQL dialect
    write='databricks',  # Target: Databricks SQL
    pretty=True
)[0]
print(databricks_sql)

# Bulk conversion — process a directory of .sql files
for sql_file in pathlib.Path('./foundry_sql/').glob('**/*.sql'):
    src = sql_file.read_text()
    out = sqlglot.transpile(src, read='spark', write='databricks', pretty=True)[0]
    dest = pathlib.Path('./databricks_sql/') / sql_file.name
    dest.write_text(out)
    print(f'Converted {sql_file.name}')

 # This is what a Foundry Python transform looks like
from transforms.api import transform_df, Input, Output

@transform_df(
    Output('/ri/foundry.main.dataset/output-rid'),
    raw=Input('/ri/foundry.main.dataset/input-rid'),
)
def compute(raw):
    return raw.filter(raw.status == 'ACTIVE') \
               .groupBy('region') \
               .agg({'revenue': 'sum'})

 # After LLM rewrite — clean Databricks PySpark notebook cell
# Input table: main.bronze.input_table
# Output table: main.silver.output_table

from pyspark.sql import functions as F

raw = spark.table('main.bronze.input_table')

result = (
    raw.filter(F.col('status') == 'ACTIVE')
       .groupBy('region')
       .agg(F.sum('revenue').alias('revenue'))
)

result.write.format('delta').mode('overwrite').saveAsTable('main.silver.output_table')

# pip install anthropic  (or: openai for GPT-4o, or use Databricks DBRX endpoint)
import anthropic, pathlib, json, concurrent.futures
client = anthropic.Anthropic()  # uses ANTHROPIC_API_KEY env var

SYSTEM_PROMPT = '''You are an expert Spark/Databricks data engineer.
Convert the Palantir Foundry Python transform below to a clean Databricks PySpark notebook.
Rules:
1. Remove all imports from transforms.api
2. Remove @transform_df and all decorator arguments
3. Replace Input('/path/to/dataset') with spark.table('') — infer table name from path
4. Replace Output write calls with df.write.format('delta').mode('overwrite').saveAsTable('')
5. Keep ALL business logic (filters, joins, aggregations) identical
6. Return ONLY the Python code, no markdown, no commentary'''

def convert_transform(py_file: pathlib.Path) -> str:
    source_code = py_file.read_text()
    msg = client.messages.create(
        model='claude-sonnet-4-6',
        max_tokens=4096,
        system=SYSTEM_PROMPT,
        messages=[{'role': 'user', 'content': source_code}]
    )
    return msg.content[0].text

# Batch convert all .py files in the extracted transform directory
transform_files = list(pathlib.Path('./foundry_transforms/').glob('**/*.py'))
print(f'Converting {len(transform_files)} transform files...')

with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
    futures = {executor.submit(convert_transform, f): f for f in transform_files}
    for future in concurrent.futures.as_completed(futures):
        src_file = futures[future]
        converted = future.result()
        out_path = pathlib.Path('./databricks_notebooks/') / src_file.name
        out_path.write_text(converted)
        print(f'  ✓ {src_file.name}')

# Alternative: use Databricks DBRX endpoint (keeps data in your environment)
# Replace client with: OpenAI(api_key=TOKEN, base_url=DATABRICKS_ENDPOINT)
# Model: 'databricks-meta-llama-3-70b-instruct' or 'databricks-dbrx-instruct'

# Step 1: Export lineage graph from Foundry REST API
# GET /foundry-catalog/api/datasets/{rid}/lineage
# → returns JSON with upstream_datasets[] and downstream_datasets[]

# Step 2: Build Databricks Workflow JSON from lineage
import json

def build_workflow_from_lineage(lineage_json: dict) -> dict:
    tasks = []
    for transform in lineage_json['transforms']:
        task = {
            'task_key': transform['name'].replace('/', '_'),
            'notebook_task': {
                'notebook_path': f'/Repos/migration/{transform["name"]}',
                'base_parameters': {}
            },
            'existing_cluster_id': 'YOUR_CLUSTER_ID',
            'depends_on': [
                {'task_key': dep.replace('/', '_')}
                for dep in transform.get('upstream_transforms', [])
            ]
        }
        tasks.append(task)

    return {
        'name': 'migrated_pipeline',
        'tasks': tasks,
        'schedule': {
            'quartz_cron_expression': transform.get('cron', '0 0 6 * * ?'),
            'timezone_id': 'UTC'
        }
    }

# Step 3: POST to Databricks Workflows API to create the job
import requests
workflow_json = build_workflow_from_lineage(lineage_data)
resp = requests.post(
    'https://.azuredatabricks.net/api/2.1/jobs/create',
    headers={'Authorization': f'Bearer {DATABRICKS_TOKEN}'},
    json=workflow_json
)
print('Created job id:', resp.json()['job_id'])

import mlflow
import mlflow.pyfunc

# Register an exported AIP model artifact into MLflow
mlflow.set_tracking_uri('databricks')  # Uses Databricks MLflow
mlflow.set_experiment('/migrations/aip_model_registry')

with mlflow.start_run():
    # Log model from local path (exported from AIP)
    mlflow.pyfunc.log_model(
        artifact_path='model',
        python_model=mlflow.pyfunc.load_model('./exported_aip_model/'),
        registered_model_name='my_model_name',  # same name as in AIP
        pip_requirements='./exported_aip_model/requirements.txt'
    )
    mlflow.log_param('migrated_from', 'Palantir AIP')
print('Model registered in MLflow Model Registry')

# Great Expectations for automated data validation
# pip install great_expectations
import great_expectations as gx

context = gx.get_context()
batch = context.sources.add_spark('databricks') \
               .add_dataframe_asset('migrated_table') \
               .build_batch_request(dataframe=spark.table('main.silver.my_table'))

suite = context.add_expectation_suite('migration_suite')
validator = context.get_validator(batch_request=batch, expectation_suite=suite)

# Assert row count matches Foundry source (store expected count from export step)
validator.expect_table_row_count_to_equal(value=1_234_567)

# Assert no nulls on critical columns
validator.expect_column_values_to_not_be_null('customer_id')

# Assert value ranges are preserved
validator.expect_column_values_to_be_between('revenue', min_value=0, max_value=1e9)

results = validator.validate()
print('Validation passed:', results.success)