49 lines
1.2 KiB
Python
49 lines
1.2 KiB
Python
# -*- coding: utf-8 -*-
|
|
import pandas as pd
|
|
import pytz
|
|
|
|
from sqlalchemy import create_engine
|
|
from sqlalchemy.exc import ProgrammingError
|
|
from sqlalchemy.types import DateTime
|
|
|
|
|
|
def get_engine(uri: str):
|
|
return create_engine(uri, pool_recycle=3600)
|
|
|
|
|
|
def get_connection(uri: str):
|
|
return get_engine(uri).connect()
|
|
|
|
|
|
def get_df(uri, table_name):
|
|
"""Get dataframe and the first time create DB."""
|
|
connection = get_connection(uri)
|
|
try:
|
|
return pd.read_sql_table(
|
|
table_name=table_name,
|
|
con=connection,
|
|
index_col='index',
|
|
columns=['Token', 'Text', 'Link', 'Datetime discover', 'Datetime announcement'],
|
|
)
|
|
except ValueError as e:
|
|
return None
|
|
finally:
|
|
connection.close()
|
|
|
|
|
|
def save_df(df, uri, table_name):
|
|
"""Save dataframe on DB."""
|
|
connection = get_connection(uri)
|
|
try:
|
|
df.to_sql(
|
|
name=table_name,
|
|
con=connection,
|
|
index=True,
|
|
index_label='index',
|
|
if_exists='replace',
|
|
dtype={"Datetime discover": DateTime(timezone=pytz.utc),
|
|
"Datetime announcement": DateTime(timezone=pytz.utc)}
|
|
)
|
|
finally:
|
|
connection.close()
|