always use returning statement if supported (#170)

src/components/gui/main-connection.tsx

@@ -38,7 +38,7 @@ function MainConnectionContainer() {
    */
   useLayoutEffect(() => {
     console.info("Injecting message into window object");
-    window.internalPubSub = new InternalPubSub();
+    if (!window.internalPubSub) window.internalPubSub = new InternalPubSub();
   }, [driver]);
 
   useEffect(() => {

src/components/gui/query-result.tsx

@@ -20,7 +20,10 @@ export default function QueryResult({
       return { _tag: "EXPLAIN", value: result.result } as const;
     }
 
-    const state = OptimizeTableState.createFromResult(result.result);
+    const state = OptimizeTableState.createFromResult(
+      databaseDriver,
+      result.result
+    );
     state.setReadOnlyMode(true);
     state.mismatchDetection = databaseDriver.getFlags().mismatchDetection;
     return { _tag: "QUERY", value: state } as const;

src/components/gui/table-optimized/OptimizeTableState.tsx

@@ -2,6 +2,7 @@ import { selectArrayFromIndexList } from "@/components/lib/export-helper";
 import { OptimizeTableHeaderProps } from ".";
 import { LucideKey, LucideKeySquare, LucideSigma } from "lucide-react";
 import {
+  BaseDriver,
   DatabaseResultSet,
   DatabaseTableSchema,
   TableColumnDataType,
@@ -37,14 +38,19 @@ export default class OptimizeTableState {
   protected changeLogs: Record<number, OptimizeTableRowValue> = {};
 
   static createFromResult(
+    driver: BaseDriver,
     dataResult: DatabaseResultSet,
     schemaResult?: DatabaseTableSchema
   ) {
     return new OptimizeTableState(
       dataResult.headers.map((header) => {
+        const headerData = schemaResult
+          ? schemaResult.columns.find((c) => c.name === header.name)
+          : undefined;
+
         let initialSize = 150;
         const headerName = header.name;
-        const dataType = header.type;
+        const dataType = header.type ?? driver.inferTypeFromHeader(headerData);
 
         if (
           dataType === TableColumnDataType.INTEGER ||
@@ -67,10 +73,6 @@ export default class OptimizeTableState {
           initialSize = Math.max(150, Math.min(500, maxSize * 8));
         }
 
-        const headerData = schemaResult
-          ? schemaResult.columns.find((c) => c.name === header.name)
-          : undefined;
-
         // --------------------------------------
         // Matching foreign key
         // --------------------------------------

src/components/gui/tabs/table-data-tab.tsx

@@ -89,6 +89,7 @@ export default function TableDataWindow({
           });
 
         const tableState = OptimizeTableState.createFromResult(
+          databaseDriver,
           dataResult,
           schemaResult
         );

src/components/gui/toolbar.tsx

@@ -48,7 +48,7 @@ export function ToolbarButton({
   if (tooltip) {
     return (
       <Tooltip>
-        <TooltipTrigger>{buttonContent}</TooltipTrigger>
+        <TooltipTrigger asChild>{buttonContent}</TooltipTrigger>
         <TooltipContent>{tooltip}</TooltipContent>
       </Tooltip>
     );

src/drivers/base-driver.ts

@@ -207,6 +207,12 @@ export interface DriverFlags {
   supportModifyColumn: boolean;
   mismatchDetection: boolean;
   dialect: SupportedDialect;
+
+  // If database supports this, we don't need
+  // to make a separate request to get updated
+  // data when update
+  supportInsertReturning: boolean;
+  supportUpdateReturning: boolean;
 }
 
 export interface DatabaseTableColumnChange {
@@ -252,6 +258,10 @@ export abstract class BaseDriver {
     tableName: string
   ): Promise<DatabaseTableSchema>;
 
+  abstract inferTypeFromHeader(
+    header?: DatabaseTableColumn
+  ): TableColumnDataType | undefined;
+
   abstract trigger(
     schemaName: string,
     name: string

src/drivers/common-sql-imp.ts

@@ -36,11 +36,25 @@ export default abstract class CommonSQLImplement extends BaseDriver {
 
     const sqls = ops.map((op) => {
       if (op.operation === "INSERT")
-        return insertInto(this, schemaName, tableName, op.values);
+        return insertInto(
+          this,
+          schemaName,
+          tableName,
+          op.values,
+          this.getFlags().supportInsertReturning
+        );
+
       if (op.operation === "DELETE")
         return deleteFrom(this, schemaName, tableName, op.where);
 
-      return updateTable(this, schemaName, tableName, op.values, op.where);
+      return updateTable(
+        this,
+        schemaName,
+        tableName,
+        op.values,
+        op.where,
+        this.getFlags().supportInsertReturning
+      );
     });
 
     const result = await this.transaction(sqls);
@@ -57,18 +71,29 @@ export default abstract class CommonSQLImplement extends BaseDriver {
       }
 
       if (op.operation === "UPDATE") {
-        const selectResult = await this.findFirst(
-          schemaName,
-          tableName,
-          op.where
-        );
+        if (r.rows.length === 1)
+          // This is when database support RETURNING
+          tmp.push({
+            record: r.rows[0],
+          });
+        else {
+          const selectResult = await this.findFirst(
+            schemaName,
+            tableName,
+            op.where
+          );
 
-        tmp.push({
-          lastId: r.lastInsertRowid,
-          record: selectResult.rows[0],
-        });
+          tmp.push({
+            lastId: r.lastInsertRowid,
+            record: selectResult.rows[0],
+          });
+        }
       } else if (op.operation === "INSERT") {
-        if (op.autoIncrementPkColumn) {
+        if (r.rows.length === 1) {
+          tmp.push({
+            record: r.rows[0],
+          });
+        } else if (op.autoIncrementPkColumn) {
           const selectResult = await this.findFirst(schemaName, tableName, {
             [op.autoIncrementPkColumn]: r.lastInsertRowid,
           });

src/drivers/mysql/mysql-driver.ts

@@ -5,6 +5,7 @@ import {
   DriverFlags,
   DatabaseSchemaItem,
   DatabaseTableColumn,
+  TableColumnDataType,
 } from "../base-driver";
 import CommonSQLImplement from "../common-sql-imp";
 import { escapeSqlValue } from "../sqlite/sql-helper";
@@ -53,6 +54,9 @@ export default abstract class MySQLLikeDriver extends CommonSQLImplement {
       mismatchDetection: false,
       supportCreateUpdateTable: false,
       dialect: "mysql",
+
+      supportInsertReturning: false,
+      supportUpdateReturning: false,
     };
   }
 
@@ -153,4 +157,8 @@ export default abstract class MySQLLikeDriver extends CommonSQLImplement {
   createUpdateTableSchema(): string[] {
     throw new Error("Not implemented");
   }
+
+  inferTypeFromHeader(): TableColumnDataType | undefined {
+    return undefined;
+  }
 }

src/drivers/query-builder.ts

@@ -64,12 +64,14 @@ export function insertInto(
   dialect: BaseDriver,
   schema: string,
   table: string,
-  value: Record<string, unknown>
+  value: Record<string, unknown>,
+  supportReturning: boolean
 ) {
   return [
     "INSERT INTO",
     `${dialect.escapeId(schema)}.${dialect.escapeId(table)}`,
     generateInsertValue(dialect, value),
+    supportReturning ? "RETURNING *" : "",
   ].join(" ");
 }
 
@@ -78,14 +80,16 @@ export function updateTable(
   schema: string,
   table: string,
   value: Record<string, unknown>,
-  where: Record<string, unknown>
+  where: Record<string, unknown>,
+  supportReturning: boolean
 ): string {
   return [
     "UPDATE",
     `${dialect.escapeId(schema)}.${dialect.escapeId(table)}`,
     "SET",
     generateSet(dialect, value),
     generateWhere(dialect, where),
+    supportReturning ? "RETURNING *" : "",
   ]
     .filter(Boolean)
     .join(" ");

src/drivers/sqlite-base-driver.ts

@@ -9,8 +9,9 @@ import type {
   DatabaseValue,
   DriverFlags,
   SelectFromTableOptions,
+  TableColumnDataType,
 } from "./base-driver";
-import { escapeSqlValue } from "@/drivers/sqlite/sql-helper";
+import { convertSqliteType, escapeSqlValue } from "@/drivers/sqlite/sql-helper";
 
 import { parseCreateTableScript } from "@/drivers/sqlite/sql-parse-table";
 import { parseCreateTriggerScript } from "@/drivers/sqlite/sql-parse-trigger";
@@ -32,6 +33,8 @@ export abstract class SqliteLikeBaseDriver extends CommonSQLImplement {
     return {
       supportBigInt: false,
       supportModifyColumn: false,
+      supportInsertReturning: true,
+      supportUpdateReturning: true,
       defaultSchema: "main",
       optionalSchema: true,
       mismatchDetection: false,
@@ -170,6 +173,13 @@ export abstract class SqliteLikeBaseDriver extends CommonSQLImplement {
     // do nothing
   }
 
+  inferTypeFromHeader(
+    header?: DatabaseTableColumn
+  ): TableColumnDataType | undefined {
+    if (!header) return undefined;
+    return convertSqliteType(header.type);
+  }
+
   async tableSchema(
     schemaName: string,
     tableName: string

src/drivers/sqljs-driver.ts

@@ -84,12 +84,6 @@ export default class SqljsDriver extends SqliteLikeBaseDriver {
         rowsWritten: null,
         queryDurationMs: endTime - startTime,
       },
-      lastInsertRowid:
-        headers.length > 0
-          ? undefined
-          : (this.db.exec("select last_insert_rowid();")[0].values[0][0] as
-              | number
-              | undefined),
     };
   }