marcboeker / go-duckdb

go-duckdb provides a database/sql driver for the DuckDB database engine.
MIT License
646 stars 97 forks source link

Scan enums as strings #130

Closed k-anshul closed 9 months ago

k-anshul commented 9 months ago

This PR adds ability to scan duckDB enum types as strings by default:

  1. This is similar to how julia pkg for duckDB handles enums : https://github.com/duckdb/duckdb/blob/958f1898805822793fefef2b9a0c0e953c1f9fd0/tools/juliapkg/src/result.jl#L126
  2. This is also similar to how other go SQL drivers like pgx handle enums out of box :
    
    package main

import ( "context" "database/sql" "fmt" "log"

_ "github.com/jackc/pgx/v5/stdlib"

)

// Define a custom type for the enum type Status string

const ( Pending Status = "pending" Approved Status = "approved" Rejected Status = "rejected" )

func main() {

// Connect to the PostgreSQL database using pgxpool
connString := "postgresql://postgres:postgres@localhost:5432/postgres"
conn, err := sql.Open("pgx", connString)
if err != nil {
    log.Fatal(err)
}
defer conn.Close()

// Create the table with an enum column
_, err = conn.ExecContext(context.Background(), `
    CREATE TYPE status_enum AS ENUM ('pending', 'approved', 'rejected');
    CREATE TABLE your_table (
        id SERIAL PRIMARY KEY,
        name VARCHAR(50),
        status status_enum
    );
`)
if err != nil {
    log.Fatal(err)
}

// Insert some data into the table
_, err = conn.ExecContext(context.Background(), "INSERT INTO your_table (name, status) VALUES ($1, $2)", "Item 1", Pending)
if err != nil {
    log.Fatal(err)
}

// Query the data and scan the results
rows, err := conn.QueryContext(context.Background(), "SELECT id, name, status FROM your_table")
if err != nil {
    log.Fatal(err)
}
defer rows.Close()

for rows.Next() {
    var id int
    var name string
    var status Status

    err := rows.Scan(&id, &name, &status)
    if err != nil {
        log.Fatal(err)
    }

    fmt.Printf("ID: %d, Name: %s, Status: %s\n", id, name, status)
}

if err := rows.Err(); err != nil {
    log.Fatal(err)
}

}

marcboeker commented 9 months ago

Great addition, thanks @k-anshul. I've refactored the tests a little bit. Could you please provide a reference or documentation link for further understanding of the following part?

internalType := C.duckdb_enum_internal_type(ty)
    switch internalType {
    case C.DUCKDB_TYPE_UTINYINT:
        idx = uint64(get[uint8](vector, rowIdx))
    case C.DUCKDB_TYPE_USMALLINT:
        idx = uint64(get[uint16](vector, rowIdx))
    case C.DUCKDB_TYPE_UINTEGER:
        idx = uint64(get[uint32](vector, rowIdx))
    case C.DUCKDB_TYPE_UBIGINT:
        idx = get[uint64](vector, rowIdx)
    default:
        return "", errInvalidType
    }
k-anshul commented 9 months ago

Actually I couldn't find any documentation. I looked at the test case around enums in the CAPI src code here and then modified it as per how we scan decimal in this driver.