def check_conn(self, conn):
"""Make sure the object is a valid SQL connection"""
raiseifnot(
hasattr(conn, "cursor") or is_sqlalchemy_conn(conn),
"Connection must have a cursor() method or be a SQLAlchemy connection: %s"
% conn,
)
def execute(self, conn, cursor, sql, params=None, **kwargs):
"""Executes the sql statement and returns an object that can fetch results
Parameters
----------
conn
A SQL database connection object
cursor
A SQL database cursor
sql : str
A sql query to execute
params : tuple, optional
A tuple of params to pass to the execute method of the conn or cursor
**kwargs
kwargs passed through to execute()
Returns
-------
cursor
cursor object that has executed but not fetched a query.
"""
params = params or ()
if is_sqlalchemy_conn(conn):
qr = conn.execute(sql, *params, **kwargs)
return qr
qr = cursor.execute(sql, params, **kwargs)
return cursor
def commit(self, obj):
"""Commit any currently active transactions"""
if hasattr(obj, "commit"):
obj.commit()
elif is_sqlalchemy_conn(obj):
# Hack: I don't want to have to pass around the transaction
# between nodes since that requirement is really specific to
# SQLAlchemy, and SQLAlchemy doesn't seem to provide a standard
# way of obtaining the current transaction, so this approach is
# lifted from the SQLAlchemy internals.
raiseifnot(
hasattr(obj, "_Connection__transaction"),
"Could not find transaction attribute on SQLAlchemy object: %s"
% obj,
)
if getattr(obj, "_Connection__transaction", None):
obj._Connection__transaction.commit()
else:
# SQLAlchemy connections autocommit by default, so we assume
# that happened.
pass
else:
raise AssertionError(
"Could not determine how to commit with object: %s" % obj)
def transaction(self, conn, cursor=None):
"""Start a transaction. If conn is a SQLAlchemy conn return a
reference to the transaction object, otherwise just return the conn
which should have commit/rollback methods."""
dbg("starting transaction: %s" % conn)
if is_sqlalchemy_conn(conn):
return conn.begin()
# For SQLite and DBAPI connections we explicitly call begin.
# https://docs.python.org/3/library/sqlite3.html#sqlite3-controlling-transactions
if not cursor:
cursor = self.get_sql_executor(conn)
cursor.execute("BEGIN")
return conn
def get_bulk_statement(self, conn, stmt_type, table, rows, odku=False):
"""Get a bulk execution SQL statement
Parameters
----------
conn
A SQL database connection object
stmt_type : str
Type of SQL statement to use (REPLACE, INSERT, etc.)
table : str
name of a SQL table
rows
An iterable of dict rows. The first row is used to determine
column names.
odku : bool or list, optional
If true, add ON DUPLICATE KEY UPDATE clause for all columns. If a
list then only add it for the specified columns. **Note:** Backend
support for this varies.
Returns
-------
A SQL bulk load query of the given stmt_type
"""
if is_sqlalchemy_conn(conn):
return get_bulk_statement(stmt_type,
table,
rows[0].keys(),
dicts=False,
odku=odku)
if isinstance(conn, sqlite3.Connection):
raiseifnot(isinstance(rows[0], sqlite3.Row),
"Only sqlite3.Row rows are supported")
return get_bulk_statement(
stmt_type,
table,
rows[0].keys(),
dicts=False,
value_string="?",
odku=odku,
)
raiseif(
isinstance(rows[0], tuple),
"Dict rows expected, got tuple. Please use a dict cursor.",
)
return get_bulk_statement(stmt_type, table, rows[0].keys(), odku=odku)
def rollback(self, obj):
"""Rollback any currently active transactions"""
dbg("rolling back transaction: %s" % obj)
if hasattr(obj, "rollback"):
obj.rollback()
elif is_sqlalchemy_conn(obj):
# See note above about this hack
raiseifnot(
hasattr(obj, "_Connection__transaction"),
"Could not find transaction attribute on SQLAlchemy object: %s"
% obj,
)
if getattr(obj, "_Connection__transaction", None):
obj._Connection__transaction.rollback()
else:
raise AssertionError(
"Trying to rollback a transaction but the SQLAlchemy "
"conn was not in a transaction. It may have "
"autocommitted.")
else:
raise AssertionError(
"Could not determine how to rollback with object: %s" % obj)
def executemany(self, conn, cursor, sql, rows):
"""Bulk executes the sql statement and returns an object that can fetch results
Parameters
----------
conn
A SQL database connection object
cursor
A SQL database cursor
sql : str
A sql query to execute
rows
Rows of data to bulk execute
Returns
-------
cursor
cursor object that has executed but not fetched a query.
"""
if is_sqlalchemy_conn(conn):
qr = conn.execute(sql, rows)
return qr
qr = cursor.executemany(sql, rows)
return cursor
def get_sql_executor(self, conn, cursor_type=None):
"""Get the object that can execute queries"""
if is_sqlalchemy_conn(conn):
return conn
return conn.cursor(cursor_type) if cursor_type else conn.cursor()