mirror of
https://github.com/dolthub/dolt.git
synced 2026-01-16 03:00:24 -06:00
66 lines
2.1 KiB
Python
66 lines
2.1 KiB
Python
import csv
|
|
import sys
|
|
import time
|
|
import mysql.connector
|
|
|
|
from io import StringIO
|
|
from multiprocessing import Process
|
|
|
|
def _connect(user, password, host, port, database):
|
|
return mysql.connector.connect(user=user, password=password, host=host, port=port, database=database, allow_local_infile=True)
|
|
|
|
def _print_err_and_exit(e):
|
|
print(e, file=sys.stderr)
|
|
sys.exit(1)
|
|
|
|
class DoltConnection(object):
|
|
def __init__(self, user='root', password=None, host='127.0.0.1', port=3306, database='dolt', auto_commit=False):
|
|
self.user = user
|
|
self.password = password
|
|
self.host = host
|
|
self.port = port
|
|
self.database = database
|
|
self.auto_commit = auto_commit
|
|
self.cnx = None
|
|
|
|
def connect(self):
|
|
try:
|
|
self.cnx = _connect(self.user, self.password, self.host, self.port, self.database)
|
|
self.cnx.autocommit=self.auto_commit
|
|
except BaseException as e:
|
|
_print_err_and_exit(e)
|
|
|
|
def close(self):
|
|
self.cnx.close()
|
|
|
|
class InfiniteRetryConnection(DoltConnection):
|
|
def connect(self):
|
|
while True:
|
|
try:
|
|
self.cnx = _connect(user=self.user, password=self.password, host=self.host, port=self.port, database=self.database)
|
|
|
|
try:
|
|
self.cnx.close()
|
|
except BaseException:
|
|
pass
|
|
|
|
return
|
|
|
|
except BaseException:
|
|
pass
|
|
|
|
def wait_for_connection(user='root', password=None, host='127.0.0.1', port=3306, database='dolt', timeout_ms=5000):
|
|
timeoutf = timeout_ms / 1000.0
|
|
exit_zero_on_connect = InfiniteRetryConnection(user=user, password=password, host=host, port=port, database=database)
|
|
|
|
cnx_proc = Process(target=exit_zero_on_connect.connect)
|
|
cnx_proc.start()
|
|
cnx_proc.join(timeout=timeoutf)
|
|
|
|
if cnx_proc.exitcode is None:
|
|
cnx_proc.terminate()
|
|
_print_err_and_exit(Exception("Failed to establish connection in time."))
|
|
elif cnx_proc.exitcode != 0:
|
|
_print_err_and_exit(Exception("Connection process exited with exit code %d." % cnx_proc.exitcode))
|
|
|