mb_sql 2.0.1__py3-none-any.whl

mb/sql/conn.py

@@ -0,0 +1,85 @@
+import sqlalchemy as sa
+from sqlalchemy.orm import sessionmaker,declarative_base
+from mb.utils.logging import logg
+
+__all__ = ['get_engine', 'get_session','get_base','get_metadata']
+
+
+def get_engine(name:str ='postgresql+psycopg2' , db: str = 'postgres',
+                user: str = 'postgres', password: str  = 'postgres', 
+                host: str = 'localhost', port=5432 , 
+                logger=None, echo=False):
+    """Get a SQLAlchemy engine object.
+
+    Args:
+        name (str): Name of the database.
+        db (str): Database name.
+        user (str): Username for the database.
+        password (str): Password for the database.
+        host (str): Hostname for the database.
+        port (str): Port for the database. Default: 5432
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+        echo (bool): Echo SQL statements to stdout. Default: False
+    Returns:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+
+    """
+    try:
+        engine = sa.create_engine(f'{name}://{user}:{password}@{host}:{port}/{db}', echo=echo)
+        logg.info(f'Engine created for {name} database.', logger=logger)
+        return engine
+    except Exception as e:
+        logg.error(f'Error creating engine for {name} database.', logger=logger)
+        raise e    
+
+def get_session(engine, logger=None):
+    """Get a SQLAlchemy session object.
+
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        session (sqlalchemy.orm.session.Session): Session object.
+
+    """
+    try:
+        session = sessionmaker(bind=engine)()
+        logg.info(f'Session created for {engine.url.database} database.', logger=logger)
+        return session
+    except Exception as e:
+        logg.error(f'Error creating session for {engine.url.database} database.', logger=logger)
+        raise e
+    
+def get_base(logger=None):
+    """Get a SQLAlchemy declarative base object.
+    
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        base (sqlalchemy.ext.declarative.api.DeclarativeMeta): Declarative base object.
+    """
+    try:
+        Base = declarative_base()
+        logg.info('Base created for database.', logger=logger)
+        return Base
+    except Exception:
+        logg.error('Error creating base for database.', logger=logger)
+    
+def get_metadata(base,conn, logger=None):
+    """Get a SQLAlchemy metadata object.
+    
+    Args:
+        base (sqlalchemy.ext.declarative.api.DeclarativeMeta): Declarative base object.
+        conn (sqlalchemy.engine.base.Connection): Connection object.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        metadata (sqlalchemy.sql.schema.MetaData): Metadata object.
+    """
+    try:
+        metadata = base.metadata.create_all(conn)
+        logg.info('Metadata created for database.', logger=logger)
+        return metadata
+    except Exception:
+        logg.error('Error creating metadata for database.', logger=logger)
+            

mb/sql/slack.py

@@ -0,0 +1,24 @@
+
+import requests
+import json
+from mb.utils.logging import logg
+
+__all__ = ['slack_msg']
+
+def slack_msg(webhook,msg,logger=None):
+    """
+    Send a message to a slack channel
+    Args:
+        webhook (str): Slack webhook URL
+        msg (str): Message to send
+        logger (logging.Logger): Logger to use
+    Returns:
+        None
+    """
+    
+    response = requests.post(
+        url=webhook,
+        data=json.dumps(msg),
+        headers={'Content-Type': 'application/json'})
+
+    logg.info('Slack response: %s', response.text, logger=logger)

mb/sql/snapshot.py

