Skip to content
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension


Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
4 changes: 4 additions & 0 deletions docker/.env.example
Original file line number Diff line number Diff line change
Expand Up @@ -46,3 +46,7 @@ AI_SERVICE_FORWARD_PORT=5555

# Wren UI
EXPERIMENTAL_ENGINE_RUST_VERSION=false

# Wren Engine
# OPTIONAL: set if you want to use local storage for the Wren Engine
LOCAL_STORAGE=.
2 changes: 2 additions & 0 deletions docker/docker-compose.yaml
Original file line number Diff line number Diff line change
Expand Up @@ -41,6 +41,8 @@ services:
- ${IBIS_SERVER_PORT}
environment:
WREN_ENGINE_ENDPOINT: http://wren-engine:${WREN_ENGINE_PORT}
volumes:
- ${LOCAL_STORAGE:-.}:/usr/src/app/data
networks:
- wren

Expand Down
71 changes: 71 additions & 0 deletions wren-launcher/commands/dbt.go
Original file line number Diff line number Diff line change
@@ -0,0 +1,71 @@
package commands

import (
"flag"
"os"

"github.com/Canner/WrenAI/wren-launcher/commands/dbt"
"github.com/pterm/pterm"
)

// DbtAutoConvert automatically searches for dbt profiles and catalog.json,
// then converts them to WrenDataSource and Wren MDL format
func DbtAutoConvert() {
var opts struct {
ProjectPath string
OutputDir string
ProfileName string
Target string
}

// Define command line flags
flag.StringVar(&opts.ProjectPath, "path", "", "Path to the dbt project root directory")
flag.StringVar(&opts.OutputDir, "output", "", "Output directory for generated JSON files")
flag.StringVar(&opts.ProfileName, "profile", "", "Specific profile name to use (optional, uses first found if not provided)")
flag.StringVar(&opts.Target, "target", "", "Specific target to use (optional, uses profile default if not provided)")
flag.Parse()

// Validate required parameters
if opts.ProjectPath == "" {
pterm.Error.Println("Error: --path parameter is required")
pterm.Info.Println("Usage: wren-launcher dbt-auto-convert --path /path/to/dbt/project --output /path/to/output")
os.Exit(1)
}

if opts.OutputDir == "" {
pterm.Error.Println("Error: --output parameter is required")
pterm.Info.Println("Usage: wren-launcher dbt-auto-convert --path /path/to/dbt/project --output /path/to/output")
os.Exit(1)
}

// ConvertOptions struct for core conversion logic
convertOpts := dbt.ConvertOptions{
ProjectPath: opts.ProjectPath,
OutputDir: opts.OutputDir,
ProfileName: opts.ProfileName,
Target: opts.Target,
RequireCatalog: true, // DbtAutoConvert requires catalog.json to exist
}

// Call the core conversion logic
_, err := dbt.ConvertDbtProjectCore(convertOpts)
if err != nil {
pterm.Error.Printf("Error: Conversion failed: %v\n", err)
os.Exit(1)
}
}

// DbtConvertProject is a public wrapper function for processDbtProject to use
// It converts a dbt project without requiring catalog.json to exist
func DbtConvertProject(projectPath, outputDir, profileName, target string, usedByContainer bool) (*dbt.ConvertResult, error) {
convertOpts := dbt.ConvertOptions{
ProjectPath: projectPath,
OutputDir: outputDir,
ProfileName: profileName,
Target: target,
RequireCatalog: false, // Allow processDbtProject to continue without catalog.json
UsedByContainer: usedByContainer,
}

return dbt.ConvertDbtProjectCore(convertOpts)
}
116 changes: 116 additions & 0 deletions wren-launcher/commands/dbt/README.md
Original file line number Diff line number Diff line change
@@ -0,0 +1,116 @@
# How to Support a New Data Source

