connections.py 52 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170117111721173117411751176117711781179118011811182118311841185118611871188118911901191119211931194119511961197119811991200120112021203120412051206120712081209121012111212121312141215121612171218121912201221122212231224122512261227122812291230123112321233123412351236123712381239124012411242124312441245124612471248124912501251125212531254125512561257125812591260126112621263126412651266126712681269127012711272127312741275127612771278127912801281128212831284128512861287128812891290129112921293129412951296129712981299130013011302130313041305130613071308130913101311131213131314131513161317131813191320132113221323132413251326132713281329133013311332133313341335133613371338133913401341134213431344134513461347134813491350135113521353135413551356135713581359136013611362136313641365136613671368136913701371137213731374137513761377137813791380138113821383138413851386138713881389139013911392139313941395139613971398139914001401140214031404140514061407140814091410141114121413141414151416141714181419142014211422142314241425142614271428142914301431
  1. # Python implementation of the MySQL client-server protocol
  2. # http://dev.mysql.com/doc/internals/en/client-server-protocol.html
  3. # Error codes:
  4. # https://dev.mysql.com/doc/refman/5.5/en/error-handling.html
  5. import errno
  6. import os
  7. import socket
  8. import struct
  9. import sys
  10. import traceback
  11. import warnings
  12. from . import _auth
  13. from .charset import charset_by_name, charset_by_id
  14. from .constants import CLIENT, COMMAND, CR, ER, FIELD_TYPE, SERVER_STATUS
  15. from . import converters
  16. from .cursors import Cursor
  17. from .optionfile import Parser
  18. from .protocol import (
  19. dump_packet,
  20. MysqlPacket,
  21. FieldDescriptorPacket,
  22. OKPacketWrapper,
  23. EOFPacketWrapper,
  24. LoadLocalPacketWrapper,
  25. )
  26. from . import err, VERSION_STRING
  27. try:
  28. import ssl
  29. SSL_ENABLED = True
  30. except ImportError:
  31. ssl = None
  32. SSL_ENABLED = False
  33. try:
  34. import getpass
  35. DEFAULT_USER = getpass.getuser()
  36. del getpass
  37. except (ImportError, KeyError):
  38. # KeyError occurs when there's no entry in OS database for a current user.
  39. DEFAULT_USER = None
  40. DEBUG = False
  41. TEXT_TYPES = {
  42. FIELD_TYPE.BIT,
  43. FIELD_TYPE.BLOB,
  44. FIELD_TYPE.LONG_BLOB,
  45. FIELD_TYPE.MEDIUM_BLOB,
  46. FIELD_TYPE.STRING,
  47. FIELD_TYPE.TINY_BLOB,
  48. FIELD_TYPE.VAR_STRING,
  49. FIELD_TYPE.VARCHAR,
  50. FIELD_TYPE.GEOMETRY,
  51. }
  52. DEFAULT_CHARSET = "utf8mb4"
  53. MAX_PACKET_LEN = 2**24 - 1
  54. def _pack_int24(n):
  55. return struct.pack("<I", n)[:3]
  56. # https://dev.mysql.com/doc/internals/en/integer.html#packet-Protocol::LengthEncodedInteger
  57. def _lenenc_int(i):
  58. if i < 0:
  59. raise ValueError(
  60. "Encoding %d is less than 0 - no representation in LengthEncodedInteger" % i
  61. )
  62. elif i < 0xFB:
  63. return bytes([i])
  64. elif i < (1 << 16):
  65. return b"\xfc" + struct.pack("<H", i)
  66. elif i < (1 << 24):
  67. return b"\xfd" + struct.pack("<I", i)[:3]
  68. elif i < (1 << 64):
  69. return b"\xfe" + struct.pack("<Q", i)
  70. else:
  71. raise ValueError(
  72. f"Encoding {i:x} is larger than {1 << 64:x} - no representation in LengthEncodedInteger"
  73. )
  74. class Connection:
  75. """
  76. Representation of a socket with a mysql server.
  77. The proper way to get an instance of this class is to call
  78. connect().
  79. Establish a connection to the MySQL database. Accepts several
  80. arguments:
  81. :param host: Host where the database server is located.
  82. :param user: Username to log in as.
  83. :param password: Password to use.
  84. :param database: Database to use, None to not use a particular one.
  85. :param port: MySQL port to use, default is usually OK. (default: 3306)
  86. :param bind_address: When the client has multiple network interfaces, specify
  87. the interface from which to connect to the host. Argument can be
  88. a hostname or an IP address.
  89. :param unix_socket: Use a unix socket rather than TCP/IP.
  90. :param read_timeout: The timeout for reading from the connection in seconds.
  91. (default: None - no timeout)
  92. :param write_timeout: The timeout for writing to the connection in seconds.
  93. (default: None - no timeout)
  94. :param str charset: Charset to use.
  95. :param str collation: Collation name to use.
  96. :param sql_mode: Default SQL_MODE to use.
  97. :param read_default_file:
  98. Specifies my.cnf file to read these parameters from under the [client] section.
  99. :param conv:
  100. Conversion dictionary to use instead of the default one.
  101. This is used to provide custom marshalling and unmarshalling of types.
  102. See converters.
  103. :param use_unicode:
  104. Whether or not to default to unicode strings.
  105. This option defaults to true.
  106. :param client_flag: Custom flags to send to MySQL. Find potential values in constants.CLIENT.
  107. :param cursorclass: Custom cursor class to use.
  108. :param init_command: Initial SQL statement to run when connection is established.
  109. :param connect_timeout: The timeout for connecting to the database in seconds.
  110. (default: 10, min: 1, max: 31536000)
  111. :param ssl: A dict of arguments similar to mysql_ssl_set()'s parameters or an ssl.SSLContext.
  112. :param ssl_ca: Path to the file that contains a PEM-formatted CA certificate.
  113. :param ssl_cert: Path to the file that contains a PEM-formatted client certificate.
  114. :param ssl_disabled: A boolean value that disables usage of TLS.
  115. :param ssl_key: Path to the file that contains a PEM-formatted private key for
  116. the client certificate.
  117. :param ssl_key_password: The password for the client certificate private key.
  118. :param ssl_verify_cert: Set to true to check the server certificate's validity.
  119. :param ssl_verify_identity: Set to true to check the server's identity.
  120. :param read_default_group: Group to read from in the configuration file.
  121. :param autocommit: Autocommit mode. None means use server default. (default: False)
  122. :param local_infile: Boolean to enable the use of LOAD DATA LOCAL command. (default: False)
  123. :param max_allowed_packet: Max size of packet sent to server in bytes. (default: 16MB)
  124. Only used to limit size of "LOAD LOCAL INFILE" data packet smaller than default (16KB).
  125. :param defer_connect: Don't explicitly connect on construction - wait for connect call.
  126. (default: False)
  127. :param auth_plugin_map: A dict of plugin names to a class that processes that plugin.
  128. The class will take the Connection object as the argument to the constructor.
  129. The class needs an authenticate method taking an authentication packet as
  130. an argument. For the dialog plugin, a prompt(echo, prompt) method can be used
  131. (if no authenticate method) for returning a string from the user. (experimental)
  132. :param server_public_key: SHA256 authentication plugin public key value. (default: None)
  133. :param binary_prefix: Add _binary prefix on bytes and bytearray. (default: False)
  134. :param compress: Not supported.
  135. :param named_pipe: Not supported.
  136. :param db: **DEPRECATED** Alias for database.
  137. :param passwd: **DEPRECATED** Alias for password.
  138. See `Connection <https://www.python.org/dev/peps/pep-0249/#connection-objects>`_ in the
  139. specification.
  140. """
  141. _sock = None
  142. _auth_plugin_name = ""
  143. _closed = False
  144. _secure = False
  145. def __init__(
  146. self,
  147. *,
  148. user=None, # The first four arguments is based on DB-API 2.0 recommendation.
  149. password="",
  150. host=None,
  151. database=None,
  152. unix_socket=None,
  153. port=0,
  154. charset="",
  155. collation=None,
  156. sql_mode=None,
  157. read_default_file=None,
  158. conv=None,
  159. use_unicode=True,
  160. client_flag=0,
  161. cursorclass=Cursor,
  162. init_command=None,
  163. connect_timeout=10,
  164. read_default_group=None,
  165. autocommit=False,
  166. local_infile=False,
  167. max_allowed_packet=16 * 1024 * 1024,
  168. defer_connect=False,
  169. auth_plugin_map=None,
  170. read_timeout=None,
  171. write_timeout=None,
  172. bind_address=None,
  173. binary_prefix=False,
  174. program_name=None,
  175. server_public_key=None,
  176. ssl=None,
  177. ssl_ca=None,
  178. ssl_cert=None,
  179. ssl_disabled=None,
  180. ssl_key=None,
  181. ssl_key_password=None,
  182. ssl_verify_cert=None,
  183. ssl_verify_identity=None,
  184. compress=None, # not supported
  185. named_pipe=None, # not supported
  186. passwd=None, # deprecated
  187. db=None, # deprecated
  188. ):
  189. if db is not None and database is None:
  190. # We will raise warning in 2022 or later.
  191. # See https://github.com/PyMySQL/PyMySQL/issues/939
  192. # warnings.warn("'db' is deprecated, use 'database'", DeprecationWarning, 3)
  193. database = db
  194. if passwd is not None and not password:
  195. # We will raise warning in 2022 or later.
  196. # See https://github.com/PyMySQL/PyMySQL/issues/939
  197. # warnings.warn(
  198. # "'passwd' is deprecated, use 'password'", DeprecationWarning, 3
  199. # )
  200. password = passwd
  201. if compress or named_pipe:
  202. raise NotImplementedError(
  203. "compress and named_pipe arguments are not supported"
  204. )
  205. self._local_infile = bool(local_infile)
  206. if self._local_infile:
  207. client_flag |= CLIENT.LOCAL_FILES
  208. if read_default_group and not read_default_file:
  209. if sys.platform.startswith("win"):
  210. read_default_file = "c:\\my.ini"
  211. else:
  212. read_default_file = "/etc/my.cnf"
  213. if read_default_file:
  214. if not read_default_group:
  215. read_default_group = "client"
  216. cfg = Parser()
  217. cfg.read(os.path.expanduser(read_default_file))
  218. def _config(key, arg):
  219. if arg:
  220. return arg
  221. try:
  222. return cfg.get(read_default_group, key)
  223. except Exception:
  224. return arg
  225. user = _config("user", user)
  226. password = _config("password", password)
  227. host = _config("host", host)
  228. database = _config("database", database)
  229. unix_socket = _config("socket", unix_socket)
  230. port = int(_config("port", port))
  231. bind_address = _config("bind-address", bind_address)
  232. charset = _config("default-character-set", charset)
  233. if not ssl:
  234. ssl = {}
  235. if isinstance(ssl, dict):
  236. for key in ["ca", "capath", "cert", "key", "password", "cipher"]:
  237. value = _config("ssl-" + key, ssl.get(key))
  238. if value:
  239. ssl[key] = value
  240. self.ssl = False
  241. if not ssl_disabled:
  242. if ssl_ca or ssl_cert or ssl_key or ssl_verify_cert or ssl_verify_identity:
  243. ssl = {
  244. "ca": ssl_ca,
  245. "check_hostname": bool(ssl_verify_identity),
  246. "verify_mode": ssl_verify_cert
  247. if ssl_verify_cert is not None
  248. else False,
  249. }
  250. if ssl_cert is not None:
  251. ssl["cert"] = ssl_cert
  252. if ssl_key is not None:
  253. ssl["key"] = ssl_key
  254. if ssl_key_password is not None:
  255. ssl["password"] = ssl_key_password
  256. if ssl:
  257. if not SSL_ENABLED:
  258. raise NotImplementedError("ssl module not found")
  259. self.ssl = True
  260. client_flag |= CLIENT.SSL
  261. self.ctx = self._create_ssl_ctx(ssl)
  262. self.host = host or "localhost"
  263. self.port = port or 3306
  264. if type(self.port) is not int:
  265. raise ValueError("port should be of type int")
  266. self.user = user or DEFAULT_USER
  267. self.password = password or b""
  268. if isinstance(self.password, str):
  269. self.password = self.password.encode("latin1")
  270. self.db = database
  271. self.unix_socket = unix_socket
  272. self.bind_address = bind_address
  273. if not (0 < connect_timeout <= 31536000):
  274. raise ValueError("connect_timeout should be >0 and <=31536000")
  275. self.connect_timeout = connect_timeout or None
  276. if read_timeout is not None and read_timeout <= 0:
  277. raise ValueError("read_timeout should be > 0")
  278. self._read_timeout = read_timeout
  279. if write_timeout is not None and write_timeout <= 0:
  280. raise ValueError("write_timeout should be > 0")
  281. self._write_timeout = write_timeout
  282. self.charset = charset or DEFAULT_CHARSET
  283. self.collation = collation
  284. self.use_unicode = use_unicode
  285. self.encoding = charset_by_name(self.charset).encoding
  286. client_flag |= CLIENT.CAPABILITIES
  287. if self.db:
  288. client_flag |= CLIENT.CONNECT_WITH_DB
  289. self.client_flag = client_flag
  290. self.cursorclass = cursorclass
  291. self._result = None
  292. self._affected_rows = 0
  293. self.host_info = "Not connected"
  294. # specified autocommit mode. None means use server default.
  295. self.autocommit_mode = autocommit
  296. if conv is None:
  297. conv = converters.conversions
  298. # Need for MySQLdb compatibility.
  299. self.encoders = {k: v for (k, v) in conv.items() if type(k) is not int}
  300. self.decoders = {k: v for (k, v) in conv.items() if type(k) is int}
  301. self.sql_mode = sql_mode
  302. self.init_command = init_command
  303. self.max_allowed_packet = max_allowed_packet
  304. self._auth_plugin_map = auth_plugin_map or {}
  305. self._binary_prefix = binary_prefix
  306. self.server_public_key = server_public_key
  307. self._connect_attrs = {
  308. "_client_name": "pymysql",
  309. "_client_version": VERSION_STRING,
  310. "_pid": str(os.getpid()),
  311. }
  312. if program_name:
  313. self._connect_attrs["program_name"] = program_name
  314. if defer_connect:
  315. self._sock = None
  316. else:
  317. self.connect()
  318. def __enter__(self):
  319. return self
  320. def __exit__(self, *exc_info):
  321. del exc_info
  322. self.close()
  323. def _create_ssl_ctx(self, sslp):
  324. if isinstance(sslp, ssl.SSLContext):
  325. return sslp
  326. ca = sslp.get("ca")
  327. capath = sslp.get("capath")
  328. hasnoca = ca is None and capath is None
  329. ctx = ssl.create_default_context(cafile=ca, capath=capath)
  330. ctx.check_hostname = not hasnoca and sslp.get("check_hostname", True)
  331. verify_mode_value = sslp.get("verify_mode")
  332. if verify_mode_value is None:
  333. ctx.verify_mode = ssl.CERT_NONE if hasnoca else ssl.CERT_REQUIRED
  334. elif isinstance(verify_mode_value, bool):
  335. ctx.verify_mode = ssl.CERT_REQUIRED if verify_mode_value else ssl.CERT_NONE
  336. else:
  337. if isinstance(verify_mode_value, str):
  338. verify_mode_value = verify_mode_value.lower()
  339. if verify_mode_value in ("none", "0", "false", "no"):
  340. ctx.verify_mode = ssl.CERT_NONE
  341. elif verify_mode_value == "optional":
  342. ctx.verify_mode = ssl.CERT_OPTIONAL
  343. elif verify_mode_value in ("required", "1", "true", "yes"):
  344. ctx.verify_mode = ssl.CERT_REQUIRED
  345. else:
  346. ctx.verify_mode = ssl.CERT_NONE if hasnoca else ssl.CERT_REQUIRED
  347. if "cert" in sslp:
  348. ctx.load_cert_chain(
  349. sslp["cert"], keyfile=sslp.get("key"), password=sslp.get("password")
  350. )
  351. if "cipher" in sslp:
  352. ctx.set_ciphers(sslp["cipher"])
  353. ctx.options |= ssl.OP_NO_SSLv2
  354. ctx.options |= ssl.OP_NO_SSLv3
  355. return ctx
  356. def close(self):
  357. """
  358. Send the quit message and close the socket.
  359. See `Connection.close() <https://www.python.org/dev/peps/pep-0249/#Connection.close>`_
  360. in the specification.
  361. :raise Error: If the connection is already closed.
  362. """
  363. if self._closed:
  364. raise err.Error("Already closed")
  365. self._closed = True
  366. if self._sock is None:
  367. return
  368. send_data = struct.pack("<iB", 1, COMMAND.COM_QUIT)
  369. try:
  370. self._write_bytes(send_data)
  371. except Exception:
  372. pass
  373. finally:
  374. self._force_close()
  375. @property
  376. def open(self):
  377. """Return True if the connection is open."""
  378. return self._sock is not None
  379. def _force_close(self):
  380. """Close connection without QUIT message."""
  381. if self._sock:
  382. try:
  383. self._sock.close()
  384. except: # noqa
  385. pass
  386. self._sock = None
  387. self._rfile = None
  388. __del__ = _force_close
  389. def autocommit(self, value):
  390. self.autocommit_mode = bool(value)
  391. current = self.get_autocommit()
  392. if value != current:
  393. self._send_autocommit_mode()
  394. def get_autocommit(self):
  395. return bool(self.server_status & SERVER_STATUS.SERVER_STATUS_AUTOCOMMIT)
  396. def _read_ok_packet(self):
  397. pkt = self._read_packet()
  398. if not pkt.is_ok_packet():
  399. raise err.OperationalError(
  400. CR.CR_COMMANDS_OUT_OF_SYNC,
  401. "Command Out of Sync",
  402. )
  403. ok = OKPacketWrapper(pkt)
  404. self.server_status = ok.server_status
  405. return ok
  406. def _send_autocommit_mode(self):
  407. """Set whether or not to commit after every execute()."""
  408. self._execute_command(
  409. COMMAND.COM_QUERY, "SET AUTOCOMMIT = %s" % self.escape(self.autocommit_mode)
  410. )
  411. self._read_ok_packet()
  412. def begin(self):
  413. """Begin transaction."""
  414. self._execute_command(COMMAND.COM_QUERY, "BEGIN")
  415. self._read_ok_packet()
  416. def commit(self):
  417. """
  418. Commit changes to stable storage.
  419. See `Connection.commit() <https://www.python.org/dev/peps/pep-0249/#commit>`_
  420. in the specification.
  421. """
  422. self._execute_command(COMMAND.COM_QUERY, "COMMIT")
  423. self._read_ok_packet()
  424. def rollback(self):
  425. """
  426. Roll back the current transaction.
  427. See `Connection.rollback() <https://www.python.org/dev/peps/pep-0249/#rollback>`_
  428. in the specification.
  429. """
  430. self._execute_command(COMMAND.COM_QUERY, "ROLLBACK")
  431. self._read_ok_packet()
  432. def show_warnings(self):
  433. """Send the "SHOW WARNINGS" SQL command."""
  434. self._execute_command(COMMAND.COM_QUERY, "SHOW WARNINGS")
  435. result = MySQLResult(self)
  436. result.read()
  437. return result.rows
  438. def select_db(self, db):
  439. """
  440. Set current db.
  441. :param db: The name of the db.
  442. """
  443. self._execute_command(COMMAND.COM_INIT_DB, db)
  444. self._read_ok_packet()
  445. def escape(self, obj, mapping=None):
  446. """Escape whatever value is passed.
  447. Non-standard, for internal use; do not use this in your applications.
  448. """
  449. if isinstance(obj, str):
  450. return "'" + self.escape_string(obj) + "'"
  451. if isinstance(obj, (bytes, bytearray)):
  452. ret = self._quote_bytes(obj)
  453. if self._binary_prefix:
  454. ret = "_binary" + ret
  455. return ret
  456. return converters.escape_item(obj, self.charset, mapping=mapping)
  457. def literal(self, obj):
  458. """Alias for escape().
  459. Non-standard, for internal use; do not use this in your applications.
  460. """
  461. return self.escape(obj, self.encoders)
  462. def escape_string(self, s):
  463. if self.server_status & SERVER_STATUS.SERVER_STATUS_NO_BACKSLASH_ESCAPES:
  464. return s.replace("'", "''")
  465. return converters.escape_string(s)
  466. def _quote_bytes(self, s):
  467. if self.server_status & SERVER_STATUS.SERVER_STATUS_NO_BACKSLASH_ESCAPES:
  468. return "'{}'".format(
  469. s.replace(b"'", b"''").decode("ascii", "surrogateescape")
  470. )
  471. return converters.escape_bytes(s)
  472. def cursor(self, cursor=None):
  473. """
  474. Create a new cursor to execute queries with.
  475. :param cursor: The type of cursor to create. None means use Cursor.
  476. :type cursor: :py:class:`Cursor`, :py:class:`SSCursor`, :py:class:`DictCursor`,
  477. or :py:class:`SSDictCursor`.
  478. """
  479. if cursor:
  480. return cursor(self)
  481. return self.cursorclass(self)
  482. # The following methods are INTERNAL USE ONLY (called from Cursor)
  483. def query(self, sql, unbuffered=False):
  484. # if DEBUG:
  485. # print("DEBUG: sending query:", sql)
  486. if isinstance(sql, str):
  487. sql = sql.encode(self.encoding, "surrogateescape")
  488. self._execute_command(COMMAND.COM_QUERY, sql)
  489. self._affected_rows = self._read_query_result(unbuffered=unbuffered)
  490. return self._affected_rows
  491. def next_result(self, unbuffered=False):
  492. self._affected_rows = self._read_query_result(unbuffered=unbuffered)
  493. return self._affected_rows
  494. def affected_rows(self):
  495. return self._affected_rows
  496. def kill(self, thread_id):
  497. arg = struct.pack("<I", thread_id)
  498. self._execute_command(COMMAND.COM_PROCESS_KILL, arg)
  499. return self._read_ok_packet()
  500. def ping(self, reconnect=True):
  501. """
  502. Check if the server is alive.
  503. :param reconnect: If the connection is closed, reconnect.
  504. :type reconnect: boolean
  505. :raise Error: If the connection is closed and reconnect=False.
  506. """
  507. if self._sock is None:
  508. if reconnect:
  509. self.connect()
  510. reconnect = False
  511. else:
  512. raise err.Error("Already closed")
  513. try:
  514. self._execute_command(COMMAND.COM_PING, "")
  515. self._read_ok_packet()
  516. except Exception:
  517. if reconnect:
  518. self.connect()
  519. self.ping(False)
  520. else:
  521. raise
  522. def set_charset(self, charset):
  523. """Deprecated. Use set_character_set() instead."""
  524. # This function has been implemented in old PyMySQL.
  525. # But this name is different from MySQLdb.
  526. # So we keep this function for compatibility and add
  527. # new set_character_set() function.
  528. self.set_character_set(charset)
  529. def set_character_set(self, charset, collation=None):
  530. """
  531. Set charaset (and collation)
  532. Send "SET NAMES charset [COLLATE collation]" query.
  533. Update Connection.encoding based on charset.
  534. """
  535. # Make sure charset is supported.
  536. encoding = charset_by_name(charset).encoding
  537. if collation:
  538. query = f"SET NAMES {charset} COLLATE {collation}"
  539. else:
  540. query = f"SET NAMES {charset}"
  541. self._execute_command(COMMAND.COM_QUERY, query)
  542. self._read_packet()
  543. self.charset = charset
  544. self.encoding = encoding
  545. self.collation = collation
  546. def connect(self, sock=None):
  547. self._closed = False
  548. try:
  549. if sock is None:
  550. if self.unix_socket:
  551. sock = socket.socket(socket.AF_UNIX, socket.SOCK_STREAM)
  552. sock.settimeout(self.connect_timeout)
  553. sock.connect(self.unix_socket)
  554. self.host_info = "Localhost via UNIX socket"
  555. self._secure = True
  556. if DEBUG:
  557. print("connected using unix_socket")
  558. else:
  559. kwargs = {}
  560. if self.bind_address is not None:
  561. kwargs["source_address"] = (self.bind_address, 0)
  562. while True:
  563. try:
  564. sock = socket.create_connection(
  565. (self.host, self.port), self.connect_timeout, **kwargs
  566. )
  567. break
  568. except OSError as e:
  569. if e.errno == errno.EINTR:
  570. continue
  571. raise
  572. self.host_info = "socket %s:%d" % (self.host, self.port)
  573. if DEBUG:
  574. print("connected using socket")
  575. sock.setsockopt(socket.IPPROTO_TCP, socket.TCP_NODELAY, 1)
  576. sock.setsockopt(socket.SOL_SOCKET, socket.SO_KEEPALIVE, 1)
  577. sock.settimeout(None)
  578. self._sock = sock
  579. self._rfile = sock.makefile("rb")
  580. self._next_seq_id = 0
  581. self._get_server_information()
  582. self._request_authentication()
  583. # Send "SET NAMES" query on init for:
  584. # - Ensure charaset (and collation) is set to the server.
  585. # - collation_id in handshake packet may be ignored.
  586. # - If collation is not specified, we don't know what is server's
  587. # default collation for the charset. For example, default collation
  588. # of utf8mb4 is:
  589. # - MySQL 5.7, MariaDB 10.x: utf8mb4_general_ci
  590. # - MySQL 8.0: utf8mb4_0900_ai_ci
  591. #
  592. # Reference:
  593. # - https://github.com/PyMySQL/PyMySQL/issues/1092
  594. # - https://github.com/wagtail/wagtail/issues/9477
  595. # - https://zenn.dev/methane/articles/2023-mysql-collation (Japanese)
  596. self.set_character_set(self.charset, self.collation)
  597. if self.sql_mode is not None:
  598. c = self.cursor()
  599. c.execute("SET sql_mode=%s", (self.sql_mode,))
  600. c.close()
  601. if self.init_command is not None:
  602. c = self.cursor()
  603. c.execute(self.init_command)
  604. c.close()
  605. if self.autocommit_mode is not None:
  606. self.autocommit(self.autocommit_mode)
  607. except BaseException as e:
  608. self._rfile = None
  609. if sock is not None:
  610. try:
  611. sock.close()
  612. except: # noqa
  613. pass
  614. if isinstance(e, (OSError, IOError)):
  615. exc = err.OperationalError(
  616. CR.CR_CONN_HOST_ERROR,
  617. f"Can't connect to MySQL server on {self.host!r} ({e})",
  618. )
  619. # Keep original exception and traceback to investigate error.
  620. exc.original_exception = e
  621. exc.traceback = traceback.format_exc()
  622. if DEBUG:
  623. print(exc.traceback)
  624. raise exc
  625. # If e is neither DatabaseError or IOError, It's a bug.
  626. # But raising AssertionError hides original error.
  627. # So just reraise it.
  628. raise
  629. def write_packet(self, payload):
  630. """Writes an entire "mysql packet" in its entirety to the network
  631. adding its length and sequence number.
  632. """
  633. # Internal note: when you build packet manually and calls _write_bytes()
  634. # directly, you should set self._next_seq_id properly.
  635. data = _pack_int24(len(payload)) + bytes([self._next_seq_id]) + payload
  636. if DEBUG:
  637. dump_packet(data)
  638. self._write_bytes(data)
  639. self._next_seq_id = (self._next_seq_id + 1) % 256
  640. def _read_packet(self, packet_type=MysqlPacket):
  641. """Read an entire "mysql packet" in its entirety from the network
  642. and return a MysqlPacket type that represents the results.
  643. :raise OperationalError: If the connection to the MySQL server is lost.
  644. :raise InternalError: If the packet sequence number is wrong.
  645. """
  646. buff = bytearray()
  647. while True:
  648. packet_header = self._read_bytes(4)
  649. # if DEBUG: dump_packet(packet_header)
  650. btrl, btrh, packet_number = struct.unpack("<HBB", packet_header)
  651. bytes_to_read = btrl + (btrh << 16)
  652. if packet_number != self._next_seq_id:
  653. self._force_close()
  654. if packet_number == 0:
  655. # MariaDB sends error packet with seqno==0 when shutdown
  656. raise err.OperationalError(
  657. CR.CR_SERVER_LOST,
  658. "Lost connection to MySQL server during query",
  659. )
  660. raise err.InternalError(
  661. "Packet sequence number wrong - got %d expected %d"
  662. % (packet_number, self._next_seq_id)
  663. )
  664. self._next_seq_id = (self._next_seq_id + 1) % 256
  665. recv_data = self._read_bytes(bytes_to_read)
  666. if DEBUG:
  667. dump_packet(recv_data)
  668. buff += recv_data
  669. # https://dev.mysql.com/doc/internals/en/sending-more-than-16mbyte.html
  670. if bytes_to_read < MAX_PACKET_LEN:
  671. break
  672. packet = packet_type(bytes(buff), self.encoding)
  673. if packet.is_error_packet():
  674. if self._result is not None and self._result.unbuffered_active is True:
  675. self._result.unbuffered_active = False
  676. packet.raise_for_error()
  677. return packet
  678. def _read_bytes(self, num_bytes):
  679. self._sock.settimeout(self._read_timeout)
  680. while True:
  681. try:
  682. data = self._rfile.read(num_bytes)
  683. break
  684. except OSError as e:
  685. if e.errno == errno.EINTR:
  686. continue
  687. self._force_close()
  688. raise err.OperationalError(
  689. CR.CR_SERVER_LOST,
  690. f"Lost connection to MySQL server during query ({e})",
  691. )
  692. except BaseException:
  693. # Don't convert unknown exception to MySQLError.
  694. self._force_close()
  695. raise
  696. if len(data) < num_bytes:
  697. self._force_close()
  698. raise err.OperationalError(
  699. CR.CR_SERVER_LOST, "Lost connection to MySQL server during query"
  700. )
  701. return data
  702. def _write_bytes(self, data):
  703. self._sock.settimeout(self._write_timeout)
  704. try:
  705. self._sock.sendall(data)
  706. except OSError as e:
  707. self._force_close()
  708. raise err.OperationalError(
  709. CR.CR_SERVER_GONE_ERROR, f"MySQL server has gone away ({e!r})"
  710. )
  711. def _read_query_result(self, unbuffered=False):
  712. self._result = None
  713. if unbuffered:
  714. try:
  715. result = MySQLResult(self)
  716. result.init_unbuffered_query()
  717. except:
  718. result.unbuffered_active = False
  719. result.connection = None
  720. raise
  721. else:
  722. result = MySQLResult(self)
  723. result.read()
  724. self._result = result
  725. if result.server_status is not None:
  726. self.server_status = result.server_status
  727. return result.affected_rows
  728. def insert_id(self):
  729. if self._result:
  730. return self._result.insert_id
  731. else:
  732. return 0
  733. def _execute_command(self, command, sql):
  734. """
  735. :raise InterfaceError: If the connection is closed.
  736. :raise ValueError: If no username was specified.
  737. """
  738. if not self._sock:
  739. raise err.InterfaceError(0, "")
  740. # If the last query was unbuffered, make sure it finishes before
  741. # sending new commands
  742. if self._result is not None:
  743. if self._result.unbuffered_active:
  744. warnings.warn("Previous unbuffered result was left incomplete")
  745. self._result._finish_unbuffered_query()
  746. while self._result.has_next:
  747. self.next_result()
  748. self._result = None
  749. if isinstance(sql, str):
  750. sql = sql.encode(self.encoding)
  751. packet_size = min(MAX_PACKET_LEN, len(sql) + 1) # +1 is for command
  752. # tiny optimization: build first packet manually instead of
  753. # calling self..write_packet()
  754. prelude = struct.pack("<iB", packet_size, command)
  755. packet = prelude + sql[: packet_size - 1]
  756. self._write_bytes(packet)
  757. if DEBUG:
  758. dump_packet(packet)
  759. self._next_seq_id = 1
  760. if packet_size < MAX_PACKET_LEN:
  761. return
  762. sql = sql[packet_size - 1 :]
  763. while True:
  764. packet_size = min(MAX_PACKET_LEN, len(sql))
  765. self.write_packet(sql[:packet_size])
  766. sql = sql[packet_size:]
  767. if not sql and packet_size < MAX_PACKET_LEN:
  768. break
  769. def _request_authentication(self):
  770. # https://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol::HandshakeResponse
  771. if int(self.server_version.split(".", 1)[0]) >= 5:
  772. self.client_flag |= CLIENT.MULTI_RESULTS
  773. if self.user is None:
  774. raise ValueError("Did not specify a username")
  775. charset_id = charset_by_name(self.charset).id
  776. if isinstance(self.user, str):
  777. self.user = self.user.encode(self.encoding)
  778. data_init = struct.pack(
  779. "<iIB23s", self.client_flag, MAX_PACKET_LEN, charset_id, b""
  780. )
  781. if self.ssl and self.server_capabilities & CLIENT.SSL:
  782. self.write_packet(data_init)
  783. self._sock = self.ctx.wrap_socket(self._sock, server_hostname=self.host)
  784. self._rfile = self._sock.makefile("rb")
  785. self._secure = True
  786. data = data_init + self.user + b"\0"
  787. authresp = b""
  788. plugin_name = None
  789. if self._auth_plugin_name == "":
  790. plugin_name = b""
  791. authresp = _auth.scramble_native_password(self.password, self.salt)
  792. elif self._auth_plugin_name == "mysql_native_password":
  793. plugin_name = b"mysql_native_password"
  794. authresp = _auth.scramble_native_password(self.password, self.salt)
  795. elif self._auth_plugin_name == "caching_sha2_password":
  796. plugin_name = b"caching_sha2_password"
  797. if self.password:
  798. if DEBUG:
  799. print("caching_sha2: trying fast path")
  800. authresp = _auth.scramble_caching_sha2(self.password, self.salt)
  801. else:
  802. if DEBUG:
  803. print("caching_sha2: empty password")
  804. elif self._auth_plugin_name == "sha256_password":
  805. plugin_name = b"sha256_password"
  806. if self.ssl and self.server_capabilities & CLIENT.SSL:
  807. authresp = self.password + b"\0"
  808. elif self.password:
  809. authresp = b"\1" # request public key
  810. else:
  811. authresp = b"\0" # empty password
  812. if self.server_capabilities & CLIENT.PLUGIN_AUTH_LENENC_CLIENT_DATA:
  813. data += _lenenc_int(len(authresp)) + authresp
  814. elif self.server_capabilities & CLIENT.SECURE_CONNECTION:
  815. data += struct.pack("B", len(authresp)) + authresp
  816. else: # pragma: no cover - not testing against servers without secure auth (>=5.0)
  817. data += authresp + b"\0"
  818. if self.db and self.server_capabilities & CLIENT.CONNECT_WITH_DB:
  819. if isinstance(self.db, str):
  820. self.db = self.db.encode(self.encoding)
  821. data += self.db + b"\0"
  822. if self.server_capabilities & CLIENT.PLUGIN_AUTH:
  823. data += (plugin_name or b"") + b"\0"
  824. if self.server_capabilities & CLIENT.CONNECT_ATTRS:
  825. connect_attrs = b""
  826. for k, v in self._connect_attrs.items():
  827. k = k.encode("utf-8")
  828. connect_attrs += _lenenc_int(len(k)) + k
  829. v = v.encode("utf-8")
  830. connect_attrs += _lenenc_int(len(v)) + v
  831. data += _lenenc_int(len(connect_attrs)) + connect_attrs
  832. self.write_packet(data)
  833. auth_packet = self._read_packet()
  834. # if authentication method isn't accepted the first byte
  835. # will have the octet 254
  836. if auth_packet.is_auth_switch_request():
  837. if DEBUG:
  838. print("received auth switch")
  839. # https://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol::AuthSwitchRequest
  840. auth_packet.read_uint8() # 0xfe packet identifier
  841. plugin_name = auth_packet.read_string()
  842. if (
  843. self.server_capabilities & CLIENT.PLUGIN_AUTH
  844. and plugin_name is not None
  845. ):
  846. auth_packet = self._process_auth(plugin_name, auth_packet)
  847. else:
  848. raise err.OperationalError("received unknown auth switch request")
  849. elif auth_packet.is_extra_auth_data():
  850. if DEBUG:
  851. print("received extra data")
  852. # https://dev.mysql.com/doc/internals/en/successful-authentication.html
  853. if self._auth_plugin_name == "caching_sha2_password":
  854. auth_packet = _auth.caching_sha2_password_auth(self, auth_packet)
  855. elif self._auth_plugin_name == "sha256_password":
  856. auth_packet = _auth.sha256_password_auth(self, auth_packet)
  857. else:
  858. raise err.OperationalError(
  859. "Received extra packet for auth method %r", self._auth_plugin_name
  860. )
  861. if DEBUG:
  862. print("Succeed to auth")
  863. def _process_auth(self, plugin_name, auth_packet):
  864. handler = self._get_auth_plugin_handler(plugin_name)
  865. if handler:
  866. try:
  867. return handler.authenticate(auth_packet)
  868. except AttributeError:
  869. if plugin_name != b"dialog":
  870. raise err.OperationalError(
  871. CR.CR_AUTH_PLUGIN_CANNOT_LOAD,
  872. f"Authentication plugin '{plugin_name}'"
  873. f" not loaded: - {type(handler)!r} missing authenticate method",
  874. )
  875. if plugin_name == b"caching_sha2_password":
  876. return _auth.caching_sha2_password_auth(self, auth_packet)
  877. elif plugin_name == b"sha256_password":
  878. return _auth.sha256_password_auth(self, auth_packet)
  879. elif plugin_name == b"mysql_native_password":
  880. data = _auth.scramble_native_password(self.password, auth_packet.read_all())
  881. elif plugin_name == b"client_ed25519":
  882. data = _auth.ed25519_password(self.password, auth_packet.read_all())
  883. elif plugin_name == b"mysql_old_password":
  884. data = (
  885. _auth.scramble_old_password(self.password, auth_packet.read_all())
  886. + b"\0"
  887. )
  888. elif plugin_name == b"mysql_clear_password":
  889. # https://dev.mysql.com/doc/internals/en/clear-text-authentication.html
  890. data = self.password + b"\0"
  891. elif plugin_name == b"dialog":
  892. pkt = auth_packet
  893. while True:
  894. flag = pkt.read_uint8()
  895. echo = (flag & 0x06) == 0x02
  896. last = (flag & 0x01) == 0x01
  897. prompt = pkt.read_all()
  898. if prompt == b"Password: ":
  899. self.write_packet(self.password + b"\0")
  900. elif handler:
  901. resp = "no response - TypeError within plugin.prompt method"
  902. try:
  903. resp = handler.prompt(echo, prompt)
  904. self.write_packet(resp + b"\0")
  905. except AttributeError:
  906. raise err.OperationalError(
  907. CR.CR_AUTH_PLUGIN_CANNOT_LOAD,
  908. f"Authentication plugin '{plugin_name}'"
  909. f" not loaded: - {handler!r} missing prompt method",
  910. )
  911. except TypeError:
  912. raise err.OperationalError(
  913. CR.CR_AUTH_PLUGIN_ERR,
  914. f"Authentication plugin '{plugin_name}'"
  915. f" {handler!r} didn't respond with string. Returned '{resp!r}' to prompt {prompt!r}",
  916. )
  917. else:
  918. raise err.OperationalError(
  919. CR.CR_AUTH_PLUGIN_CANNOT_LOAD,
  920. f"Authentication plugin '{plugin_name}' not configured",
  921. )
  922. pkt = self._read_packet()
  923. pkt.check_error()
  924. if pkt.is_ok_packet() or last:
  925. break
  926. return pkt
  927. else:
  928. raise err.OperationalError(
  929. CR.CR_AUTH_PLUGIN_CANNOT_LOAD,
  930. "Authentication plugin '%s' not configured" % plugin_name,
  931. )
  932. self.write_packet(data)
  933. pkt = self._read_packet()
  934. pkt.check_error()
  935. return pkt
  936. def _get_auth_plugin_handler(self, plugin_name):
  937. plugin_class = self._auth_plugin_map.get(plugin_name)
  938. if not plugin_class and isinstance(plugin_name, bytes):
  939. plugin_class = self._auth_plugin_map.get(plugin_name.decode("ascii"))
  940. if plugin_class:
  941. try:
  942. handler = plugin_class(self)
  943. except TypeError:
  944. raise err.OperationalError(
  945. CR.CR_AUTH_PLUGIN_CANNOT_LOAD,
  946. f"Authentication plugin '{plugin_name}'"
  947. f" not loaded: - {plugin_class!r} cannot be constructed with connection object",
  948. )
  949. else:
  950. handler = None
  951. return handler
  952. # _mysql support
  953. def thread_id(self):
  954. return self.server_thread_id[0]
  955. def character_set_name(self):
  956. return self.charset
  957. def get_host_info(self):
  958. return self.host_info
  959. def get_proto_info(self):
  960. return self.protocol_version
  961. def _get_server_information(self):
  962. i = 0
  963. packet = self._read_packet()
  964. data = packet.get_all_data()
  965. self.protocol_version = data[i]
  966. i += 1
  967. server_end = data.find(b"\0", i)
  968. self.server_version = data[i:server_end].decode("latin1")
  969. i = server_end + 1
  970. self.server_thread_id = struct.unpack("<I", data[i : i + 4])
  971. i += 4
  972. self.salt = data[i : i + 8]
  973. i += 9 # 8 + 1(filler)
  974. self.server_capabilities = struct.unpack("<H", data[i : i + 2])[0]
  975. i += 2
  976. if len(data) >= i + 6:
  977. lang, stat, cap_h, salt_len = struct.unpack("<BHHB", data[i : i + 6])
  978. i += 6
  979. # TODO: deprecate server_language and server_charset.
  980. # mysqlclient-python doesn't provide it.
  981. self.server_language = lang
  982. try:
  983. self.server_charset = charset_by_id(lang).name
  984. except KeyError:
  985. # unknown collation
  986. self.server_charset = None
  987. self.server_status = stat
  988. if DEBUG:
  989. print("server_status: %x" % stat)
  990. self.server_capabilities |= cap_h << 16
  991. if DEBUG:
  992. print("salt_len:", salt_len)
  993. salt_len = max(12, salt_len - 9)
  994. # reserved
  995. i += 10
  996. if len(data) >= i + salt_len:
  997. # salt_len includes auth_plugin_data_part_1 and filler
  998. self.salt += data[i : i + salt_len]
  999. i += salt_len
  1000. i += 1
  1001. # AUTH PLUGIN NAME may appear here.
  1002. if self.server_capabilities & CLIENT.PLUGIN_AUTH and len(data) >= i:
  1003. # Due to Bug#59453 the auth-plugin-name is missing the terminating
  1004. # NUL-char in versions prior to 5.5.10 and 5.6.2.
  1005. # ref: https://dev.mysql.com/doc/internals/en/connection-phase-packets.html#packet-Protocol::Handshake
  1006. # didn't use version checks as mariadb is corrected and reports
  1007. # earlier than those two.
  1008. server_end = data.find(b"\0", i)
  1009. if server_end < 0: # pragma: no cover - very specific upstream bug
  1010. # not found \0 and last field so take it all
  1011. self._auth_plugin_name = data[i:].decode("utf-8")
  1012. else:
  1013. self._auth_plugin_name = data[i:server_end].decode("utf-8")
  1014. def get_server_info(self):
  1015. return self.server_version
  1016. Warning = err.Warning
  1017. Error = err.Error
  1018. InterfaceError = err.InterfaceError
  1019. DatabaseError = err.DatabaseError
  1020. DataError = err.DataError
  1021. OperationalError = err.OperationalError
  1022. IntegrityError = err.IntegrityError
  1023. InternalError = err.InternalError
  1024. ProgrammingError = err.ProgrammingError
  1025. NotSupportedError = err.NotSupportedError
  1026. class MySQLResult:
  1027. def __init__(self, connection):
  1028. """
  1029. :type connection: Connection
  1030. """
  1031. self.connection = connection
  1032. self.affected_rows = None
  1033. self.insert_id = None
  1034. self.server_status = None
  1035. self.warning_count = 0
  1036. self.message = None
  1037. self.field_count = 0
  1038. self.description = None
  1039. self.rows = None
  1040. self.has_next = None
  1041. self.unbuffered_active = False
  1042. def __del__(self):
  1043. if self.unbuffered_active:
  1044. self._finish_unbuffered_query()
  1045. def read(self):
  1046. try:
  1047. first_packet = self.connection._read_packet()
  1048. if first_packet.is_ok_packet():
  1049. self._read_ok_packet(first_packet)
  1050. elif first_packet.is_load_local_packet():
  1051. self._read_load_local_packet(first_packet)
  1052. else:
  1053. self._read_result_packet(first_packet)
  1054. finally:
  1055. self.connection = None
  1056. def init_unbuffered_query(self):
  1057. """
  1058. :raise OperationalError: If the connection to the MySQL server is lost.
  1059. :raise InternalError:
  1060. """
  1061. self.unbuffered_active = True
  1062. first_packet = self.connection._read_packet()
  1063. if first_packet.is_ok_packet():
  1064. self._read_ok_packet(first_packet)
  1065. self.unbuffered_active = False
  1066. self.connection = None
  1067. elif first_packet.is_load_local_packet():
  1068. self._read_load_local_packet(first_packet)
  1069. self.unbuffered_active = False
  1070. self.connection = None
  1071. else:
  1072. self.field_count = first_packet.read_length_encoded_integer()
  1073. self._get_descriptions()
  1074. # Apparently, MySQLdb picks this number because it's the maximum
  1075. # value of a 64bit unsigned integer. Since we're emulating MySQLdb,
  1076. # we set it to this instead of None, which would be preferred.
  1077. self.affected_rows = 18446744073709551615
  1078. def _read_ok_packet(self, first_packet):
  1079. ok_packet = OKPacketWrapper(first_packet)
  1080. self.affected_rows = ok_packet.affected_rows
  1081. self.insert_id = ok_packet.insert_id
  1082. self.server_status = ok_packet.server_status
  1083. self.warning_count = ok_packet.warning_count
  1084. self.message = ok_packet.message
  1085. self.has_next = ok_packet.has_next
  1086. def _read_load_local_packet(self, first_packet):
  1087. if not self.connection._local_infile:
  1088. raise RuntimeError(
  1089. "**WARN**: Received LOAD_LOCAL packet but local_infile option is false."
  1090. )
  1091. load_packet = LoadLocalPacketWrapper(first_packet)
  1092. sender = LoadLocalFile(load_packet.filename, self.connection)
  1093. try:
  1094. sender.send_data()
  1095. except:
  1096. self.connection._read_packet() # skip ok packet
  1097. raise
  1098. ok_packet = self.connection._read_packet()
  1099. if (
  1100. not ok_packet.is_ok_packet()
  1101. ): # pragma: no cover - upstream induced protocol error
  1102. raise err.OperationalError(
  1103. CR.CR_COMMANDS_OUT_OF_SYNC,
  1104. "Commands Out of Sync",
  1105. )
  1106. self._read_ok_packet(ok_packet)
  1107. def _check_packet_is_eof(self, packet):
  1108. if not packet.is_eof_packet():
  1109. return False
  1110. # TODO: Support CLIENT.DEPRECATE_EOF
  1111. # 1) Add DEPRECATE_EOF to CAPABILITIES
  1112. # 2) Mask CAPABILITIES with server_capabilities
  1113. # 3) if server_capabilities & CLIENT.DEPRECATE_EOF:
  1114. # use OKPacketWrapper instead of EOFPacketWrapper
  1115. wp = EOFPacketWrapper(packet)
  1116. self.warning_count = wp.warning_count
  1117. self.has_next = wp.has_next
  1118. return True
  1119. def _read_result_packet(self, first_packet):
  1120. self.field_count = first_packet.read_length_encoded_integer()
  1121. self._get_descriptions()
  1122. self._read_rowdata_packet()
  1123. def _read_rowdata_packet_unbuffered(self):
  1124. # Check if in an active query
  1125. if not self.unbuffered_active:
  1126. return
  1127. # EOF
  1128. packet = self.connection._read_packet()
  1129. if self._check_packet_is_eof(packet):
  1130. self.unbuffered_active = False
  1131. self.connection = None
  1132. self.rows = None
  1133. return
  1134. row = self._read_row_from_packet(packet)
  1135. self.affected_rows = 1
  1136. self.rows = (row,) # rows should tuple of row for MySQL-python compatibility.
  1137. return row
  1138. def _finish_unbuffered_query(self):
  1139. # After much reading on the MySQL protocol, it appears that there is,
  1140. # in fact, no way to stop MySQL from sending all the data after
  1141. # executing a query, so we just spin, and wait for an EOF packet.
  1142. while self.unbuffered_active:
  1143. try:
  1144. packet = self.connection._read_packet()
  1145. except err.OperationalError as e:
  1146. if e.args[0] in (
  1147. ER.QUERY_TIMEOUT,
  1148. ER.STATEMENT_TIMEOUT,
  1149. ):
  1150. # if the query timed out we can simply ignore this error
  1151. self.unbuffered_active = False
  1152. self.connection = None
  1153. return
  1154. raise
  1155. if self._check_packet_is_eof(packet):
  1156. self.unbuffered_active = False
  1157. self.connection = None # release reference to kill cyclic reference.
  1158. def _read_rowdata_packet(self):
  1159. """Read a rowdata packet for each data row in the result set."""
  1160. rows = []
  1161. while True:
  1162. packet = self.connection._read_packet()
  1163. if self._check_packet_is_eof(packet):
  1164. self.connection = None # release reference to kill cyclic reference.
  1165. break
  1166. rows.append(self._read_row_from_packet(packet))
  1167. self.affected_rows = len(rows)
  1168. self.rows = tuple(rows)
  1169. def _read_row_from_packet(self, packet):
  1170. row = []
  1171. for encoding, converter in self.converters:
  1172. try:
  1173. data = packet.read_length_coded_string()
  1174. except IndexError:
  1175. # No more columns in this row
  1176. # See https://github.com/PyMySQL/PyMySQL/pull/434
  1177. break
  1178. if data is not None:
  1179. if encoding is not None:
  1180. data = data.decode(encoding)
  1181. if DEBUG:
  1182. print("DEBUG: DATA = ", data)
  1183. if converter is not None:
  1184. data = converter(data)
  1185. row.append(data)
  1186. return tuple(row)
  1187. def _get_descriptions(self):
  1188. """Read a column descriptor packet for each column in the result."""
  1189. self.fields = []
  1190. self.converters = []
  1191. use_unicode = self.connection.use_unicode
  1192. conn_encoding = self.connection.encoding
  1193. description = []
  1194. for i in range(self.field_count):
  1195. field = self.connection._read_packet(FieldDescriptorPacket)
  1196. self.fields.append(field)
  1197. description.append(field.description())
  1198. field_type = field.type_code
  1199. if use_unicode:
  1200. if field_type == FIELD_TYPE.JSON:
  1201. # When SELECT from JSON column: charset = binary
  1202. # When SELECT CAST(... AS JSON): charset = connection encoding
  1203. # This behavior is different from TEXT / BLOB.
  1204. # We should decode result by connection encoding regardless charsetnr.
  1205. # See https://github.com/PyMySQL/PyMySQL/issues/488
  1206. encoding = conn_encoding # SELECT CAST(... AS JSON)
  1207. elif field_type in TEXT_TYPES:
  1208. if field.charsetnr == 63: # binary
  1209. # TEXTs with charset=binary means BINARY types.
  1210. encoding = None
  1211. else:
  1212. encoding = conn_encoding
  1213. else:
  1214. # Integers, Dates and Times, and other basic data is encoded in ascii
  1215. encoding = "ascii"
  1216. else:
  1217. encoding = None
  1218. converter = self.connection.decoders.get(field_type)
  1219. if converter is converters.through:
  1220. converter = None
  1221. if DEBUG:
  1222. print(f"DEBUG: field={field}, converter={converter}")
  1223. self.converters.append((encoding, converter))
  1224. eof_packet = self.connection._read_packet()
  1225. assert eof_packet.is_eof_packet(), "Protocol error, expecting EOF"
  1226. self.description = tuple(description)
  1227. class LoadLocalFile:
  1228. def __init__(self, filename, connection):
  1229. self.filename = filename
  1230. self.connection = connection
  1231. def send_data(self):
  1232. """Send data packets from the local file to the server"""
  1233. if not self.connection._sock:
  1234. raise err.InterfaceError(0, "")
  1235. conn: Connection = self.connection
  1236. try:
  1237. with open(self.filename, "rb") as open_file:
  1238. packet_size = min(
  1239. conn.max_allowed_packet, 16 * 1024
  1240. ) # 16KB is efficient enough
  1241. while True:
  1242. chunk = open_file.read(packet_size)
  1243. if not chunk:
  1244. break
  1245. conn.write_packet(chunk)
  1246. except OSError:
  1247. raise err.OperationalError(
  1248. ER.FILE_NOT_FOUND,
  1249. f"Can't find file '{self.filename}'",
  1250. )
  1251. finally:
  1252. if not conn._closed:
  1253. # send the empty packet to signify we are done sending data
  1254. conn.write_packet(b"")