@@ -0,0 +1,197 @@
+'''
+fix this file. 
+it hsould be able to take any snapshot of the current data and save it to s3 and vice versa.
+ID should be given and stored in s3/DB as well along with time-data.
+'''
+from mb.utils import s3
+from typing import Optional
+import pandas as pd
+from mb.utils.logging import logg
+from mb.pandas.dfload import load_any_df
+
+__api__ = [
+    "get_s3url",
+    "list_names",
+    "list_ids",
+    "latest_id",
+    "make_new_id",
+    "save",
+    "load",
+]
+
+
+def get_s3url(df_name,id,location='s3://snapshots/dataframes/',file_type='csv'):
+    """Gets the s3url of a dataframe snapshot.
+
+    Parameters
+    ----------
+    df_name : str
+        dataframe name
+    id : int
+        snapshot id
+    location : str
+        s3cmd url to the snapshots folder
+    file_type : str
+        file type of the snapshot
+
+    Returns
+    -------
+    s3url : str
+        s3cmd url to the snapshot parquet file
+    """
+
+    return f"{location}/{df_name}/{id}.{file_type}"
+
+
+def list_names(location='s3://snapshots/dataframes/'):
+    """Returns the list of names of the dataframes that have snapshots.
+
+    Parameters
+    ----------
+    location : str
+        s3cmd url to the snapshots folder
+
+    Returns
+    -------
+    list
+        list of dataframe names available in S3
+    """
+
+    s3url = f"{location}"
+    list_s3urls = s3.list_objects(s3url)
+    return list_s3urls
+
+def list_ids(df_name: str, show_progress: bool = False) -> list:
+    """Returns the list of snapshots available for a given dataframe.
+
+    Parameters
+    ----------
+    df_name : str
+        dataframe name
+    show_progress : bool
+        show a progress spinner in the terminal
+
+    Returns
+    -------
+    list
+        list of snapshot ids available for the dataframe
+    """
+
+    s3url = f"s3://vision-ml-datasets/snapshots/dataframes/{df_name}/"
+    l_s3urls = s3.list_objects(s3url, show_progress=show_progress)
+    l_filenames = [x[len(s3url) :] for x in l_s3urls]
+    l_ids = [
+        int(x[2:-8])
+        for x in l_filenames
+        if x.startswith("sn") and x.endswith(".parquet")
+    ]
+    return l_ids
+
+
+def latest_id(df_name: str) -> int:
+    """Returns the list of snapshots available for a given dataframe.
+
+    Parameters
+    ----------
+    df_name : str
+        dataframe name
+
+    Returns
+    -------
+    list
+        list of snapshot ids available for the dataframe
+
+    Raises
+    ------
+    IndexError
+        if there is not snapshot available
+    """
+
+    l_ids = list_ids(df_name, show_progress=False)
+    if l_ids:
+        return max(l_ids)
+    raise IndexError(f"No snapshot available for dataframe '{df_name}'.")
+
+
+def make_new_id(df_name: str) -> int:
+    """Makes a new snapshot id based on the current date.
+
+    Parameters
+    ----------
+    df_name : str
+        dataframe name
+
+    Returns
+    -------
+    int
+        a snapshot id based on the current date
+    """
+    ts = pd.Timestamp.utcnow()
+    return (ts.year % 100) * 10000 + ts.month * 100 + ts.day
+
+
+def save(
+    df: pd.DataFrame,
+    df_name: str,
+    sn_id: Optional[int] = None,
+    local_path: Optional[str] = './snapshot.parquet',
+    logger = None,
+):
+    """Snapshots a dataframe.
+
+    Parameters
+    ----------
+    df : pandas.DataFrame
+        the dataframe to be saved
+    df_name : str
+        dataframe name
+    sn_id : int, optional
+        snapshot id to which the dataframe will be saved. If not provided, the id will be
+        automatically generated.
+    local_path : str, optional
+        local path to save the snapshot. Default is './snapshot.parquet'
+    logger : optional
+        logger to log the snapshotting process. If not provided, no logging will be done.
+    """
+
+    if sn_id is None:
+        sn_id = make_new_id(df_name)
+
+    msg = f"Snapshotting dataframe {df_name} with id {sn_id}"
+    s3url_local_path = get_s3url(df_name, sn_id)
+    logg.info(msg, logger=logger)
+    file = load_any_df(s3url_local_path)
+    file.to_parquet(local_path, index=False)
+    return file
+
+
+def load(
+    df_name: str,
+    sn_id: Optional[int] = None,
+    logger: Optional = None,
+) -> pd.DataFrame:
+    """Loads a dataframe snapshot.
+
+    Parameters
+    ----------
+    df_name : str
+        dataframe name
+    sn_id : int, optional
+        snapshot id to which the dataframe will be saved. If not provided, the id corresponding to
+        the latest snapshot will be used.
+
+    Returns
+    -------
+    pandas.DataFrame
+        the snapnotted dataframe
+    """
+
+    if sn_id is None:
+        sn_id = latest_id(df_name)
+
+    msg = f"Loading a dataframe {df_name} with snapshot id {sn_id}"
+    with logg.scoped_info(msg, logger=logger):
+        s3url = get_s3url(df_name, sn_id)
+        # filepath = cache(s3url, logger=logger)
+
+        # return pd.dfload(filepath, show_progress=bool(logger), unpack=False)

mb/sql/sql.py

@@ -0,0 +1,97 @@
+import sqlalchemy as sa
+import pandas as pd
+from mb.utils.logging import logg
+
+__all__ = ['read_sql','engine_execute','to_sql']
+
+def trim_sql_query(sql_query: str) -> str:
+    """
+    Remove extra whitespace from a SQL query.
+    """
+    sql_query = " ".join(sql_query.splitlines())
+    sql_query = " ".join(sql_query.split())
+    return sql_query
+
+def read_sql(query,engine,index_col=None,chunk_size=10000,logger=None):
+    """Read SQL query into a DataFrame.
+    
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        query (str): SQL query.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        df (pandas.core.frame.DataFrame): DataFrame object.
+    """
+    try:
+
+        if isinstance(query, str):
+            query = trim_sql_query(query)
+            query = sa.text(query)
+        elif isinstance(query, sa.sql.selectable.Select):
+            query = query
+        
+        if chunk_size==None or chunk_size==0:
+            with engine.begin() as conn:
+                df = pd.read_sql(query,conn,index_col=index_col)
+            return df
+            
+        with engine.begin() as conn:
+            df= pd.DataFrame()
+            for chunk in pd.read_sql(query,conn,index_col=index_col,chunksize=chunk_size):
+                df = pd.concat([df,chunk],ignore_index=True)
+            
+        logg.info('SQL query executed successfully.',logger=logger)
+        return df
+    except sa.exc.SQLAlchemyError as e:
+        logg.error('Error executing SQL query.',logger=logger)
+        raise e
+    
+def engine_execute(engine, query_str):
+    """
+    Execute a query on a SQLAlchemy engine object.
+    
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        query_str (str): Query string.
+    Returns:
+        result (sqlalchemy.engine.result.ResultProxy): Result object.
+    """
+    if isinstance(query_str, str):
+        query = sa.text(query_str)
+    else:
+        query = query_str
+    
+    if isinstance(engine, sa.engine.Engine):
+        with engine.begin() as conn:
+            return conn.execute(query)
+    elif isinstance(engine, sa.engine.Connection):
+        return engine.execute(query)
+    
+    
+def to_sql(df,engine,table_name,schema=None,if_exists='replace',index=False,index_label=None,chunksize=10000,logger=None):
+    """Write records stored in a DataFrame to a SQL database.
+    
+    Args:
+        df (pandas.core.frame.DataFrame): DataFrame object.
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        table_name (str): Name of the table.
+        schema (str): Name of the schema. Default: None
+        if_exists (str): How to behave if the table already exists. Default: 'replace'
+        index (bool): Write DataFrame index as a column. Default: False
+        index_label (str): Column label for index column(s). If None is given (default) and index is True, then the index names are used. A sequence should be given if the DataFrame uses MultiIndex. Default: None
+        chunksize (int): Number of rows to write at a time. Default: 10000
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        None
+    """
+    try:
+        if index:
+            if index_label is None:
+                index_label = df.index.name
+        df.to_sql(table_name,engine,schema=schema,if_exists=if_exists,index=index,index_label=index_label,chunksize=chunksize)
+        logg.info(f'DataFrame written to {table_name} table.',logger=logger)
+    except Exception as e:
+        logg.error(f'Error writing DataFrame to {table_name} table.',logger=logger)
+        raise e
+    
+    

