Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Script to purge all unprocessed messages #623

Merged
merged 2 commits into from
Jan 8, 2025
Merged

Script to purge all unprocessed messages #623

Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
10c4cb8
Script to purge all unprocessed messages
kongzii Jan 7, 2025
b9b5e4e
fix
kongzii Jan 7, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
33 changes: 33 additions & 0 deletions scripts/purge_messages.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -0,0 +1,33 @@
import typer
from prediction_market_agent_tooling.gtypes import private_key_type

from prediction_market_agent.db.agent_communication import (
fetch_count_unprocessed_transactions,
pop_message,
)
from prediction_market_agent.utils import APIKeys


def main(private_key: str) -> None:
keys = APIKeys(BET_FROM_PRIVATE_KEY=private_key_type(private_key))
n_messages = fetch_count_unprocessed_transactions(
consumer_address=keys.bet_from_address
)
Comment on lines +13 to +15
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

⚠️ Potential issue

Add error handling for database operations.

The database operations should be wrapped in try-except blocks to handle potential connection issues.

+    try:
         n_messages = fetch_count_unprocessed_transactions(
             consumer_address=keys.bet_from_address
         )
+    except Exception as e:
+        raise ConnectionError(f"Failed to fetch unprocessed messages: {str(e)}")
+
+    if n_messages == 0:
+        print("No messages to purge.")
+        return
📝 Committable suggestion

‼️ IMPORTANT
Carefully review the code before committing. Ensure that it accurately replaces the highlighted code, contains no missing lines, and has no issues with indentation. Thoroughly test & benchmark the code to ensure it meets the requirements.

Suggested change
n_messages = fetch_count_unprocessed_transactions(
consumer_address=keys.bet_from_address
)
try:
n_messages = fetch_count_unprocessed_transactions(
consumer_address=keys.bet_from_address
)
except Exception as e:
raise ConnectionError(f"Failed to fetch unprocessed messages: {str(e)}")
if n_messages == 0:
print("No messages to purge.")
return


if (
input(
f"Are you sure you want to purge all {n_messages} messages for agent {keys.bet_from_address}? (y/n): "
)
!= "y"
):
return

popped = 0
while fetch_count_unprocessed_transactions(consumer_address=keys.bet_from_address):
pop_message(api_keys=keys)
Comment on lines +26 to +27
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

I would have expected that fetch_count_unprocessed_transactions is called only once, yielding N unprocessed transactions. Afterwards, a for loop is created with exactly N, and those are then deleted.
Any reason for continuously sending RPC calls?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

In case it goes out of sync for some reason. But it's just script to be run "once in a decade", does that really matter?

popped += 1
print(f"Popped {popped} messages.")
Comment on lines +17 to +29
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

🛠️ Refactor suggestion

Improve robustness of the purging logic.

The current implementation has several potential issues:

  1. No timeout mechanism for the while loop
  2. No error handling during message purging
  3. No batch processing for efficiency
  4. No progress indication for large message counts

Apply this diff to improve the implementation:

+    import time
+    from tqdm import tqdm
+
     if (
         input(
             f"Are you sure you want to purge all {n_messages} messages for agent {keys.bet_from_address}? (y/n): "
         )
         != "y"
     ):
         return

     popped = 0
-    while fetch_count_unprocessed_transactions(consumer_address=keys.bet_from_address):
-        pop_message(api_keys=keys)
-        popped += 1
-        print(f"Popped {popped} messages.")
+    start_time = time.time()
+    timeout = 300  # 5 minutes timeout
+    
+    with tqdm(total=n_messages, desc="Purging messages") as pbar:
+        while (
+            fetch_count_unprocessed_transactions(consumer_address=keys.bet_from_address)
+            and time.time() - start_time < timeout
+        ):
+            try:
+                pop_message(api_keys=keys)
+                popped += 1
+                pbar.update(1)
+            except Exception as e:
+                print(f"\nError while purging message {popped + 1}: {str(e)}")
+                if input("\nContinue purging? (y/n): ") != "y":
+                    break
+
+    if popped < n_messages:
+        print(f"\nPurge incomplete. {n_messages - popped} messages remaining.")

Committable suggestion skipped: line range outside the PR's diff.



if __name__ == "__main__":
typer.run(main)
Comment on lines +1 to +33
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

💡 Codebase verification

Add security warning and restrict file permissions

The script requires security improvements:

  • Change file permissions to 600 (chmod 600 scripts/purge_messages.py) to restrict access since it handles private keys
  • Add a warning message at the start of the script:
def main(private_key: str) -> None:
    """
    WARNING: This script performs irreversible deletion of messages.
    Handle with care as it uses private keys and permanently removes data.
    """
🔗 Analysis chain

Add security considerations and file permissions.

  1. Add a warning about the destructive nature of this script
  2. Ensure the script has appropriate file permissions

Run this script to check and fix file permissions:

🏁 Scripts executed

The following scripts were executed for the analysis:

Script:

#!/bin/bash
# Description: Check and fix file permissions for the script

# Check current permissions
current_perms=$(stat -c "%a" scripts/purge_messages.py)
echo "Current permissions: $current_perms"

# Ensure the script is not world-readable/writable due to sensitive data handling
if [[ $current_perms =~ ^[0-7]{3}$ && $current_perms != "600" ]]; then
    echo "Warning: Script should have restricted permissions (600) due to handling sensitive data"
fi

Length of output: 229

Script:

#!/bin/bash
# Check file permissions using ls and search for warning messages
echo "File permissions:"
ls -l scripts/purge_messages.py

echo -e "\nChecking for existing warning messages:"
rg -i "warning|caution|danger" scripts/purge_messages.py

Length of output: 306

Loading