class BaseClient(abc.ABC):
"""Create BaseClient for DBs."""
default_dialect: Optional[str] = None # To be defined by subclasses.
default_driver: Optional[str] = None
def __init__(
self,
username=None,
database=None,
host=None,
dialect=None,
port=None,
driver=None,
password=None,
connect_args=None,
):
if dialect is None:
dialect = self.default_dialect
if driver is None and self.default_driver is not None:
driver = self.default_driver
self.connect_args = connect_args
self.conn_url = URL(
drivername=format_drivername(dialect, driver),
username=username,
password=password,
host=host,
port=port,
database=database,
)
# This variable is used to list class attributes to be forwarded to
# `self.conn_url` for backward compatibility.
# This should be removed in the future as well as the extra setter/getter
# logic.
_conn_url_delegated = (
"username",
"database",
"host",
"dialect",
"port",
"driver",
"password",
)
def __getattr__(self, attr):
if attr in self._conn_url_delegated:
if attr == 'dialect':
return self.conn_url.get_backend_name()
elif attr == 'driver':
return self.conn_url.get_driver_name()
else:
_cls = URL
target = self.conn_url
else:
_cls = object
target = self
return _cls.__getattribute__(target, attr)
def __setattr__(self, attr, value):
if attr in self._conn_url_delegated:
_cls = URL
target = self.conn_url
if attr in ('dialect', 'driver'):
if attr == 'dialect':
value = format_drivername(value, self.driver)
elif attr == 'driver':
value = format_drivername(self.dialect, value)
attr = 'drivername'
else:
_cls = object
target = self
return _cls.__setattr__(target, attr, value)
@property
def conn_str(self):
return str(self.conn_url)
@abc.abstractmethod
def _connect(self):
raise NotImplementedError
@staticmethod
def _cursor_columns(cursor):
# TODO: This can be moved out of the class
if hasattr(cursor, 'keys'):
return cursor.keys()
else:
return [c[0] for c in cursor.description]
@parse_sql_statement_decorator
def execute(self, sql, params=None, connection=None): # pylint: disable=W0613
"""Execute sql statement."""
if params is not None:
sql = sql.format(**params)
if connection is not None:
# It's not our job to close the passed connection
return connection.execute(sql)
with self._connect() as connection:
# Create and destroy a connection, as to avoid dangling connections
return connection.execute(sql)
def to_frame(self, sql, params=None, connection=None):
"""Execute SQL statement and return as Pandas dataframe."""
with closing(self.execute(sql, params, connection)) as cursor:
if not cursor:
return
data = cursor.fetchall()
if data:
df = pd.DataFrame(data, columns=self._cursor_columns(cursor))
else:
df = pd.DataFrame()
return df
def insert_from_frame(self,
df,
table,
if_exists='append',
index=False,
connection=None,
**kwargs):
"""Insert from a Pandas dataframe."""
# TODO: Validate types here?
column_names = df.columns.tolist()
chunksize = kwargs.get("chunksize", 10_000)
if if_exists == "fail" or if_exists == "replace":
raise NotImplementedError
if index:
raise NotImplementedError
insert_stmt = """
INSERT INTO {table} ({columns})
VALUES {values_chunk}
"""
def _to_sql_tuple(d):
return f"({', '.join(map(str, d.values()))})"
def df_chunksize_iterator(df, chunksize=10_000):
for start in range(0, len(df), chunksize):
yield df[start:start + chunksize]
for chunk in df_chunksize_iterator(df, chunksize):
values = map(_to_sql_tuple, chunk.to_dict(orient="records"))
chunk_insert_stmt = insert_stmt.format(
table=table,
columns=column_names,
values_chunk=",\n".join(values))
self.execute(chunk_insert_stmt, connection=connection)
def is_db_uri_mysql(db_uri: URL) -> bool:
backend_name = db_uri.get_backend_name()
return backend_name == MYSQL_DRIVER_NAME_PREFIX