mb/sql/tables.py

@@ -0,0 +1,64 @@
+## Tables to be updated every night by the cron job
+
+import sqlalchemy as sa
+import typing as tp
+from .conn import get_engine
+
+
+class TableConfig:
+    """
+    Table configuration object.
+    """
+    def __init__(
+        self,
+        schema: str,
+        table: str,
+        index_col: str,
+        chunk_size: int,
+        updated_col: str,
+        dst_engine: str = "mb_public2",
+        dtype: tp.Optional[dict] = None,
+    ):
+        self.schema = schema
+        self.table = table
+        self.index_col = index_col
+        self.chunk_size = chunk_size
+        self.updated_col = updated_col
+        self.dst_engine = dst_engine
+        self.dtype = dtype
+
+    def get_src_engine(self):
+        if self.schema == "mb_public1":
+            self.src_engine = get_engine(name='postgresql' , db= 'postgres', user='postgres' , password= 'postgres', host= 'localhost', port= 5432, echo=False)
+            return self.src_engine
+        
+    def get_dst_engine(self):
+        if self.schema == "mb_public2":
+            self.dst_engine =get_engine(name='postgresql' , db= 'postgres_2', user='postgres' , password= 'postgres', host= 'localhost', port= 5432, echo=False)
+        return self.dst_engine
+
+
+
+mutable_tables = {
+    'table_to_update1': TableConfig('mb_public1',
+        "test2",
+        None,
+        10000,
+        None,
+        dtype={
+            'name': sa.Text,
+            'age' : sa.Integer
+            },
+    ),
+    
+    'table_to_update2': TableConfig('mb_public1',
+        "test3",
+        'id',
+        10000,
+        None,
+        dtype={
+            'id':sa.Integer,
+            'num': sa.Integer,
+            'data': sa.Text
+        },)   
+    }

mb/sql/update.py

