Changes in Database

FileStream/Database/Elements.py

@@ -0,0 +1,54 @@
+
+
+def NewTG_Files(details:dict) -> dict:
+    return {
+        "user_id": details['user_id'] if details['user_id'] else None,
+        "user_type": details['user_type'] if details['user_type'] else None,
+        "message_id": details['message_id'] if details['message_id'] else None,
+        "location":details['location'] if details['location'] else None,
+        "IMDB_id": "Unknown",
+        "poster": "Unknown",   
+        "title": "Unknown",
+        "type": "Unknown",
+        "description":"Unknown",
+        "genre":"Unknown",
+        "quality":"Unknown",
+        "time": details['time'] if details['time'] else None,
+        "privacy_type": details['privacy_type'] if details['privacy_type'] else None,
+        "file_ids":{},
+        "file": {
+            "file_id":details['file']['file_id'] if details['file']['file_id'] else None,
+            "caption":"",
+            "file_unique_id": details['file']['file_unique_id'] if details['file']['file_unique_id'] else None,
+            "file_name":details['file']['file_name'] if details['file']['file_name'] else None,
+            "file_size":details['file']['file_size'] if details['file']['file_size'] else None,
+            "mime_type":details['file']['mime_type']if details['file']['mime_type'] else None,
+            "taged_users": {}
+        },
+    }
+
+
+def UserSchema(id: str) -> dict:
+
+    return dict(
+        telegram_id=id,
+        access="USER",
+        tele_status={
+            "status": "ACTIVE",
+            "activity": None,
+            "joined": Time_ISTKolNow()
+        },
+        file={
+            "links": 0,
+            "private_files": 0,
+            "public_files": 0,
+        },
+        site_id="None",
+        site_status={
+            "status": None,
+            "activity": None,
+            "password": None,
+            "links": 0,
+            "joined": "None"
+        },
+    )

FileStream/Database/__init__.py

@@ -1 +1,4 @@
 from .database import Database
+from FileStream.config import Telegram
+
+DB_Controller = Database(Telegram.DATABASE_URL, Telegram.SESSION_NAME)

FileStream/Database/database.py

@@ -10,6 +10,8 @@ from bson.json_util import dumps
 from FileStream.Exceptions import FileNotFound
 from FileStream.Tools import Time_ISTKolNow
 
+from .Elements import UserSchema, NewTG_Files
+
 
 class Database:
 
@@ -21,67 +23,13 @@ class Database:
     self.files = self.db.Public_Files
     self.pfile = self.db.Private_Files
     self.web_upload = self.db.Web_Files
-    
 
 #---------------------[ SCHEMAS ]------------------------------#
 #---------------------[ NEW USER ]---------------------#
-
-  def new_user(self, id):
-    return dict(
-        telegram_id=id,
-        access="USER",
-        tele_status={
-            "status": "ACTIVE",
-            "activity": None,
-            "joined": Time_ISTKolNow()
-        },
-        file={
-            "links": 0,
-            "private_files": 0,
-            "public_files": 0,
-        },
-        site_id="None",
-        site_status={
-            "status": None,
-            "activity": None,
-            "password": None,
-            "links": 0,
-            "joined": "None"
-        },
-    )
-#------------------------------------------------
-
-  def NewTG_Files(self, details):
-    return {
-        "user_id": details['user_id'] if details['user_id'] else None,
-        "user_type": details['user_type'] if details['user_type'] else None,
-        "message_id": details['message_id'] if details['message_id'] else None,
-        "location":details['location'] if details['location'] else None,
-        "IMDB_id": "Unknown",
-        "poster": "Unknown",   
-        "title": "Unknown",
-        "type": "Unknown",
-        "description":"Unknown",
-        "genre":"Unknown",
-        "quality":"Unknown",
-        "time": details['time'] if details['time'] else None,
-        "privacy_type": details['privacy_type'] if details['privacy_type'] else None,
-        "file_ids":{},
-        "file": {
-            "file_id":details['file']['file_id'] if details['file']['file_id'] else None,
-            "caption":"",
-            "file_unique_id": details['file']['file_unique_id'] if details['file']['file_unique_id'] else None,
-            "file_name":details['file']['file_name'] if details['file']['file_name'] else None,
-            "file_size":details['file']['file_size'] if details['file']['file_size'] else None,
-            "mime_type":details['file']['mime_type']if details['file']['mime_type'] else None,
-            "taged_users": {}
-        },
-    }
-
 # ---------------------[ ADD USER ]---------------------#
 
   async def add_user(self, id):
-    user = self.new_user(id)
+    user = UserSchema(id)
     await self.users.insert_one(user)
 
   async def add_admin(self, id):
@@ -89,7 +37,7 @@ class Database:
     if user:
       await self.users.update_one({"_id": user['_id']}, {"$set": {"access":"ADMIN" }})
     else:
-      user = self.new_user(id)
+      user = UserSchema(id)
       user['access']="ADMIN"
       await self.users.insert_one(user)
 
@@ -147,13 +95,25 @@ class Database:
       return False
 
   # ---------------------[ ADD FILE TO DB ]---------------------#
-  async def add_file(self, file_info):
+  async def add_file(self, file_info:dict, db_type:str):
     file_info["time"] = Time_ISTKolNow()