This document outlines the steps required to add support for a new data source to the dbt project converter.
The target data source must be supported by both dbt and the Wren engine:
- [dbt supported databases](https://docs.getdbt.com/docs/supported-data-platforms)
- [Wren engine supported data sources](https://docs.getwren.ai/oss/wren_engine_api#tag/AthenaConnectionInfo)

## 1. Implement the DataSource Interface

The first step is to define a new struct for your data source and implement the `DataSource` interface defined in `data_source.go`.

The `DataSource` interface is as follows:

```go
type DataSource interface {
GetType() string
Validate() error
MapType(sourceType string) string
}
```

### Steps:

1. **Define Your Struct**: Create a new struct that represents the connection properties for your data source. The fields in this struct should correspond to the properties defined in the [Wren engine's API documentation](https://docs.getwren.ai/oss/wren_engine_api#tag/SnowflakeConnectionInfo) for the target data source.

For example, to add support for `Snowflake`, you would define the following struct:

```go
type WrenSnowflakeDataSource struct {
Account string `json:"account"`
User string `json:"user"`
Password string `json:"password"`
Database string `json:"database"`
Warehouse string `json:"warehouse"`
// ... other properties
}
```

2. **Implement `GetType()`**: This method should return a string that identifies your data source type (e.g., `"snowflake"`).

3. **Implement `Validate()`**: This method should check if the essential properties of your data source are set and valid. Return an error if validation fails.

4. **Implement `MapType()`**: This method is crucial for mapping data types from the source system (as defined in `catalog.json`) to Wren's supported data types (e.g., `integer`, `varchar`, `timestamp`).

## 2. Add Conversion Logic in `data_source.go`

After implementing the interface, you need to integrate your new data source into the conversion logic. This is done by updating the `convertConnectionToDataSource` function in `data_source.go`.

Add a new `case` to the `switch` statement that matches the `type` field from the dbt `profiles.yml` file. This new case will be responsible for creating an instance of your new data source struct from the dbt connection details.

### Example:

```go
// in data_source.go

func convertConnectionToDataSource(conn DbtConnection, dbtHomePath, profileName, outputName string) (DataSource, error) {
switch strings.ToLower(conn.Type) {
case "postgres", "postgresql":
return convertToPostgresDataSource(conn)
case "duckdb":
return convertToLocalFileDataSource(conn, dbtHomePath)
// Add your new case here
case "snowflake":
return convertToSnowflakeDataSource(conn) // Implement this function
default:
// ...
}
}

// Implement the conversion function
func convertToSnowflakeDataSource(conn DbtConnection) (*WrenSnowflakeDataSource, error) {
// Logic to extract snowflake properties from conn
// and return a new *WrenSnowflakeDataSource
}
```

## 3. Handle the New Data Source in `ConvertDbtProjectCore`

The `ConvertDbtProjectCore` function in `converter.go` is responsible for generating the `wren-datasource.json` file. You must add your new data source to the `switch` statement within this function to ensure it is correctly serialized.

### Steps:

1. **Locate the `switch` statement**: Find the `switch typedDS := ds.(type)` block inside `ConvertDbtProjectCore`.
2. **Add a new `case`**: Add a new `case` for your data source struct. Inside this case, construct the `wrenDataSource` map with the correct `type` and `properties`.

### Example:

```go
// in converter.go's ConvertDbtProjectCore function

// ...
switch typedDS := ds.(type) {
case *WrenPostgresDataSource:
// ...
case *WrenLocalFileDataSource:
// ...
// Add your new case here
case *WrenSnowflakeDataSource:
wrenDataSource = map[string]interface{}{
"type": "snowflake",
"properties": map[string]interface{}{
"account": typedDS.Account,
"user": typedDS.User,
"password": typedDS.Password,
"database": typedDS.Database,
"warehouse": typedDS.Warehouse,
// ... other properties
},
}
default:
// ...
}
// ...
```

**Note on File-Based Data Sources**: If your data source is file-based (like `duckdb`), you also need to add logic to set the `localStoragePath` variable correctly within `ConvertDbtProjectCore`. This path tells the Wren engine where to find the data files.
Loading
Loading