@@ -0,0 +1,709 @@
+import sqlalchemy as sa
+import pandas as pd
+import sqlalchemy.dialects.postgresql
+from .sql import read_sql
+from .utils import list_tables
+import boto3
+
+__all__ = ['get_last_updated_timestamp']
+
+
+def get_last_updated_timestamp(engine,table_name: str,schema : str,updated_col: str = "updated") -> pd.Timestamp:
+    """
+    Gets the last updated timestamp from all records of an sqlite table containing the 'updated' field.
+
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        table_name (str): tables containing the 'updated' field. 
+        schema (str): Name of the schema. Default: None
+        updated_col (str): name of the 'updated' field. Default: 'updated'
+    Returns:
+        ts (pandas.Timestamp): last updated timestamp
+    """
+
+    query_str = f"SELECT MAX({updated_col}) AS last_updated FROM {table_name};"
+    df = read_sql(query_str, engine)
+
+    ts = pd.Timestamp(df["last_updated"][0])
+    if pd.isnull(ts):
+        ts = pd.Timestamp("2014-01-01")
+    return ts
+
+def get_last_updated_data_id(engine,table_name: str,schema : str,updated_col: str = "updated") -> pd.Timestamp:
+    """
+    Gets all the id , updated timestamp from all records of an sqlite table containing the 'updated' field.
+    
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        table_name (str): tables containing the 'updated' field.
+        schema (str): Name of the schema. Default: None
+        updated_col (str): name of the 'updated' field. Default: 'updated'
+    Returns:
+        ts (pandas.Timestamp): last updated timestamp
+    """
+    query_str = f"SELECT id, {updated_col} FROM {table_name};"
+    df = read_sql(query_str, engine)
+    
+    for row in df.iterrows():
+        if pd.isnull(row[updated_col]):
+            row[updated_col] = pd.Timestamp("2014-01-01")
+    return df
+
+
+def get_new_updated_timestamps_and_id_from_s3(bucket_name,table_name, key, last_updated,
+                access_key='your_access_key',
+                secret_key='your_secret_key',region_name='your_region_name'):
+    s3 = boto3.client('s3', aws_access_key_id=access_key,
+                      aws_secret_access_key=secret_key, region_name=region_name)
+    obj = s3.get_object(Bucket=bucket_name, Key=key)
+    data = pd.read_csv(obj['Body'])
+    new_data = data[data['last_updated'] > last_updated]
+    new_rows = []
+    for index, row in new_data.iterrows():
+        if row['id'] not in [r[0] for r in get_last_updated_data_id(table_name)]:
+            new_rows.append(row)
+    return new_rows, new_data.iloc[-1]['last_updated']
+
+# Get new updated timestamps from S3
+def get_new_updated_timestamps_from_s3(bucket_name, key, last_updated,
+                               access_key='your_access_key',
+                               secret_key='your_secret_key',region_name='your_region_name'):
+    """
+    Get new updated timestamps from S3 bucket
+    
+    Args:
+        bucket_name (str): name of the S3 bucket
+        key (str): name of the file in the S3 bucket
+        last_updated (pandas.Timestamp): last updated timestamp
+    Returns:
+        new_data (pandas.DataFrame): dataframe containing 'updated' column.
+    """
+    s3 = boto3.client('s3', aws_access_key_id=access_key,
+                      aws_secret_access_key=secret_key, region_name=region_name)
+    obj = s3.get_object(Bucket=bucket_name, Key=key)
+    data = pd.read_csv(obj['Body'])
+    new_data = data[data['last_updated'] > last_updated]
+    return new_data
+
+#####
+'''
+
+# Update local table
+def update(table_name, data):
+    for index, row in data.iterrows():
+        query = f"UPDATE {table_name} SET column1 = {row['column1']}, column2 = {row['column2']} WHERE id = {row['id']}"
+        cur.execute(query)
+    
+# Read sync via id
+def readsync_via_id(table_name, id):
+    query = f"SELECT * from {table_name} WHERE id = {id}"
+    cur.execute(query)
+    rows = cur.fetchall()
+    return rows
+
+# Merge data
+def merge(table_name, data):
+    conn = create_conn()
+    cur = conn.cursor()
+    for index, row in data.iterrows():
+        query = f"INSERT INTO {table_name} (id, column1, column2) VALUES ({row['id']}, {row['column1']}, {row['column2']}) ON CONFLICT (id) DO UPDATE SET column1 = {row['column1']}, column2 = {row['column2']}"
+        cur.execute(query)
+    conn.commit()
+    conn.close()
+
+# Get duplicates
+def get_duplicates(table_name):
+    query = f"SELECT id, COUNT(*) FROM {table_name} GROUP BY id HAVING COUNT(*) > 1"
+    cur.execute(query)
+    rows = cur.fetchall()
+    return rows
+
+# Get updated fields
+def get_updated_fields(table_name, id):
+    query = f"SELECT column1, column2 from {table_name} WHERE id = {id}"
+    cur.execute(query)
+    rows = cur.fetchall()
+    return rows
+
+
+######
+
+
+def get_new_updated_timestamps(mutable_name: str, mutable_table: dict, last_uts: pd.Timestamp, limit: int = 10000) -> pd.Timestamp:
+    """
+    Obtains all new timestamps and their counts from the 'updated' field of the table.
+
+    Args:
+        mutable_name (str) : name of the table
+        table (dict) : dict containing the 'mutable_tables'.
+        last_uts (pandas.Timestamp) : the last updated timestamp of the corresponding table
+        limit (int) : maximum number of timestamps to be queried
+
+    Returns:
+        df (pd.DataFrame) : dataframe containing 'updated', 'count' columns. The 'updated' field is sorted in ascending order
+    """
+
+    if mutable_name not in mutable_table:
+        raise ValueError("Unknown mutable with name '{}'.".format(mutable_name))
+
+    t = mutable_table[mutable_name]
+    remote_engine = t.get_src_engine()
+    if t.scheme:
+        final_table = f"{t.scheme}.{t.table}"
+    else:
+        final_table = t.table
+    
+    query_str = "SELECT {updated_col}, COUNT(*) AS count FROM {} WHERE {updated_col} > '{}' GROUP BY {updated_col} ORDER BY {updated_col} LIMIT {};".format(
+        final_table, str(last_uts), limit, updated_col=t.updated_col
+    )
+    df = read_sql(query_str, remote_engine)
+
+    return df
+
+
+def _to_sql(df, table_name, engine, dtype=None, from_mysql: bool = False, **kwargs):
+    if from_mysql and isinstance(dtype, dict):
+        for k, v in dtype.items():
+            if v != sa.Text:
+                continue
+            df[k] = df[k].apply(
+                lambda x: x.decode() if isinstance(x, (bytes, bytearray)) else x
+            )
+    return df.to_sql(table_name, engine, dtype=dtype, **kwargs)
+
+
+def update(df: pd.DataFrame, mutable_name: str, mutable_table: dict, index_col: str, is_new_table: bool) -> int:
+    if mutable_name not in mutable_table:
+        raise ValueError("Unknown mutable with name '{}'.".format(mutable_name))
+
+    t = mutable_table[mutable_name]
+    engine = t.get_dst_engine()
+
+    with engine.begin() as conn:  # 
+        if not is_new_table:
+            query_str = ",".join((str(x) for x in df[index_col].tolist()))
+            query_str = "DELETE FROM {} WHERE {} IN ({});".format(
+                mutable_name, index_col, query_str
+            )
+            conn.execute(sa.text(query_str))
+
+        df = df.set_index(index_col, drop=True)
+        _to_sql(
+            df,
+            mutable_name,
+            conn,
+            if_exists="append",
+            dtype=t.dtype
+        )
+
+        return len(df)
+
+
+def readsync_via_id(mutable_name: str,mutable_table: dict,logger=None):
+    """Read-sync a muv1db table containing the 'updated' field.
+
+    Parameters
+    ----------
+    mutable_name : str
+        one of the muv1db tables containing the 'updated' field. Must be a key of `mutable_map`
+        module variable.
+    logger : mt.logg.IndentedLoggerAdapter, optional
+        logger for debugging purposes
+    """
+
+    if mutable_name not in mutable_table:
+        raise ValueError("Unknown mutable with name '{}'.".format(mutable_name))
+
+    t = mutable_table[mutable_name]
+    engine = t.get_dst_engine()
+    schema = t.schema
+    table = t.table
+    index_col = t.index_col
+    chunk_size = t.chunk_size
+
+    remote_engine = t.get_src_engine()
+    if schema:
+        final_table = f"{schema}.{table}"
+    else:
+        final_table = table
+    
+    
+    if mutable_name in list_tables(engine):
+        is_new_table = False
+        query_str = "SELECT MAX({index_col}) AS last_id FROM {mutable_name}".format(
+            index_col=index_col, mutable_name=mutable_name
+        )
+        last_id = read_sql(query_str, engine)["last_id"][0]
+    else:
+        is_new_table = True
+        last_id = 0
+
+    msg = "Read-syncing mutable '{}' with last id {}".format(mutable_name, last_id)
+    with logg.scoped_info(msg, logger=logger):
+        first_id = last_id
+        query_str = "SELECT MAX({index_col}) AS last_id FROM {final_table}".format(
+            index_col=index_col, frame_str=final_table
+        )
+        last_id = read_sql(query_str, remote_engine)["last_id"][0]
+        if logger:
+            logger.debug("Remote has last id {}.".format(last_id))
+
+        query_str = "SELECT COUNT(*) AS cnt FROM {frame_str} WHERE {index_col} > {first_id} AND {index_col} <= {last_id};".format(
+            index_col=index_col, frame_str=final_table, first_id=first_id, last_id=last_id
+        )
+        count = read_sql(query_str, remote_engine)["cnt"][0]
+        if count == 0:
+            return
+
+        # determine new fold
+        fold = (count // chunk_size) + 1
+        if logger:
+            logger.info(
+                "  Found {} records to be downloaded in {} chunks.".format(count, fold)
+            )
+
+        for i in range(fold):
+            start_ofs = first_id + (last_id - first_id) * i // fold
+            end_ofs = first_id + (last_id - first_id) * (i + 1) // fold
+
+            query_str = "SELECT * FROM {frame_str} WHERE {index_col}>{first_id} and {index_col}<={last_id};".format(
+                frame_str=frame_str,
+                index_col=index_col,
+                first_id=start_ofs,
+                last_id=end_ofs,
+            )
+            df = ss.read_sql(query_str, remote_engine)
+
+            if len(df) == 0:
+                continue
+
+            if logger:
+                logger.info(
+                    "  {}: Downloaded {} records with id in [{},{}].".format(
+                        i + 1, len(df), start_ofs, end_ofs
+                    )
+                )
+
+            update(df, mutable_name, index_col, is_new_table)
+
+
+def readsync(mutable_name: str, logger=None):
+    """Read-sync a muv1db table containing the 'updated' field.
+
+    Parameters
+    ----------
+    mutable_name : str
+        one of the muv1db tables containing the 'updated' field. Must be a key of `mutable_map`
+        module variable.
+    logger : mt.logg.IndentedLoggerAdapter, optional
+        logger for debugging purposes
+    """
+
+    t = mutable_map[mutable_name]
+
+    if t.updated_col is None:
+        return readsync_via_id(mutable_name, logger=logger)
+
+    schema = t.schema
+    table = t.table
+    index_col = t.index_col
+    chunk_size = t.chunk_size
+    updated_col = t.updated_col
+
+    remote_engine = t.get_src_engine()
+    frame_str = ss.frame_sql(table, schema=schema)
+
+    last_uts = get_last_updated_timestamp(mutable_name, updated_col=updated_col)
+    is_new_table = last_uts < pd.Timestamp("2017-01-02")
+
+    msg = "Read-syncing mutable '{}' with timestamp '{}'".format(
+        mutable_name, str(last_uts)
+    )
+    with logg.scoped_info(msg, logger=logger):
+        count = 0
+        while True:
+            thresh_uts = pd.Timestamp.utcnow().tz_localize(None) - pd.Timedelta(
+                1, "hours"
+            )
+            if last_uts > thresh_uts:
+                break
+
+            # large enough so that the probing+sorting op is neglible compared to the downloading op
+            fold = 16
+
+            if logger:
+                logger.info(
+                    "Inspecting maximum {} timestamps after '{}':".format(
+                        chunk_size * fold, str(last_uts)
+                    )
+                )
+            df0 = get_new_updated_timestamps(
+                mutable_name, last_uts, limit=chunk_size * fold
+            )
+            if len(df0) == 0:
+                break
+
+            # determine new fold
+            count_sum = df0["count"].sum()
+            fold = (count_sum // chunk_size) + 1
+            if logger:
+                logger.info(
+                    "  Found {} records to be downloaded in {} chunks.".format(
+                        count_sum, fold
+                    )
+                )
+
+            df0_cnt = len(df0)
+            for i in range(fold):
+                start_ofs = df0_cnt * i // fold
+                end_ofs = (df0_cnt * (i + 1)) // fold
+
+                df = df0.iloc[start_ofs:end_ofs]
+
+                first_uts = last_uts
+                last_uts = df[updated_col].max()
+                chunk_cnt = df["count"].sum()
+
+                if logger:
+                    logger.info(
+                        "  {}: Downloading {} records timestamped in ['{}','{}'].".format(
+                            i + 1, chunk_cnt, str(first_uts), str(last_uts)
+                        )
+                    )
+                query_str = "SELECT * FROM {frame_str} WHERE {updated_col}>'{first_uts}' and {updated_col}<='{last_uts}';".format(
+                    frame_str=frame_str,
+                    updated_col=updated_col,
+                    first_uts=str(first_uts),
+                    last_uts=str(last_uts),
+                )
+                df = ss.read_sql(query_str, remote_engine)
+
+                count += update(df, mutable_name, index_col, is_new_table)
+
+        if logger:
+            logger.info(
+                "Downloaded {} records and updated timestamp to '{}'.".format(
+                    count, str(last_uts)
+                )
+            )
+
+
+def merge_from(mutable_name: str, other_engine, logger=None):
+    """Merges a mutable from another sqlite database to this database.
+
+    Parameters
+    ----------
+    mutable_name : str
+        mutable to merge. Must be a key of `mutable_map` module variable.
+    other_engine : sqlalchemy.engine.Engine
+        connection engine to an sqlite3 wml database that contains the mutable to merge from
+    logger : mt.logg.IndentedLoggerAdapter, optional
+        logger for debugging purposes
+
+    Returns
+    -------
+    bool
+        whether the table has been successfully merged
+    """
+
+    if not mutable_name in mutable_map:
+        raise ValueError(
+            "First argument must be a valid mutable name in the 'mutable_map'. Got: '{}'.".format(
+                mutable_name
+            )
+        )
+    t = mutable_map[mutable_name]
+    if t.dst_engine != "wml":
+        logger.info(f"Skipping mutable '{mutable_name}' as it is not in wml db.")
+        return False
+
+    msg = "Merging mutable '{}' from '{}'".format(mutable_name, other_engine.url)
+    with logg.scoped_info(msg, logger):
+        if not mutable_name in ss.list_tables(other_engine):
+            if logger:
+                logger.info("Remote mutable does not exist.")
+            return False
+
+        # merge structure
+        remote_sql_str = ss.get_table_sql_code(mutable_name, other_engine)
+        if mutable_name in ss.list_tables(wml_engine):  # local table exists
+            local_sql_str = ss.get_table_sql_code(mutable_name, wml_engine)
+            if local_sql_str != remote_sql_str:
+                if logger:
+                    logger.debug(
+                        "Remote table has a different sql definition from the local one:"
+                    )
+                    logger.debug("  Local : {}".format(local_sql_str))
+                    logger.debug("  Remote: {}".format(remote_sql_str))
+                    logger.info("Skipping.")
+                return False
+        else:  # local table does not exist
+            if logger:
+                logger.info("Creating local mutable.")
+                logger.debug("  SQL: {}".format(remote_sql_str))
+            ss.engine_execute(wml_engine, remote_sql_str)
+
+        # merge index
+        remote_indices = ss.list_indices(other_engine)
+        remote_indices = getattr(remote_indices, mutable_name, {})
+        local_indices = ss.list_indices(wml_engine)
+        local_indices = getattr(local_indices, mutable_name, {})
+        for index_col in remote_indices:
+            if index_col in local_indices:
+                continue
+            index_sql_str = remote_indices[index_col]
+            if logger:
+                logger.info("Creating index '{}'.".format(index_col))
+                logger.debug("  SQL: {}".format(index_sql_str))
+            ss.engine_execute(wml_engine, index_sql_str)
+
+        # identify duplicate records
+        t = mutable_map[mutable_name]
+        with logg.scoped_info("Comparing local <-> remote headers", logger=logger):
+            if t.updated_col is None:
+                if logger:
+                    logger.info("Loading local ids:")
+                query_str = "SELECT {index_col} FROM {mutable_name}".format(
+                    index_col=t.index_col, mutable_name=mutable_name
+                )
+                local_id_df = ss.read_sql(query_str, wml_engine)
+                local_id_list = local_id_df[index_col].tolist()
+                if logger:
+                    logger.debug("  {} ids loaded.".format(len(local_id_df)))
+                    logger.info("Loading remote ids:")
+
+                query_str = "SELECT {index_col} FROM {mutable_name}".format(
+                    index_col=t.index_col, mutable_name=mutable_name
+                )
+                remote_id_df = ss.read_sql(query_str, other_engine)
+                remote_id_list = remote_id_df[t.index_col].tolist()
+                if logger:
+                    logger.debug("  {} ids loaded.".format(len(remote_id_df)))
+
+                local_set = set(local_id_list)
+                remote_set = set(remote_id_list)
+                drop_list = []
+                insert_list = list(remote_set - local_set)
+
+                if logger:
+                    logger.info(
+                        "{} remote records to be inserted.".format(len(insert_list))
+                    )
+            else:
+                if logger:
+                    logger.info("Loading local headers:")
+                query_str = "SELECT {index_col}, {updated_col} AS local_updated FROM {mutable_name}".format(
+                    index_col=t.index_col,
+                    updated_col=t.updated_col,
+                    mutable_name=mutable_name,
+                )
+                local_id_df = ss.read_sql(query_str, wml_engine, index_col=t.index_col)
+                if logger:
+                    logger.debug("  {} headers loaded.".format(len(local_id_df)))
+                    logger.info("Loading remote headers:")
+
+                query_str = "SELECT {index_col}, {updated_col} AS remote_updated FROM {mutable_name}".format(
+                    index_col=t.index_col,
+                    updated_col=t.updated_col,
+                    mutable_name=mutable_name,
+                )
+                remote_id_df = ss.read_sql(
+                    query_str, other_engine, index_col=t.index_col
+                )
+                if logger:
+                    logger.debug("  {} headers loaded.".format(len(remote_id_df)))
+
+                # records in local but not in remote
+                id_df = local_id_df.join(remote_id_df, how="outer").reset_index(
+                    drop=False
+                )
+                drop_list = []
+                insert_list = []
+                the_list = id_df[id_df["remote_updated"].isnull()][t.index_col].tolist()
+                if logger:
+                    logger.debug(
+                        "{} local records not in remote.".format(len(the_list))
+                    )
+
+                # records in remote but not in local
+                id_df = id_df[id_df["remote_updated"].notnull()]
+                the_list = id_df[id_df["local_updated"].isnull()][t.index_col].tolist()
+                insert_list.extend(the_list)
+                if logger:
+                    logger.debug(
+                        "{} remote records not in local.".format(len(the_list))
+                    )
+
+                # duplicate records
+                id_df = id_df[id_df["local_updated"].notnull()]
+                common_list = id_df[id_df["local_updated"] == id_df["remote_updated"]][
+                    t.index_col
+                ].tolist()
+                if logger:
+                    logger.info(
+                        "{} common records to be kept.".format(len(common_list))
+                    )
+
+                # new local records
+                id_df = id_df[id_df["local_updated"] != id_df["remote_updated"]]
+                the_list = id_df[id_df["local_updated"] > id_df["remote_updated"]][
+                    t.index_col
+                ].tolist()
+                if logger:
+                    logger.debug("{} new local records detected.".format(len(the_list)))
+
+                # new remote records
+                id_df = id_df[id_df["local_updated"] < id_df["remote_updated"]]
+                the_list = id_df[t.index_col].tolist()
+                drop_list.extend(the_list)
+                insert_list.extend(the_list)
+                if logger:
+                    logger.debug(
+                        "{} new remote records detected.".format(len(the_list))
+                    )
+                    logger.info(
+                        "{} local records to be dropped.".format(len(drop_list))
+                    )
+                    logger.info(
+                        "{} remote records to be inserted.".format(len(insert_list))
+                    )
+
+        # drop records
+        if drop_list:
+            if logger:
+                logger.info("Dropping {} local records.".format(len(drop_list)))
+            query_str = ",".join((str(x) for x in drop_list))
+            query_str = (
+                "DELETE FROM {mutable_name} WHERE {index_col} IN ({values});".format(
+                    mutable_name=mutable_name, index_col=t.index_col, values=query_str
+                )
+            )
+            ss.engine_execute(wml_engine, query_str)
+
+        # insert remote records
+        if insert_list:
+            msg = "Upserting {} remote records".format(len(insert_list))
+            with logg.scoped_info(msg, logger=logger):
+                while insert_list:
+                    the_list = insert_list[:1000000]
+                    insert_list = insert_list[1000000:]
+
+                    query_str = ",".join((str(x) for x in the_list))
+                    query_str = "SELECT * FROM {mutable_name} WHERE {index_col} IN ({values});".format(
+                        mutable_name=mutable_name,
+                        index_col=t.index_col,
+                        values=query_str,
+                    )
+                    if logger:
+                        logger.info(
+                            "Loading {} remote records, {} remaining.".format(
+                                len(the_list), len(insert_list)
+                            )
+                        )
+                    df = ss.read_sql(query_str, other_engine, logger=logger)
+                    if logger:
+                        logger.info("Inserting {} new records.".format(len(df)))
+                    _to_sql(
+                        df,
+                        mutable_name,
+                        wml_engine,
+                        index=False,
+                        if_exists="append",
+                        dtype=t.dtype,
+                        from_mysql=t.schema == "winnow_db",
+                    )
+
+        if logger:
+            logger.info("Merge completed.")
+
+    return True
+
+
+def vacuum(level: str = "full", logger=None):
+    """Cleans up duplicate records in mutables and then make the wml database compact.
+
+    Parameters
+    ----------
+    level : {'full', 'minimal'}
+        If 'minimal', only the mutables are vacuumed. Otherwise, in addition to 'minima', the whole
+        wml database is vacuumed.
+    logger : mt.logg.IndentedLoggerAdapter, optional
+        logger for debugging purposes
+    """
+
+    wml_tables = ss.list_tables(wml_engine)
+
+    for mutable_name, t in mutable_map.items():
+        if (mutable_name not in wml_tables) or (t.dst_engine != "wml"):
+            continue
+
+        msg = "Vacuuming mutable '{}'".format(mutable_name)
+        with logg.scoped_info(msg, logger=logger):
+            schema = t.schema
+            table = t.table
+            index_col = t.index_col
+            chunk_size = t.chunk_size
+            updated_col = t.updated_col
+            query_str = """WITH t1 AS (SELECT {index_col}, COUNT(*) AS id_cnt FROM {mutable_name} GROUP BY {index_col})
+                SELECT {mutable_name}.* FROM {mutable_name}
+                    LEFT JOIN t1 ON {mutable_name}.{index_col}=t1.{index_col}
+                  WHERE t1.id_cnt > 1
+                ;""".format(
+                mutable_name=mutable_name,
+                index_col=index_col,
+            )
+            df = ss.read_sql(query_str, wml_engine, logger=logger)
+            if len(df) == 0:
+                if logger:
+                    logger.info("The table is clean.")
+                continue
+
+            id_list = df[index_col].drop_duplicates().tolist()
+            if logger:
+                logger.info(
+                    "Detected {} duplicate ids spanning over {} records.".format(
+                        len(id_list), len(df)
+                    )
+                )
+
+            query_str = ",".join((str(x) for x in id_list))
+            query_str = (
+                "DELETE FROM {mutable_name} WHERE {index_col} IN ({values});".format(
+                    mutable_name=mutable_name,
+                    index_col=index_col,
+                    values=query_str,
+                )
+            )
+            ss.engine_execute(wml_engine, query_str)
+            if logger:
+                logger.info("Deleted {} duplicate records.".format(len(df)))
+
+            if updated_col is None:
+                df = df.groupby(index_col).head(1)
+            else:
+                df = (
+                    df.sort_values([index_col, updated_col], ascending=[True, False])
+                    .groupby(index_col)
+                    .head(1)
+                )
+            df = df.set_index(index_col, drop=True)
+            _to_sql(
+                df,
+                mutable_name,
+                wml_engine,
+                if_exists="append",
+                dtype=t.dtype,
+                from_mysql=t.schema == "winnow_db",
+            )
+            if logger:
+                logger.info("Reinserted {} clean records.".format(len(df)))
+
+    if level == "full":
+        if logger:
+            logger.info("Vacuuming the wml database.")
+        ss.vacuum(wml_engine)
+
+    if logger:
+        logger.info("Done.")
+        
+'''

