Minor updates

This commit is contained in:
codez0mb1e 2022-06-19 23:01:08 +00:00
parent 153fc5a230
commit b51d4a0162

View File

@ -1,16 +1,18 @@
# %% Import dependencies ---- # %% Import dependencies ----
from dataclasses import dataclass from dataclasses import dataclass
from typing import Dict, Any, Iterable from typing import Dict, Any
from pandas import DataFrame
from sqlalchemy import create_engine, inspect from sqlalchemy import create_engine, inspect
import pandas as pd
import urllib import urllib
# %% # %% Models
@dataclass(frozen=True) @dataclass(frozen=True)
class ConnectionSettings: class ConnectionSettings:
"""Connection Settings.""" """Connection Settings"""
server: str server: str
database: str database: str
username: str username: str
@ -19,10 +21,10 @@ class ConnectionSettings:
timeout: int = 30 timeout: int = 30
# %% Connection
class AzureDbConnection: class AzureDbConnection:
""" """Azure SQL database connection."""
Azure SQL database connection.
"""
def __init__(self, conn_settings: ConnectionSettings, echo: bool = False) -> None: def __init__(self, conn_settings: ConnectionSettings, echo: bool = False) -> None:
conn_params = urllib.parse.quote_plus( conn_params = urllib.parse.quote_plus(
'Driver=%s;' % conn_settings.driver + 'Driver=%s;' % conn_settings.driver +
@ -36,29 +38,29 @@ class AzureDbConnection:
) )
conn_string = f'mssql+pyodbc:///?odbc_connect={conn_params}' conn_string = f'mssql+pyodbc:///?odbc_connect={conn_params}'
self.db = create_engine(conn_string, echo=echo) self._db = create_engine(conn_string, echo=echo)
def connect(self) -> None: def connect(self) -> None:
"""Estimate connection.""" """Estimate connection"""
self.conn = self.db.connect() self._conn = self._db.connect()
def get_tables(self) -> Iterable[str]: def get_tables(self) -> list[str]:
"""Get list of tables.""" """Get list of tables"""
inspector = inspect(self.db) inspector = inspect(self._db)
return [t for t in inspector.get_table_names()] return [t for t in inspector.get_table_names()]
def insert(self, inserted_data: DataFrame, target_table: str, db_mapping: Dict[str, Any], chunksize: int = 10000) -> None: def insert(self, inserted_data: pd.DataFrame, target_table: str, db_mapping: Dict[str, Any], chunksize: int = 10000) -> None:
inserted_data.to_sql( inserted_data.to_sql(
con=self.db, con=self._db,
schema='dbo', schema='dbo',
name=target_table, name=target_table,
if_exists='append', # or replace if_exists='replace', # or append
index=False, index=False,
chunksize=chunksize, chunksize=chunksize,
dtype=db_mapping dtype=db_mapping
) )
def dispose(self) -> None: def dispose(self) -> None:
"""Dispose opened connections.""" """Dispose opened connections"""
self.conn.close() self._conn.close()
self.db.dispose() self._db.dispose()