monitor: Batch processing

- Add Slowly Changing Dimensions logic to database upload process
- Keep bot logged in. When logging out, messages can be lost / not synced properly. A logged in account can fetch all the messages that it's seen.
- Store latest timestamps per chat

monitor.py

@@ -134,13 +134,15 @@ def save_file(file_path: str, data: Any):
     with open(file_path, "wb") as f:
         pickle.dump(data, f)
 
-def download(folder: str, config: dict):
+def create_bot(config: dict) -> Account:
     """
-    Download Status App messages / info from communities and store them in pickle files.
+    Initialized a logged in bot account that will monitor the communities.
 
     Parameters:
-        - `folder` - the folder where the files will be created. Sub folders are automatically created
         - `config` - the `load_config` configuration
+
+    Output:
+        - Logged in Bot account
     """
     account = Account(**config.get("bot_params", {}))
     available_accounts = [acc["display_name"] for acc in account.available_accounts]
@@ -156,6 +158,16 @@ def download(folder: str, config: dict):
 
     account.login(**params)
     account.logger.info(f"Account Information:\nCompressed Key: {account.info['compressed_key']}\nPublic Key: {account.info['public_key']}\nURL: {account.info['url']}")
+    return account
+
+def download(account: Account, folder: str, config: dict):
+    """
+    Download Status App messages / info from communities and store them in pickle files.
+
+    Parameters:
+        - `folder` - the folder where the files will be created. Sub folders are automatically created
+        - `config` - the `load_config` configuration
+    """
     file_path = os.path.join(os.path.dirname(__file__), config["files"]["current_state"])
     latest_dates: dict[str, pd.Timestamp] = pd.read_pickle(file_path) if os.path.exists(file_path) else {}
 
@@ -197,7 +209,7 @@ def download(folder: str, config: dict):
                 save_file(file_path, messages)
                 account.logger.info(f"Created {file_path}")
 
-def store(folder: str, config: dict):
+def store(folder: str, config: dict, logger: Logger):
     """
     Upload Status App `download` file to Postgres.
     NOTE: The Postgres schema must already exist
@@ -211,10 +223,14 @@ def store(folder: str, config: dict):
     table_schema = config["postgres"]["schema"]
 
     upload: dict[str, list] = {}
-    latest_dates: dict[str, pd.Timestamp] = {}
+
+    file_path = os.path.join(os.path.dirname(__file__), config["files"]["current_state"])
+    latest_dates: dict[str, pd.Timestamp] = pd.read_pickle(file_path) if os.path.exists(file_path) else {}
+
     completed = []
 
     files = list(path.rglob("*.pkl")) + list(path.rglob("*.csv"))
+    logger.info(f"There are {len(files)} file(s) to upload")
     for file_path in files:
 
         table_name = table_name_mapping.get(file_path.parent.parent.name)
@@ -240,8 +256,8 @@ def store(folder: str, config: dict):
         upload[table_name].append(data)
         completed.append(str(file_path))
 
-    if latest_dates:
-        save_file(config["files"]["current_state"], latest_dates)
+    save_file(config["files"]["current_state"], latest_dates)
+    logger.info(f"Updated {config['files']['current_state']}")
 
     prefix = "POSTGRES_"
     params = {
@@ -258,22 +274,24 @@ def store(folder: str, config: dict):
         json_columns = [
             column
             for column in df.columns
-            if isinstance(df[column].dropna().reset_index(drop=True).iloc[0], (dict, list))
+            if len(df[column].dropna()) > 0 and isinstance(df[column].dropna().reset_index(drop=True).iloc[0], (dict, list))
         ]
         connector.insert(df, table_name, table_schema, json_columns)
+        logger.info(f"Uploaded {len(df)} record(s) to {table_schema}.{table_name}")
 
     for file_path in completed:
         os.remove(file_path)
-
+        logger.info(f"Deleted {file_path}")
 
 if __name__ == "__main__":
     folder = os.path.dirname(__file__)
     config = load_config(os.path.join(folder, "config.yaml"))
     upload_folder = os.path.join(os.path.dirname(__file__), "uploads")
     logger = Logger()
+    account = create_bot(config)
 
     while True:
-        download(upload_folder, config)
-        store(upload_folder, config)
+        download(account, upload_folder, config)
+        store(upload_folder, config, logger)
         logger.info(f"Sleeping for {config['sleep']} minute(s)")
-        time.sleep(config["sleep"])
+        time.sleep(config["sleep"] * 60)

postgres.py

@@ -57,6 +57,16 @@ class Postgres:
                 for json_column in json_columns
             }
 
+        # Add new columns as they come
+        existing_columns = self.get_columns(schema, table_name)
+
+        if existing_columns:
+            for column in data.columns:
+                if column in existing_columns:
+                    continue
+                #  NOTE: New values will have to be transformed
+                self.execute(f"ALTER TABLE {schema}.{table_name} ADD COLUMN {column} TEXT")
+
         data.to_sql(**params)
 
     def execute(self, query: str):
@@ -69,6 +79,29 @@ class Postgres:
         self.__execute(query)
         self.__conn.commit()
 
+    def to_pandas(self, query: str, batch_size: int = 50_000, uppercase: bool = True) -> pd.DataFrame:
+        """
+        Create a DataFrame from the given query
+
+        Parameters:
+            - `query` - the PostgreSQL query
+            - `batch_size` - how many rows will be fetched at once
+            - `uppercase` - if `True` then the columns will be uppercase. If `False` the columns will be lowercase
+        Output:
+            - DataFrame for the executed query
+        """
+        self.__execute(query)
+        columns = [column.name.upper() if uppercase else column.name.lower() for column in self.__cursor.description]
+        chunks = []
+
+        while True:
+            rows = self.__cursor.fetchmany(batch_size)
+            if not rows:
+                break
+            chunks.append(pd.DataFrame(rows, columns=columns))
+
+        return pd.concat(chunks, ignore_index=True) if chunks else pd.DataFrame(columns=columns)
+
     def close(self):
         self.__cursor.close()
         self.__conn.close()
@@ -93,3 +126,24 @@ class Postgres:
 
         if failed:
             self.__cursor.execute(query)
+
+
+    def get_columns(self, schema: str, table_name: str) -> list[str]:
+        """
+        Get the column names in the correct order for the given table.
+
+        Parameters:
+            - `table_name` - the name of the table
+            - `schema` - the name of the schema
+
+        Output:
+            - the table's columns in the correct order
+        """
+        query = f"""
+        SELECT column_name
+        FROM information_schema.columns
+        WHERE table_name = '{table_name}'
+        AND table_schema = '{schema}'
+        ORDER BY ordinal_position ASC
+        """
+        return self.to_pandas(query)["COLUMN_NAME"].to_list()