mb/sql/utils.py

@@ -0,0 +1,148 @@
+## file for extra sql functions
+
+import pandas as pd
+from .sql import read_sql,engine_execute
+import os
+from mb.utils.logging import logg
+
+__all__ = ['list_schemas','rename_table','drop_table','drop_schema','create_schema','create_index','clone_db']
+
+def list_schemas(engine,logger=None) -> pd.DataFrame:
+    """
+    Returns list of schemas in database.
+    
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        df (pandas.core.frame.DataFrame): DataFrame object.
+    """
+    q1 ="SELECT schema_name FROM information_schema.schemata WHERE schema_name NOT IN ('information_schema', 'mysql', 'performance_schema') ORDER BY schema_name;"
+    return read_sql(q1,engine,logger=logger)
+
+def list_tables(engine,schema=None,logger=None) -> pd.DataFrame:
+    """
+    Returns list of tables in database.
+    
+    Args:
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        schema (str): Name of the schema. Default: None
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        df (pandas.core.frame.DataFrame): DataFrame object.
+    
+    """
+    q1 = f"SELECT table_name FROM information_schema.tables WHERE table_schema = '{schema}' ORDER BY table_name;"
+    return read_sql(q1,engine,logger=logger)
+    
+
+def rename_table(new_table_name,old_table_name,engine,schema_name=None,logger=None):
+    """
+    Rename table in database.
+    
+    Args:
+        new_table_name (str): New name of the table.
+        old_table_name (str): Old name of the table.
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        schema_name (str): Name of the schema. Default: None
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        None
+    """
+    if schema_name:
+        q1 = f"ALTER TABLE {schema_name}.{old_table_name} RENAME TO {new_table_name};"
+    else:
+        q1 = f"ALTER TABLE {old_table_name} RENAME TO {new_table_name};"
+    engine_execute(engine,q1)
+    logg.info(f'Table {old_table_name} renamed to {new_table_name}.',logger=logger)
+
+def drop_table(table_name,engine,schema_name=None,logger=None):
+    """
+    Drop table in database.
+    
+    Args:
+        table_name (str): Name of the table.
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        schema_name (str): Name of the schema. Default: None
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        None
+    """
+    if schema_name:
+        q1 = f"DROP TABLE {schema_name}.{table_name};"
+    else:
+        q1 = f"DROP TABLE {table_name};"
+    engine_execute(engine,q1)
+    logg.info(f'Table {table_name} dropped.',logger=logger)
+        
+def drop_schema(schema_name,engine,logger=None):
+    """
+    Drop schema in database.
+    
+    Args:
+        schema_name (str): Name of the schema.
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        None
+    """
+    q1 = f"DROP SCHEMA {schema_name};"
+    engine_execute(engine,q1)
+    logg.info(f'Schema {schema_name} dropped.',logger=logger)
+        
+def create_schema(schema_name,engine,logger=None):
+    """
+    Create schema in database.
+    
+    Args:
+        schema_name (str): Name of the schema.
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        None
+    """
+    q1 = f"CREATE SCHEMA {schema_name};"
+    engine_execute(engine,q1)
+    logg.info(f'Schema {schema_name} created.',logger=logger)
+
+def create_index(table,index_col,engine,logger=None):
+    """
+    Create an Index for a table in database.
+    
+    Args:
+        table (str): Name of the table.
+        index_col (str): Name of the index_col in the table.
+        engine (sqlalchemy.engine.base.Engine): Engine object.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        None
+    """
+    q1 = f"CREATE INDEX {index_col} ON {table};"
+    engine_execute(engine,q1)
+    logg.info(f'Index {index_col} created for table {table}.',logger=logger)
+        
+def clone_db(ori_db_location,copy_db_location, logger=None):
+    """
+    Clone a database.
+    
+    Args:
+        ori_db_location (str): Location of the original database.
+        copy_db_location (str): Location for the new database.
+        logger (logging.Logger): Logger object. Default: mb_utils.src.logging.logger
+    Returns:
+        None
+    """
+    if not os.path.exists(ori_db_location):
+        raise FileNotFoundError("The original location does not exist.")
+
+    if os.path.exists(copy_db_location):
+        copy_db_location = os.path.join(copy_db_location,'copy_db')
+    
+    if os.path.exists(copy_db_location)==False:
+        os.makedirs(copy_db_location)
+    
+    cmd = f"pg_dump -U postgres -h localhost -p 5432 {ori_db_location} | psql -U postgres -h localhost -p 5432 {copy_db_location}"
+    os.system(cmd)
+
+    logg.info(f'Database cloned from {ori_db_location} to {copy_db_location}.',logger=logger)    
+    

