Add streaming infiles

Author: kesmit13

singlestoredb/connection.py

@@ -2,6 +2,8 @@
 """SingleStoreDB connections and cursors."""
 import abc
 import inspect
+import io
+import queue
 import re
 import warnings
 import weakref
@@ -496,6 +498,14 @@ def close(self) -> None:
     def execute(
         self, query: str,
         args: Optional[Union[Sequence[Any], Dict[str, Any], Any]] = None,
+        infile_stream: Optional[
+            Union[
+                io.RawIOBase,
+                io.TextIOBase,
+                Iterator[Union[bytes, str]],
+                queue.Queue[Union[bytes, str]],
+            ]
+        ] = None,
     ) -> int:
         """
         Execute a SQL statement.
@@ -510,6 +520,8 @@ def execute(
             The SQL statement to execute
         args : Sequence or dict, optional
             Parameters to substitute into the SQL code
+        infile_stream : io.RawIOBase or io.TextIOBase or Iterator[bytes|str], optional
+            Data stream for ``LOCAL INFILE`` statement
 
         Examples
         --------

singlestoredb/http/connection.py

@@ -3,9 +3,11 @@
 import datetime
 import decimal
 import functools
+import io
 import json
 import math
 import os
+import queue
 import re
 import time
 from base64 import b64decode
@@ -420,6 +422,14 @@ def close(self) -> None:
     def execute(
         self, query: str,
         args: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
+        infile_stream: Optional[
+            Union[
+                io.RawIOBase,
+                io.TextIOBase,
+                Iterable[Union[bytes, str]],
+                queue.Queue[Union[bytes, str]],
+            ]
+        ] = None,
     ) -> int:
         """
         Execute a SQL statement.
