From 5bed0f8f195f615844d5dbe322ebfe47b76ba2f5 Mon Sep 17 00:00:00 2001
From: ZY <loozhengyuan@gmail.com>
Date: Fri, 4 Aug 2023 08:58:51 +0800
Subject: [PATCH] feat: add conn-level `infileStreamFactory` option (#2159)

---
 documentation/en/Extras.md        | 22 ++++++++++++++++++++++
 lib/connection.js                 | 17 +++++++++++++----
 lib/connection_config.js          |  2 ++
 typings/mysql/lib/Connection.d.ts |  6 ++++++
 4 files changed, 43 insertions(+), 4 deletions(-)

diff --git a/documentation/en/Extras.md b/documentation/en/Extras.md
index a90f74c808..dbe9ede205 100644
--- a/documentation/en/Extras.md
+++ b/documentation/en/Extras.md
@@ -68,6 +68,28 @@ connection.query(
 );
 ```
 
+The `infileStreamFactory` option may also be set at a connection-level:
+
+```js
+const fs = require("fs");
+const mysql = require('mysql2');
+
+const connection = mysql.createConnection({
+  user: 'test',
+  database: 'test',
+  infileStreamFactory: path => {
+    // Validate file path
+    const validPaths = ['/tmp/data.csv'];
+    if (!validPaths.includes(path)) {
+      throw new Error(`invalid file path: ${path}: expected to be one of ${validPaths.join(',')}`);
+    }
+    return fs.createReadStream(path);
+  }
+});
+
+connection.query('LOAD DATA LOCAL INFILE "/tmp/data.csv" INTO TABLE test', onInserted);
+```
+
 ## Connecting using custom stream:
 
 ```js
diff --git a/lib/connection.js b/lib/connection.js
index d6ad7a4f9b..b9f01a4cf4 100644
--- a/lib/connection.js
+++ b/lib/connection.js
@@ -615,10 +615,15 @@ class Connection extends EventEmitter {
   }
 
   execute(sql, values, cb) {
-    let options = {};
+    let options = {
+      infileStreamFactory: this.config.infileStreamFactory
+    };
     if (typeof sql === 'object') {
       // execute(options, cb)
-      options = sql;
+      options = {
+        ...options,
+        ...sql
+      };
       if (typeof values === 'function') {
         cb = values;
       } else {
@@ -899,11 +904,15 @@ class Connection extends EventEmitter {
 
   static createQuery(sql, values, cb, config) {
     let options = {
-      rowsAsArray: config.rowsAsArray
+      rowsAsArray: config.rowsAsArray,
+      infileStreamFactory: config.infileStreamFactory
     };
     if (typeof sql === 'object') {
       // query(options, cb)
-      options = sql;
+      options = {
+        ...options,
+        ...sql
+      };
       if (typeof values === 'function') {
         cb = values;
       } else if (values !== undefined) {
diff --git a/lib/connection_config.js b/lib/connection_config.js
index 236cf6747f..ec52051625 100644
--- a/lib/connection_config.js
+++ b/lib/connection_config.js
@@ -30,6 +30,7 @@ const validOptions = {
   flags: 1,
   host: 1,
   insecureAuth: 1,
+  infileStreamFactory: 1,
   isServer: 1,
   keepAliveInitialDelay: 1,
   localAddress: 1,
@@ -108,6 +109,7 @@ class ConnectionConfig {
       ? 10 * 1000
       : options.connectTimeout;
     this.insecureAuth = options.insecureAuth || false;
+    this.infileStreamFactory = options.infileStreamFactory || undefined;
     this.supportBigNumbers = options.supportBigNumbers || false;
     this.bigNumberStrings = options.bigNumberStrings || false;
     this.decimalNumbers = options.decimalNumbers || false;
diff --git a/typings/mysql/lib/Connection.d.ts b/typings/mysql/lib/Connection.d.ts
index d0e2cc6929..7112ef68a3 100644
--- a/typings/mysql/lib/Connection.d.ts
+++ b/typings/mysql/lib/Connection.d.ts
@@ -4,6 +4,7 @@
 // Modifications copyright (c) 2021, Oracle and/or its affiliates.
 
 import { EventEmitter } from 'events';
+import { Readable } from 'stream';
 import { Query, QueryError } from './protocol/sequences/Query.js';
 import { Prepare, PrepareStatementInfo } from './protocol/sequences/Prepare.js';
 import {
@@ -165,6 +166,11 @@ export interface ConnectionOptions {
    */
   insecureAuth?: boolean;
 
+  /**
+   * By specifying a function that returns a readable stream, an arbitrary stream can be sent when sending a local fs file.
+   */
+  infileStreamFactory?: (path: string) => Readable;
+
   /**
    * Determines if column values should be converted to native JavaScript types. It is not recommended (and may go away / change in the future)
    * to disable type casting, but you can currently do so on either the connection or query level. (Default: true)