def __init__(self,
db_name: str,
db_type: DB_Type = DB_Type.SQLITE,
log_level: str = 'WARNING') -> None:
self.INVALID_STATUS = Error('Unknown DB_Type value.')
self.db_name = db_name
self.db_type = db_type
self.log_level = log_level
self.status: Status
self.__connect__()
def create(cur: Cursor, stmt: str) -> Status:
"""Create table."""
status: Status
try:
cur.execute(stmt)
status = OK()
logger.info(f'Create statement executed: {stmt}')
except Exception as e:
logger.error(f'Create statement exception: {stmt}; {e}')
status = Error(str(e))
return status
def validate_insert(cur: Cursor, table: str, rows: Rows,
schema_cast: bool) -> Tuple[SqliteSchema, RowsPair]:
"""Validate insertion cols and optionally try casting to schema dtypes."""
q = f'SELECT sql FROM sqlite_master WHERE type="table" and name="{table}"'
ret, status = query(cur, q)
if not ret or status != OK():
msg = f'Schema validation query {q} failed'
logger.error(msg)
return [], ([], Error(msg))
schema = parse_schema(ret[0][0])
if len(schema) != len(rows[0]):
msg = f'Insertion validation error: {table} has '
msg += f'{len(schema)} cols vs input {len(rows[0])} cols'
logger.error(msg)
return [], ([], Error(msg))
return (schema, apply_schema(schema, rows) if schema_cast else
(rows, OK()))
def close(conn) -> Status:
"""Close DB connection object."""
status: Status
try:
conn.close()
status = OK()
logger.info('Closed DB conn.')
except Exception as e:
status = Error(str(e))
logger.error('Failed to close DB conn: {}'.format(str(e)))
return status
def compare_dims(df1: pd.DataFrame,
df2: pd.DataFrame,
cols: bool = True,
rows: bool = False) -> Status:
"""Compare DF dimensions for compatibility."""
status: Status
cols1, cols2 = len(df1.columns), len(df2.columns)
rows1, rows2 = len(df1), len(df2)
if not cols and not rows:
status = OK()
elif cols and not rows:
status = (OK() if cols1 == cols2 else
Error(f'Cols mismatch: {cols1} vs {cols2}'))
elif rows and not cols:
status = (OK() if rows1 == rows2 else
Error(f'Rows mismatch: {rows1} vs {rows2}'))
else:
status = (OK() if rows1 == rows2 and cols1 == cols2 else
Error('Matrix mismatch: ' +
f'{rows1} x {cols1} vs {rows2} x {cols2}'))
return status
def create(self, stmt: str) -> Status:
"""Create table."""
if not safe_statement(stmt):
msg = f'Safe statement check failed: {stmt}'
logger.error(msg)
return Error(msg)
if self.db_type is DB_Type.SQLITE:
status = db_sqlite.create(self.cur, stmt)
else:
status = self.INVALID_STATUS
logger.error(f'Create failed: {self.INVALID_STATUS.msg}')
return status
def apply_schema(schema: SqliteSchema, rows: Rows) -> RowsPair:
"""Attempt to cast rows to primitive types in schema."""
status: Status = OK()
try:
rows_ = []
for row in rows:
row_ = [cast(elem) for elem, (_, cast) in zip(row, schema)]
rows_.append(row_)
except Exception as e:
msg = f'Insertion validation error: exception while casting {row}: {e}'
status = Error(msg)
logger.error(msg)
return rows_, status
def query(cur: Cursor, q: str, hdr: bool = False) -> RowsPair:
"""Execute SQL query string."""
status: Status
try:
result = cur.execute(q)
if hdr:
cols = [d[0] for d in result.description]
rows = [cols] + [list(row) for row in result.fetchall()]
else:
rows = [list(row) for row in result.fetchall()]
status = OK()
logger.info(f'Query executed: {q}')
except Exception as e:
logger.error(f'Query exception: {q}; {e}')
rows, status = [], Error(str(e))
return rows, status
def gen_fks(fks: List[SchemaForeignKey]) -> Tuple[str, Status]:
"""Generate Foreign Keys substring."""
status: Status
s, status = '', OK()
for fk in fks:
fk_cols, fk_ref_cols = fk['cols'], fk['ref_cols']
if len(fk_cols) != len(fk_ref_cols):
s, status = '', Error(
f'Length mismatch {fk_cols} vs {fk_ref_cols}')
break
cols, ref_cols = ', '.join(fk_cols), ', '.join(fk_ref_cols)
ref_table = fk['ref_table']
s += f'FOREIGN KEY({cols}) REFERENCES {ref_table}({ref_cols}),\n'
return f'{s}', status
def gen_cols(cols: List[SchemaCol]) -> Tuple[str, Status]:
"""Generate columns substring."""
status: Status
if not cols: return '', OK()
s, status = '', OK()
for col in cols:
name, dtype, pk, uniq, not_null = col
if pk and uniq:
s, status = '', Error(f'Col {name} specified as both PK and Uniq')
break
dtype_ = dtype_to_str(dtype)
s += f'{name} {dtype_}'
s += ' PRIMARY KEY' if pk else ' UNIQUE' if uniq else ''
s += ' NOT NULL,\n' if not_null else ',\n'
return s, status
def query(self,
q: str,
hdr: bool = False,
df: bool = False) -> Tuple[QueryResult, Status]:
"""Run query."""
if not valid_query(q):
logger.error('Invalid query {}'.format(q))
return [], Error('Invalid query {}'.format(q))
if self.db_type is DB_Type.SQLITE:
ret, status = (db_sqlite.query(self.cur, q, hdr)
if not df else db_sqlite.query_df(self.cur, q))
else:
ret, status = [], self.INVALID_STATUS
logger.error('Query failed: {}'.format(self.INVALID_STATUS.msg))
return ret, status
def gen_create_stmt(td: TableDef) -> Tuple[str, Status]:
"""Translate TableDef into Create Table statement."""
status: Status
create, s1 = gen_create(td['if_not_exists'], td['name'])
cols, s2 = gen_cols(td['cols'])
fks, s3 = gen_fks(td['fks'])
pk_uniq, s4 = gen_pk_uniq(td['pk'], td['uniq'])
if all(s == OK() for s in (s1, s2, s3, s4)):
spec = strip_comma(f'{cols}{fks}{pk_uniq}')
stmt = f'{create}({spec});'
status = OK()
else:
stmt = ''
msg = ' | '.join(s.msg for s in (s1, s2, s3, s4) if s != OK())
status = Error(msg)
return stmt, status
def insert(conn: Conn,
cur: Cursor,
table: str,
rows: Rows,
schema_cast: bool = True) -> Status:
"""Attempt to execute SQL insertion into specified table."""
status: Status
schema, (rows_, v) = validate_insert(cur, table, rows, schema_cast)
if v != OK(): return v
try:
cols = ','.join(name for name, _ in schema)
vals = ','.join('?' * len(schema))
i = f'INSERT INTO {table}({cols}) VALUES ({vals})'
cur.executemany(i, rows_)
conn.commit()
status = OK()
logger.info('Insertion to {} executed: {}'.format(table, i))
except Exception as e:
status = Error(str(e))
logger.error('Insertion exception for {}: {}'.format(i, str(e)))
return status