feat: Add MySQL/MariaDB Support as a Wren UI Persistence Backend

docker/.env.example

@@ -51,3 +51,10 @@ EXPERIMENTAL_ENGINE_RUST_VERSION=false
 # Wren Engine
 # OPTIONAL: set if you want to use local storage for the Wren Engine
 LOCAL_STORAGE=.
+
+# MySQL Database Configuration
+MYSQL_HOST=mysql-database
+MYSQL_PORT=3306
+MYSQL_USER=wrenai_user
+MYSQL_PASSWORD=wrenai_pwd
+MYSQL_DB=wrenai_db

wren-ui/README.md

@@ -28,6 +28,31 @@ export PG_URL=postgres://user:password@localhost:5432/dbname
 ```
 -  `PG_URL` is the connection string of your postgres database.
 
+To use MySQL or MariaDB as the database, set `DB_TYPE=mysql` and provide the connection settings via environment variables below. (The `mysql2` driver used by Wren UI supports both MySQL and MariaDB.)
+
+```bash
+# windows
+SET DB_TYPE=mysql
+SET MYSQL_HOST=localhost
+SET MYSQL_PORT=3306
+SET MYSQL_USER=root
+SET MYSQL_PASSWORD=your_password
+SET MYSQL_DB=your_database
+
+# linux or mac
+export DB_TYPE=mysql
+export MYSQL_HOST=localhost
+export MYSQL_PORT=3306
+export MYSQL_USER=root
+export MYSQL_PASSWORD=your_password
+export MYSQL_DB=your_database
+```
+
+Notes:
+- Ensure the user has sufficient privileges to create tables and run migrations.
+- For MariaDB, the same variables apply; just point them to your MariaDB instance.
+- If you use a non-default port, adjust `MYSQL_PORT` accordingly.
+
 To switch back to using SQLite, you can reassign the `DB_TYPE` to `sqlite`.
 ```
 # windows

wren-ui/knexfile.js

@@ -9,6 +9,19 @@ if (process.env.DB_TYPE === 'pg') {
     client: 'pg',
     connection: process.env.PG_URL,
   };
+} else if (process.env.DB_TYPE === 'mysql') {
+    console.log('Using MySQL at '+process.env.MYSQL_HOST+':'+(process.env.MYSQL_PORT || 3306));
+    module.exports = {
+      client: 'mysql2',
+      connection: {
+        host: process.env.MYSQL_HOST,
+        port: +(process.env.MYSQL_PORT || 3306),
+        user: process.env.MYSQL_USER,
+        password: process.env.MYSQL_PASSWORD,
+        database: process.env.MYSQL_DB || 'wren_ui',
+      },
+      pool: { min: 2, max: 10 },
+    };
 } else {
   console.log('Using SQLite');
   module.exports = {

wren-ui/migrations/20240125070643_create_project_table.js

@@ -24,23 +24,48 @@ exports.up = function (knex) {
     table.string('dataset_id').nullable().comment('big query datasetId');
 
     // duckdb
-    table
-      .jsonb('init_sql')
-      .nullable()
-      .comment('init sql for establishing duckdb environment');
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('init_sql')
+        .nullable()
+        .comment('init sql for establishing duckdb environment');
+    } else {
+      table
+        .jsonb('init_sql')
+        .nullable()
+        .comment('init sql for establishing duckdb environment');
+    }
     // knex jsonb ref: https://knexjs.org/guide/schema-builder.html#json
-    table
-      .jsonb('extensions')
-      .nullable()
-      .comment(
-        'duckdb extensions, will be a array-like string like, eg: ["extension1", "extension2"]',
-      );
-    table
-      .jsonb('configurations')
-      .nullable()
-      .comment(
-        'duckdb configurations that can be set in session, eg: { "key1": "value1", "key2": "value2" }',
-      );
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('extensions')
+        .nullable()
+        .comment(
+          'duckdb extensions, will be a array-like string like, eg: ["extension1", "extension2"]',
+        );
+    } else {
+      table
+        .jsonb('extensions')
+        .nullable()
+        .comment(
+          'duckdb extensions, will be a array-like string like, eg: ["extension1", "extension2"]',
+        );
+    }
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('configurations')
+        .nullable()
+        .comment(
+          'duckdb configurations that can be set in session, eg: { "key1": "value1", "key2": "value2" }',
+        );
+    } else {
+      table
+        .jsonb('configurations')
+        .nullable()
+        .comment(
+          'duckdb configurations that can be set in session, eg: { "key1": "value1", "key2": "value2" }',
+        );
+    }
 
     // not sure to store or not, the catalog & schema in the manifest
     table.string('catalog').comment('catalog name');

