FazBrowse GitHub Viewer | Trending |
URL:
| Home
Tools: [Download Repo ZIP]   [Original HTTPS Page]

Refactor LOAD LOCAL INFILE handling by methane · Pull Request #1264 · PyMySQL/PyMySQL · GitHub

Merged
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension .py  (1) All 1 file type selected
Viewed files
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Unified
Split
Hide whitespace
Diff view
Unified
Split
Hide whitespace
74 changes: 30 additions & 44 deletions pymysql/connections.py
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters. Learn more about bidirectional Unicode characters
Original file line number Diff line number Diff line change
Expand Up @@ -728,7 +728,7 @@ def connect(self, sock=None):
except BaseException as e:
self._force_close()

if isinstance(e, (OSError, IOError)):
if isinstance(e, OSError):
exc = err.OperationalError(
CR.CR_CONN_HOST_ERROR,
f"Can't connect to MySQL server on {self.host!r} ({e})",
Expand Down Expand Up @@ -1275,8 +1275,8 @@ def init_unbuffered_query(self):
self.affected_rows = 18446744073709551615
self.unbuffered_active = True

def _read_ok_packet(self, first_packet):
ok_packet = OKPacketWrapper(first_packet)
def _read_ok_packet(self, packet):
ok_packet = OKPacketWrapper(packet)
self.affected_rows = ok_packet.affected_rows
self.insert_id = ok_packet.insert_id
self.server_status = ok_packet.server_status
Expand All @@ -1285,25 +1285,24 @@ def _read_ok_packet(self, first_packet):
self.has_next = ok_packet.has_next

def _read_load_local_packet(self, first_packet):
if not self.connection._local_infile:
conn: Connection = self.connection
if not conn._local_infile:
raise RuntimeError(
"**WARN**: Received LOAD_LOCAL packet but local_infile option is false."
)
load_packet = LoadLocalPacketWrapper(first_packet)
sender = LoadLocalFile(load_packet.filename, self.connection)
try:
sender.send_data()
except:
self.connection._read_packet() # skip ok packet
raise
_send_local_file(load_packet.filename, conn)
finally:
# send the empty packet to signify we are done sending data
conn.write_packet(b"")
ok_packet = conn._read_packet()
# If an error occurs while sending the file, exit here without handling
# the OK packet.

ok_packet = self.connection._read_packet()
if (
not ok_packet.is_ok_packet()
): # pragma: no cover - upstream induced protocol error
if not ok_packet.is_ok_packet():
Comment on lines +1296 to +1303
raise err.OperationalError(
CR.CR_COMMANDS_OUT_OF_SYNC,
"Commands Out of Sync",
CR.CR_COMMANDS_OUT_OF_SYNC, "Commands Out of Sync"
)
self._read_ok_packet(ok_packet)

Expand Down Expand Up @@ -1442,33 +1441,20 @@ def _get_descriptions(self):
self.description = tuple(description)


class LoadLocalFile:
def __init__(self, filename, connection):
self.filename = filename
self.connection = connection

def send_data(self):
"""Send data packets from the local file to the server"""
if not self.connection._sock:
raise err.InterfaceError(0, "")
conn: Connection = self.connection
def _send_local_file(filename: str, conn: Connection):
"""Send data packets from the local file to the server"""
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}'",
)
finally:
if not conn._closed:
# send the empty packet to signify we are done sending data
conn.write_packet(b"")
try:
with open(filename, "rb") as file:
# 16KB is efficient enough
while True:
chunk = file.read(packet_size)
if not chunk:
break
conn.write_packet(chunk)
except OSError as e:
raise err.OperationalError(
ER.FILE_NOT_FOUND,
f"Can't open file '{filename}': {e}",
)

Back | FazBrowse Home | New Git URL