-    fetch_old = await self.get_file_by_fileuniqueid(file_info["user_id"], file_info['file']["file_unique_id"])
-    if fetch_old:
-      return fetch_old["_id"]
-    await self.count_links(file_info["user_id"], "+")
-    return (await self.files.insert_one(file_info)).inserted_id
+    if db_type =="PUBLIC":
+      fetch_old = await self.get_file_by_fileuniqueid(file_info["user_id"], file_info['file']["file_unique_id"],  db_type ="PUBLIC")
+      if fetch_old:
+        return fetch_old["_id"]
+      await self.count_links(file_info["user_id"], "+")
+      return (await self.files.insert_one(file_info)).inserted_id
+    elif db_type=="PRIVATE":
+      pass
+    elif db_type == "TEMPORARY":
+      fetch_old = await self.get_file_by_fileuniqueid(file_info["user_id"], file_info['file']["file_unique_id"], db_type ="TEMPORARY")
+      if fetch_old:
+        return fetch_old["_id"]
+      await self.count_links(file_info["user_id"], "+")
+      return (await self.tfiles.insert_one(file_info)).inserted_id
+    else:
+      return {}
+    
 
   # ---------------------[ ADD FILE TO Temp DB ]---------------------#
   async def add_temp_file(self, file_info):
@@ -251,18 +211,38 @@ class Database:
 
   async def get_file_by_fileuniqueid_only(self, file_unique_id, privacy_type:str):
     if privacy_type=="TEMPORARY":
+      return await self.tfiles.find_one({"file.file_unique_id": file_unique_id})
+    elif privacy_type=="PUBLIC":
       return await self.files.find_one({"file.file_unique_id": file_unique_id})
+    elif privacy_type=="PRIVATE":
+      return await self.pfiles.find_one({"file.file_unique_id": file_unique_id})
     else:
-      return await self.files.find_one({"file.file_unique_id": file_unique_id})
+      return {}
 
-  async def get_file_by_fileuniqueid(self, id, file_unique_id):
-    count = await self.files.count_documents({"user_id":id,"file.file_unique_id":file_unique_id})
-    if count == 0:
-      return False
-    elif count == 1:
-      return await self.files.find_one({"user_id": id,"file.file_unique_id": file_unique_id})
+  async def get_file_by_fileuniqueid(self, id, file_unique_id, privacy_type:str):
+    if privacy_type=="TEMPORARY":
+      count = await self.tfiles.count_documents({"user_id":id,"file.file_unique_id":file_unique_id})
+      if count == 0:
+        return False
+      elif count == 1:
+        return await self.tfiles.find_one({"user_id": id,"file.file_unique_id": file_unique_id})
+      else:
+        return self.tfiles.find({"user_id": id,"file.file_unique_id": file_unique_id})
+
+    elif privacy_type=="PUBLIC":
+      count = await self.files.count_documents({"user_id":id,"file.file_unique_id":file_unique_id})
+      if count == 0:
+        return False
+      elif count == 1:
+        return await self.files.find_one({"user_id": id,"file.file_unique_id": file_unique_id})
+      else:
+        return self.files.find({"user_id": id,"file.file_unique_id": file_unique_id})
+
+    elif privacy_type=="PRIVATE":
+      return await self.pfiles.find_one({"file.file_unique_id": file_unique_id})
     else:
-      return self.files.find({"user_id": id,"file.file_unique_id": file_unique_id})
+      return {}
+    
 
   # ---------------------[ UPDATE FILE IN DB ]---------------------#
 

FileStream/bot/plugins/FileHandlers/callback.py

@@ -151,7 +151,7 @@ async def cb_data(bot: Client, update: CallbackQuery):
       
       file_info = get_file_info(message, instruction)
       # Here we are Adding the File Into the database First
-      inserted_id = await db.add_file(file_info)
+      inserted_id = await db.add_file(file_info=file_info, db_type="PUBLIC")
       await get_file_ids(False, inserted_id, message)
       #All the Time Get_file_ids should be called before update privacy or else tagged_users will be {}
       await db.update_privacy(file_info)
@@ -214,7 +214,7 @@ async def cb_data(bot: Client, update: CallbackQuery):
       print("Line No 212 ")
       file_info = get_file_info(message, instruction)
       # Here we are Adding the File Into the database First
-      inserted_id = await db.add_temp_file(file_info)
+      inserted_id = await db.add_file(file_info=file_info, db_type="TEMPORARY")
       await get_file_ids(False, inserted_id, message)
       #All the Time Get_file_ids should be called before update privacy or else tagged_users will be {}
       await db.update_privacy(file_info)
@@ -257,7 +257,7 @@ async def cb_data(bot: Client, update: CallbackQuery):
       }
       file_info = get_file_info(message, instruction)
       # Here we are Adding the File Into the database First
-      db_id = await db.add_file(file_info)
+      db_id = await db.add_file(file_info=file_info, db_type="PRIVATE")
       await get_file_ids(False, db_id, message)
       if True:
         file_info = await db.get_file(db_id)

FileStream/bot/plugins/FileHandlers/stream.py

@@ -100,7 +100,7 @@ async def channel_receive_handler(bot: Client, message: Message):
     await is_channel_exist(bot, message)
 
     try:
-        inserted_id = await db.add_file(get_file_info(message))
+        inserted_id = await db.add_file(file_info=get_file_info(message),  db_type="TEMPORARY")
         await get_file_ids(False, inserted_id, MULTI_CLIENTS, message)
         reply_markup, stream_link = await gen_link(_id=inserted_id)
         await bot.edit_message_reply_markup(

FileStream/utils/FileProcessors/custom_ul.py

@@ -225,7 +225,7 @@ class TeleUploader:
         
         # Here we are Adding the File Into the database First
         #await db.add_webfile(file_info)
-        inserted_id = await db.add_file(file_info)
+        inserted_id = await db.add_file(file_info=file_info,db_type="TEMPORARY")
         await get_file_ids(False, inserted_id, MessageFile)
         reply_markup, stream_text = await gen_link(_id=inserted_id)
         await message.edit_text(