wren-ui/migrations/20240319083758_create_deploy_table.js

@@ -8,7 +8,11 @@ exports.up = function (knex) {
     table.integer('project_id').comment('Reference to project.id');
 
     // basic info
-    table.jsonb('manifest').comment('the deployed manifest');
+    if (knex.client.config.client === 'mysql2') {
+      table.json('manifest').comment('the deployed manifest');
+    } else {
+      table.jsonb('manifest').comment('the deployed manifest');
+    }
     table.string('hash').comment('the hash of the manifest');
 
     // status

wren-ui/migrations/20240327030000_create_ask_table.js

@@ -15,7 +15,11 @@ exports.up = function (knex) {
     })
     .createTable('thread_response', (table) => {
       table.increments('id').comment('ID');
-      table.integer('thread_id').comment('Reference to thread.id');
+      if (knex.client.config.client === 'mysql2') {
+        table.integer('thread_id').unsigned().comment('Reference to thread.id');
+      } else {
+        table.integer('thread_id').comment('Reference to thread.id');
+      }
       table.foreign('thread_id').references('thread.id').onDelete('CASCADE');
 
       // query id from AI service
@@ -24,8 +28,13 @@ exports.up = function (knex) {
       // response from AI service
       table.text('question').comment('the question of the response');
       table.string('status').comment('the status of the response');
-      table.jsonb('detail').nullable().comment('the detail of the response');
-      table.jsonb('error').nullable().comment('the error message if any');
+      if (knex.client.config.client === 'mysql2') {
+        table.json('detail').nullable().comment('the detail of the response');
+        table.json('error').nullable().comment('the error message if any');
+      } else {
+        table.jsonb('detail').nullable().comment('the detail of the response');
+        table.jsonb('error').nullable().comment('the error message if any');
+      }
 
       // timestamps
       table.timestamps(true, true);

wren-ui/migrations/20240419090558_add_foreign_key_to_model_column_and_metric_measure.js

@@ -5,9 +5,15 @@
 exports.up = function (knex) {
   return knex.schema
     .alterTable('model_column', (table) => {
+      if (knex.client.config.client === 'mysql2') {
+        table.integer('model_id').unsigned().alter();
+      }
       table.foreign('model_id').references('model.id').onDelete('CASCADE');
     })
     .alterTable('metric_measure', (table) => {
+      if (knex.client.config.client === 'mysql2') {
+        table.integer('metric_id').unsigned().alter();
+      }
       table.foreign('metric_id').references('metric.id').onDelete('CASCADE');
     });
 };

wren-ui/migrations/20240446090560_update_relationship_table.js

@@ -5,12 +5,18 @@
 exports.up = function (knex) {
   return knex.schema
     .alterTable('relation', (table) => {
+      if (knex.client.config.client === 'mysql2') {
+        table.integer('from_column_id').unsigned().alter();
+      }
       table
         .foreign('from_column_id')
         .references('model_column.id')
         .onDelete('CASCADE');
     })
     .alterTable('relation', (table) => {
+      if (knex.client.config.client === 'mysql2') {
+        table.integer('to_column_id').unsigned().alter();
+      }
       table
         .foreign('to_column_id')
         .references('model_column.id')

wren-ui/migrations/20240530062133_update_project_table.js

@@ -6,10 +6,17 @@
 // create connectionInfo column in project table
 exports.up = function (knex) {
   return knex.schema.table('project', (table) => {
-    table
-      .jsonb('connection_info')
-      .nullable()
-      .comment('Connection information for the project');
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('connection_info')
+        .nullable()
+        .comment('Connection information for the project');
+    } else {
+      table
+        .jsonb('connection_info')
+        .nullable()
+        .comment('Connection information for the project');
+    }
   });
 };
 

wren-ui/migrations/20240530105955_drop_project_table_columns.js

@@ -23,12 +23,21 @@ exports.up = function (knex) {
  */
 exports.down = function (knex) {
   return knex.schema.table('project', (table) => {
-    table
-      .jsonb('configurations')
-      .nullable()
-      .comment(
-        'duckdb configurations that can be set in session, eg: { "key1": "value1", "key2": "value2" }',
-      );
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('configurations')
+        .nullable()
+        .comment(
+          'duckdb configurations that can be set in session, eg: { "key1": "value1", "key2": "value2" }',
+        );
+    } else {
+      table
+        .jsonb('configurations')
+        .nullable()
+        .comment(
+          'duckdb configurations that can be set in session, eg: { "key1": "value1", "key2": "value2" }',
+        );
+    }
     table
       .text('credentials')
       .nullable()
@@ -39,12 +48,21 @@ exports.down = function (knex) {
       .comment('gcp project id, big query specific');
     table.string('dataset_id').nullable().comment('big query datasetId');
     table.text('init_sql');
-    table
-      .jsonb('extensions')
-      .nullable()
-      .comment(
-        'duckdb extensions, will be a array-like string like, eg: ["extension1", "extension2"]',
-      );
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('extensions')
+        .nullable()
+        .comment(
+          'duckdb extensions, will be a array-like string like, eg: ["extension1", "extension2"]',
+        );
+    } else {
+      table
+        .jsonb('extensions')
+        .nullable()
+        .comment(
+          'duckdb extensions, will be a array-like string like, eg: ["extension1", "extension2"]',
+        );
+    }
     table
       .string('host')
       .nullable()

wren-ui/migrations/20240610070534_create_schema_change_table.js

@@ -8,8 +8,13 @@ exports.up = function (knex) {
     table.integer('project_id').comment('Reference to project.id');
 
     // schema change info
-    table.jsonb('change').nullable();
-    table.jsonb('resolve').nullable();
+    if (knex.client.config.client === 'mysql2') {
+      table.json('change').nullable();
+      table.json('resolve').nullable();
+    } else {
+      table.jsonb('change').nullable();
+      table.jsonb('resolve').nullable();
+    }
 
     // timestamps
     table.timestamps(true, true);

wren-ui/migrations/20240928165009_create_model_nested_column.js

@@ -6,7 +6,11 @@ exports.up = function (knex) {
   return knex.schema.createTable('model_nested_column', (table) => {
     table.increments('id').comment('ID');
     table.integer('model_id').comment('Reference to model ID');
-    table.integer('column_id').comment('Reference to column ID');
+    if (knex.client.config.client === 'mysql2') {
+      table.integer('column_id').unsigned().comment('Reference to column ID');
+    } else {
+      table.integer('column_id').comment('Reference to column ID');
+    }
     table
       .string('column_path')
       .comment(

wren-ui/migrations/20241106232204_update_project_table.js

@@ -4,10 +4,17 @@
  */
 exports.up = function (knex) {
   return knex.schema.alterTable('project', (table) => {
-    table
-      .jsonb('questions')
-      .nullable()
-      .comment('The recommended questions generated by AI');
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('questions')
+        .nullable()
+        .comment('The recommended questions generated by AI');
+    } else {
+      table
+        .jsonb('questions')
+        .nullable()
+        .comment('The recommended questions generated by AI');
+    }
     table
       .string('query_id')
       .nullable()
@@ -16,10 +23,17 @@ exports.up = function (knex) {
       .string('questions_status')
       .nullable()
       .comment('The status of the recommended question pipeline');
-    table
-      .jsonb('questions_error')
-      .nullable()
-      .comment('The error of the recommended question pipeline');
+    if (knex.client.config.client === 'mysql2') {
+      table
+        .json('questions_error')
+        .nullable()
+        .comment('The error of the recommended question pipeline');
+    } else {
+      table
+        .jsonb('questions_error')
+        .nullable()
+        .comment('The error of the recommended question pipeline');
+    }
   });
 };