Add prometheus telemetry

telegram_downloader_bot/main.py

@@ -10,8 +10,11 @@ from telegram_downloader_bot.settings import settings
 
 uvloop.install()
 
-if settings.app_env == "PROD":
+if settings.app_env == "production":
+    log.info("Starting telemetry server, in production mode.")
     init_telemetry()
+else:
+    log.info("Not starting telemetry server, not in production mode.")
 
 app = Client("downloader_bot",
              api_id=settings.api_id,
@@ -56,9 +59,9 @@ async def message_handler(_, message: Message):
         msg = f"Downloading video {i+1}/{len(urls)}..."
         log.info(msg)
         await message.reply_text(msg)
-        await utils.download_tt_video(url)
+        status = await utils.download_tt_video(url)
 
-    await message.reply_text("Done.")
+    await message.reply_text(f"Done. {status}")
 
 
 @app.on_message(filters.media)

telegram_downloader_bot/telemetry.py

@@ -1,10 +1,32 @@
-import sentry_sdk
+from prometheus_client import Histogram, start_http_server
+
+DOWNLOAD_DURATION = Histogram(
+    'tt_download_time',
+    'Time taken to download a single tiktok video',
+    ['service'],
+    buckets=[0.1, 0.5, 1, 2, 5, 10, 30, 60, 120]
+)
+
+FILE_SIZE_BYTES = Histogram(
+    'downloaded_file_size_bytes',
+    'Size of the downloaded file in bytes',
+    ['service'],
+    buckets=[
+        1e6,    # 1 MB
+        5e6,    # 5 MB
+        10e6,   # 10 MB
+        25e6,   # 25 MB
+        50e6,   # 50 MB
+        100e6,  # 100 MB
+        200e6,  # 200 MB
+        500e6,  # 500 MB
+        1e9     # 1 GB
+    ]
+)
 
 
 def init_telemetry() -> None:
-    sentry_sdk.init(
-        dsn="https://12d7a075d483fc133cde0ed82e72ac45@o4508071875313664.ingest.de.sentry.io/4508075566694480",  # noqa: E501
-        traces_sample_rate=1.0,
-        profiles_sample_rate=1.0,
-        enable_tracing=True
-    )
+    """
+    Initialize telemetry for the bot.
+    """
+    start_http_server(8000)

telegram_downloader_bot/utils.py

@@ -8,7 +8,10 @@ from pyrogram import Client
 from pyrogram.types import Message
 from tiktok_downloader import snaptik
 
+from telegram_downloader_bot.logger import log
 from telegram_downloader_bot.settings import settings
+from telegram_downloader_bot.telemetry import DOWNLOAD_DURATION
+from telegram_downloader_bot.telemetry import FILE_SIZE_BYTES
 
 
 def sanitize_name(input: str) -> str:
@@ -56,6 +59,7 @@ def get_user_folder(message: Message) -> os.path:
     return user_folder
 
 
+@DOWNLOAD_DURATION.labels(service='telegram').time()
 async def handle_media_message_contents(client: Client,
                                         message: Message):
     """Detect what kind of media is being sent over from the user.
@@ -96,36 +100,43 @@ async def handle_media_message_contents(client: Client,
     else:
         await message.reply_text("Unknown media type!")
 
+    size = os.path.getsize(file_path)
+    FILE_SIZE_BYTES.labels(service="telegram").observe(size)
 
-def get_tt_hashes() -> set:
+
+async def get_tt_hashes() -> set:
     if not os.path.exists(settings.tt_hash_file):
         return set()
 
     with open(settings.tt_hash_file, "rb+") as f:
         all_tt_hashes: set = pickle.load(f)  # nosec
-    print(all_tt_hashes)
+
     return all_tt_hashes
 
 
-def add_to_hashes(new_hash: str) -> None:
-    all_tt_hashes = get_tt_hashes()
+async def add_to_hashes(new_hash: str) -> None:
+    all_tt_hashes = await get_tt_hashes()
     all_tt_hashes.add(new_hash)
-    save_tt_hashes(all_tt_hashes)
+    await save_tt_hashes(all_tt_hashes)
 
 
-def save_tt_hashes(hashes: set) -> None:
+async def save_tt_hashes(hashes: set) -> None:
     with open(settings.tt_hash_file, "wb+") as f:
         pickle.dump(hashes,
                     f,
                     protocol=pickle.HIGHEST_PROTOCOL)
 
 
-def check_if_tt_downloaded(tt_hash: str) -> bool:
-    all_tt_hashes = get_tt_hashes()
+async def check_if_tt_downloaded(tt_hash: str) -> bool:
+    all_tt_hashes = await get_tt_hashes()
+    log.info(f"All hashes: {all_tt_hashes}")
+    log.info(f"Hash to check: {tt_hash}")
+    log.info(f"Hash exists: {tt_hash in all_tt_hashes}")
     return tt_hash in all_tt_hashes
 
 
-def download_tt_video(url: str) -> str:
+@DOWNLOAD_DURATION.labels(service='tiktok').time()
+async def download_tt_video(url: str) -> str:
     """Downloads tiktok video from a given URL.
     Makes sure the video integrity is correct."""
 
@@ -134,24 +145,31 @@ def download_tt_video(url: str) -> str:
 
     for video in videos:
         video_filename = now.strftime("video-tiktok-%Y-%m-%d_%H-%M-%S.mp4")
-        video_filepath: os.path = os.path.join(settings.storage,
-                                               "tiktok",
-                                               video_filename)
+        video_filepath = os.path.join(settings.storage,
+                                      "tiktok",
+                                      video_filename)
         video_content = video.download().getbuffer()
-
         video_hash = sha256(video_content).hexdigest()
-        print(video_hash)
 
-        if check_if_tt_downloaded(video_hash):
+        log.info(f"{video_hash}")
+        log.info(f"{video_filepath}")
+        print(video_filepath)
+
+        if await check_if_tt_downloaded(video_hash) is True:
             return "Already downloaded"
 
         with open(video_filepath, "wb") as f:
             f.write(video_content)
 
-        add_to_hashes(video_hash)
+        await add_to_hashes(video_hash)
+
+        size = os.path.getsize(video_filepath)
+        FILE_SIZE_BYTES.labels(service="tiktok").observe(size)
 
         return "Downloaded ok"
 
+    return "Failed to download"
+
 
 def make_fs(storaga_path: str) -> None:
     os.makedirs(os.path.join(storaga_path, "tiktok"), exist_ok=True)