@@ -432,7 +442,7 @@ def execute(
             Parameters to substitute into the SQL code
 
         """
-        return self._execute(query, args)
+        return self._execute(query, args, infile_stream=infile_stream)
 
     def _validate_param_subs(
         self, query: str,
@@ -496,6 +506,14 @@ def _execute(
         self, oper: str,
         params: Optional[Union[Sequence[Any], Dict[str, Any]]] = None,
         is_callproc: bool = False,
+        infile_stream: Optional[
+            Union[
+                io.RawIOBase,
+                io.TextIOBase,
+                Iterable[Union[bytes, str]],
+                queue.Queue[Union[bytes, str]],
+            ]
+        ] = None,
     ) -> int:
         self._descriptions = []
         self._schemas = []

singlestoredb/mysql/connection.py

@@ -5,12 +5,15 @@
 # https://dev.mysql.com/doc/refman/5.5/en/error-handling.html
 import errno
 import functools
+import io
 import os
+import queue
 import socket
 import struct
 import sys
 import traceback
 import warnings
+from typing import Iterable
 
 try:
     import _singlestoredb_accel
@@ -353,6 +356,7 @@ def __init__(  # noqa: C901
             )
 
         self._local_infile = bool(local_infile)
+        self._local_infile_stream = None
         if self._local_infile:
             client_flag |= CLIENT.LOCAL_FILES
         if multi_statements:
@@ -843,7 +847,7 @@ def cursor(self):
         return self.cursorclass(self)
 
     # The following methods are INTERNAL USE ONLY (called from Cursor)
-    def query(self, sql, unbuffered=False):
+    def query(self, sql, unbuffered=False, infile_stream=None):
         """
         Run a query on the server.
 
@@ -859,8 +863,10 @@ def query(self, sql, unbuffered=False):
         else:
             if isinstance(sql, str):
                 sql = sql.encode(self.encoding, 'surrogateescape')
+            self._local_infile_stream = infile_stream
             self._execute_command(COMMAND.COM_QUERY, sql)
             self._affected_rows = self._read_query_result(unbuffered=unbuffered)
+            self._local_infile_stream = None
         return self._affected_rows
 
     def next_result(self, unbuffered=False):
@@ -1871,24 +1877,82 @@ def __init__(self, filename, connection):
     def send_data(self):
         """Send data packets from the local file to the server"""
         if not self.connection._sock:
-            raise err.InterfaceError(0, '')
+            raise err.InterfaceError(0, 'Connection is closed')
+
         conn = self.connection
+        infile = conn._local_infile_stream
+
+        # 16KB is efficient enough
+        packet_size = min(conn.max_allowed_packet, 16 * 1024)
 
         try:
-            with open(self.filename, 'rb') as open_file:
-                packet_size = min(
-                    conn.max_allowed_packet, 16 * 1024,
-                )  # 16KB is efficient enough
-                while True:
-                    chunk = open_file.read(packet_size)
-                    if not chunk:
-                        break
-                    conn.write_packet(chunk)
-        except OSError:
-            raise err.OperationalError(
-                ER.FILE_NOT_FOUND,
-                f"Can't find file '{self.filename}'",
-            )
+
+            if self.filename in [':stream:', b':stream:']:
+
+                if infile is None:
+                    raise err.OperationalError(
+                        ER.FILE_NOT_FOUND,
+                        ':stream: specified for LOCAL INFILE, but no stream was supplied',
+                    )
+
+                # Binary IO
+                elif isinstance(infile, io.RawIOBase):
+                    while True:
+                        chunk = infile.read(packet_size)
+                        if not chunk:
+                            break
+                        conn.write_packet(chunk)
+
+                # Text IO
+                elif isinstance(infile, io.TextIOBase):
+                    while True:
+                        chunk = infile.read(packet_size)
+                        if not chunk:
+                            break
+                        conn.write_packet(chunk.encode('utf8'))
+
+                # Iterable of bytes or str
+                elif isinstance(infile, Iterable):
+                    for chunk in infile:
+                        if not chunk:
+                            continue
+                        if isinstance(chunk, str):
+                            conn.write_packet(chunk.encode('utf8'))
+                        else:
+                            conn.write_packet(chunk)
+
+                # Queue (empty value ends the iteration)
+                elif isinstance(infile, queue.Queue):
+                    while True:
+                        chunk = infile.get()
+                        if not chunk:
+                            break
+                        if isinstance(chunk, str):
+                            conn.write_packet(chunk.encode('utf8'))
+                        else:
+                            conn.write_packet(chunk)
+
+                else:
+                    raise err.OperationalError(
+                        ER.FILE_NOT_FOUND,
+                        ':stream: specified for LOCAL INFILE, ' +
+                        f'but stream type is unrecognized: {infile}',
+                    )
+
+            else:
+                try:
+                    with open(self.filename, 'rb') as open_file:
+                        while True:
+                            chunk = open_file.read(packet_size)
+                            if not chunk:
+                                break
+                            conn.write_packet(chunk)
+                except OSError:
+                    raise err.OperationalError(
+                        ER.FILE_NOT_FOUND,
+                        f"Can't find file '{self.filename!s}'",
+                    )
+
         finally:
             if not conn._closed:
                 # send the empty packet to signify we are done sending data

singlestoredb/mysql/cursors.py

@@ -178,7 +178,7 @@ def mogrify(self, query, args=None):
 
         return query
 
-    def execute(self, query, args=None):
+    def execute(self, query, args=None, infile_stream=None):
         """
         Execute a query.
 
@@ -192,6 +192,8 @@ def execute(self, query, args=None):
             Query to execute.
         args : Sequence[Any] or Dict[str, Any] or Any, optional
             Parameters used with query. (optional)
+        infile_stream : io.BytesIO or Iterator[bytes], optional
+            Data stream for ``LOCAL INFILE`` statements
 
         Returns
         -------
@@ -205,7 +207,7 @@ def execute(self, query, args=None):
 
         query = self.mogrify(query, args)
 
-        result = self._query(query)
+        result = self._query(query, infile_stream=infile_stream)
         self._executed = query
         return result
 
@@ -387,10 +389,10 @@ def scroll(self, value, mode='relative'):
             raise IndexError('out of range')
         self._rownumber = r
 
-    def _query(self, q):
+    def _query(self, q, infile_stream=None):
         conn = self._get_db()
         self._clear_result()
-        conn.query(q)
+        conn.query(q, infile_stream=infile_stream)
         self._do_get_result()
         return self.rowcount
 
@@ -680,10 +682,10 @@ def close(self):
 
     __del__ = close
 
-    def _query(self, q):
+    def _query(self, q, infile_stream=None):
         conn = self._get_db()
         self._clear_result()
-        conn.query(q, unbuffered=True)
+        conn.query(q, unbuffered=True, infile_stream=infile_stream)
         self._do_get_result()
         return self.rowcount