example.py

from nimbus import Nimbus, ConnInfo
import psycopg2

conn_info = ConnInfo(
    host="localhost",
    port=4566,
    user="root",
    password="",
    dbname="dev",
)

pg_source_config = {
    "hostname": "postgres",
    "port": 5432,
    "username": "postgres",
    "password": "postgres",
    "database.name": "postgres",
}

ICEBERG_TABLE_NAME = "my_iceberg_table"


def example_ingest():
    bus = Nimbus(conn_info)

    j = bus.json_ingestor()

    stream = bus.create_stream(
        name="direct_ingest",
        destination=bus.table(ICEBERG_TABLE_NAME),
        source=j,
    )
    try:
        stream.deploy()
        for i in range(10):
            j.ingest_record({"id": i, "name": f"name_{i}"})

        j.flush()
    finally:
        j.close()


def example_postgres_cdc():
    bus = Nimbus(conn_info)

    pg_source = bus.postgres_source(
        name="my_pg_source",
        config=pg_source_config,
    )

    stream = bus.create_stream(
        name="pg_cdc_ingest",
        destination=bus.table(ICEBERG_TABLE_NAME),
        source=bus.postgres_cdc(
            postgres_source=pg_source,
            postgres_table_name="public.my_pg_table",
            config={
                "snapshot": True,
                "backfill.parallelism": 4,
            },
        ),
    )

    stream.deploy()


def prepare():
    # create postgres table
    c = pg_source_config
    with psycopg2.connect(
        f"postgres://{c['username']}:{c['password']}@localhost:{c['port']}/{c['database.name']}"
    ) as pg:
        with pg.cursor() as cur:
            cur.execute(
                """
                CREATE TABLE IF NOT EXISTS my_pg_table (
                    id   INTEGER PRIMARY KEY,
                    name VARCHAR
                );
                """
            )
            pg.commit()

    # create iceberg table
    with conn_info.connect() as conn:
        conn.exec(
            f"""
            CREATE CONNECTION IF NOT EXISTS lakekeeper_catalog_conn WITH (
                type                 = 'iceberg', 
                warehouse.path       = 'risingwave-warehouse', 
                s3.region            = 'us-east-1', 
                s3.access.key        = 'minioadmin', 
                s3.secret.key        = 'minioadmin', 
                s3.endpoint          = 'http://minio:9301', 
                catalog.type         = 'rest',
                catalog.uri          = 'http://lakekeeper:8181/catalog/', 
                s3.path.style.access = true
            );

            SET iceberg_engine_connection = 'public.lakekeeper_catalog_conn';
            
            CREATE TABLE IF NOT EXISTS {ICEBERG_TABLE_NAME} (
                id   INT PRIMARY KEY, 
                name VARCHAR
            ) WITH (
                commit_checkpoint_interval = 1
            ) ENGINE = iceberg;
            """
        )


if __name__ == "__main__":
    prepare()

    example_ingest()

    example_postgres_cdc()