Medallion Architecture

import bauplan


def import_data_in_bronze_layer(table_name, import_branch, source_s3):
    client = bauplan.Client()
    try:
        # 1) Create + import table
        client.create_table(
            table=table_name, 
            search_uri=source_s3, 
            branch=import_branch, 
            namespace="bronze"
        )
        client.import_data(
            table=table_name, 
            search_uri=source_s3, 
            branch=import_branch, 
            namespace="bronze"
        )
    except bauplan.exceptions.BauplanError as e:
        raise Exception(f"🔴 Bronze import failed: {e}")
    
    # 2) Merge into main
    if not client.merge_branch(source_ref=import_branch, into_branch="main"):
        raise Exception("🔴 Merge failed.")
    print(f"✅ Bronze '{import_branch}' merged into main.")

import bauplan
from bauplan.standard_expectations import expect_column_no_nulls

@bauplan.model()
@bauplan.python('3.11', pip={'polars': '1.33.1'})
def silver_table(data=bauplan.Model('bronze_table')):
    import polars as pl
    
    # Convert input data into Polars and clean it
    df = data.to_polars()
    time_filter_utc = pl.Series(
        [pl.datetime(2022, 1, 1, 0, 0, 0, "UTC")]
    )[0]
    df = df.filter(pl.col("timestamp") >= time_filter_utc)
   
    return df.to_arrow()


@bauplan.expectation()
@bauplan.python("3.11")
def test_nulls(data=bauplan.Model('bronze_table')):
    col_to_check = "timestamp"    
    # Run expectation: Check for nulls and fail if nulls are found
    is_passed = expect_column_no_nulls(data, col_to_check)
    assert is_passed, f"❌ nulls found in {col_to_check}"
    
    return is_passed

import bauplan


def promote_to_gold(
    client: bauplan.Client,
    pipeline_dir: list,
    branch: str,
):
    # 1) Run the pipeline on a separate branch
    run_state = client.run(
        project_dir=pipeline_dir, 
        ref=branch, 
        namespace="gold"
    )
    # 2) Check for failures
    if run_state.job_status.lower() == "failed":
        raise Exception(f"{run_state.job_id} failed: {run_state.job_status}")
        
    # 3) Merge the silver branch into main (publish to gold)
    assert client.merge_branch(source_ref=branch, into_branch="main"), (
        "Something went wrong while merging into main."
    )