mb/sql/version.py

@@ -0,0 +1,5 @@
+MAJOR_VERSION = 2
+MINOR_VERSION = 0
+PATCH_VERSION = 1
+version = '{}.{}.{}'.format(MAJOR_VERSION, MINOR_VERSION, PATCH_VERSION)
+__all__  = ['MAJOR_VERSION', 'MINOR_VERSION', 'PATCH_VERSION', 'version']

mb_sql-2.0.1.dist-info/METADATA

@@ -0,0 +1,11 @@
+Metadata-Version: 2.4
+Name: mb_sql
+Version: 2.0.1
+Summary: mb SQL functions API
+Author: ['Malav Bateriwala']
+Requires-Python: >=3.8
+Requires-Dist: mb_utils
+Dynamic: author
+Dynamic: requires-dist
+Dynamic: requires-python
+Dynamic: summary

mb_sql-2.0.1.dist-info/RECORD

@@ -0,0 +1,13 @@
+mb/sql/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+mb/sql/conn.py,sha256=HTyAUDkBlp_yvDRNRS9OoTRJPvrcS2SEvnx46AXK3OA,3210
+mb/sql/slack.py,sha256=3dWbKvhnp1sIEWGk3m43WnfXbtW56Ohk41_5puZnH-A,549
+mb/sql/snapshot.py,sha256=3C7vqRE3vUwyEEQQIsSIfz2JMjli2hV5Tm4fONbgEbo,4775
+mb/sql/sql.py,sha256=cj00ta1j53o7VRf3LeXj9uGE8p1IcUF7CUFbWrVM0AA,3655
+mb/sql/tables.py,sha256=xECYrXJKvdGvCPZo9vT1o7-HO0LoFjaD8602BQo4oO4,1669
+mb/sql/update.py,sha256=Ed9MA3y9ZJ3arwSq9cbIPBe1hI5PCqmPyY6rD5huq64,26741
+mb/sql/utils.py,sha256=Bh1eadnfnckeDJg25rhnrDXJ3YVN8gjp2AedgyM8Iww,5248
+mb/sql/version.py,sha256=JJiVhzInvyc4XUc31yQR_q4UNm6rnom-oatXO_VvCfo,206
+mb_sql-2.0.1.dist-info/METADATA,sha256=tOr2lC_YFhpIcrBirbHz_tsq7DsmlTvsaX8AkmRdLzM,237
+mb_sql-2.0.1.dist-info/WHEEL,sha256=YCfwYGOYMi5Jhw2fU4yNgwErybb2IX5PEwBKV4ZbdBo,91
+mb_sql-2.0.1.dist-info/top_level.txt,sha256=2T5lqIVZs7HUr0KqUMPEaPF59QXcr0ErDy6K-J88ycM,3
+mb_sql-2.0.1.dist-info/RECORD,,

mb_sql-2.0.1.dist-info/WHEEL

@@ -0,0 +1,5 @@
+Wheel-Version: 1.0
+Generator: setuptools (82.0.0)
+Root-Is-Purelib: true
+Tag: py3-none-any
+

mb_sql-2.0.1.dist-info/top_level.txt

@@ -0,0 +1 @@
+mb