task added

.gitignore

@@ -124,6 +124,7 @@ dmypy.json
 .pyre/
 .idea
 
-gramup/anon.session
-anon.session
+gramup/*.session
+gramup/*.db
+*.session
 gramup/anon.session-journal

gramup/main.py

@@ -1,9 +1,10 @@
 import asyncio
 import glob
 import os
-from typing import List, Literal
+from typing import List
 
-from fastapi import FastAPI, HTTPException, Request
+import shortuuid
+from fastapi import FastAPI, HTTPException, Request, BackgroundTasks
 from starlette import status
 from starlette.middleware.cors import CORSMiddleware
 from starlette.staticfiles import StaticFiles
@@ -12,6 +13,7 @@
 
 from constants import API_ID, API_HASH
 from gramup.models import File, Task, TaskRequest
+from gramup.tasks import pull_all_to_db
 
 client = TelegramClient('anon', API_ID, API_HASH)
 app = FastAPI()
@@ -52,15 +54,15 @@ async def websocket_endpoint(websocket: WebSocket):
 
 
 @app.get("/api/files/")
-async def files(path="") -> List[File]:
+async def files(path="") -> List:
     if not client.is_connected() or not await client.is_user_authorized():
         raise HTTPException(status_code=status.HTTP_401_UNAUTHORIZED, detail="User not authorized")
 
     return []
 
 
 @app.get("/api/local_files/")
-async def local_files(path="") -> List[File]:
+async def local_files(path="") -> List:
     return [
         File(
             folder=os.path.isdir(file),
@@ -73,14 +75,40 @@ async def local_files(path="") -> List[File]:
 
 
 @app.post("/api/tasks/")
-async def tasks(request: Request) -> Task:
+async def tasks(request: Request, background_tasks: BackgroundTasks):
     body = TaskRequest(**(await request.json()))
 
+    # background_tasks.add_task(
+
     return Task(
         id="",
         name="",
         status="running",
     )
 
 
+@app.get("/api/tasks/")
+async def tasks_list():
+    ret = []
+
+    for task in Task.select().order_by(Task.schedule_time.desc()):
+        ret.append({
+            "id": task.id,
+            "name": task.name,
+            "status": task.status,
+            "schedule_time": task.schedule_time
+        })
+
+    return ret
+
+
+@app.get("/api/pull_all/")
+async def pull_all(background_tasks: BackgroundTasks):
+    task, _ = Task.get_or_create(status="running", name="Pull all files")
+    task.save()
+    background_tasks.add_task(pull_all_to_db, task=task)
+
+    return task.id
+
+
 app.mount("/", StaticFiles(directory="static"), name="static")

gramup/models.py

@@ -1,19 +1,41 @@
+from datetime import datetime
 from typing import Literal, TypedDict
 
-from pydantic import BaseModel
+import peewee
+import shortuuid
+
+db = peewee.SqliteDatabase('database.db')
+db.autoconnect = True
+
+
+class BaseModel(peewee.Model):
+    id = peewee.CharField(primary_key=True,max_length=128, default=lambda: shortuuid.uuid(pad_length=128))
+
+    class Meta:
+        database = db
 
 
 class File(BaseModel):
-    folder: bool
-    name: str
-    path: str
-    id: str
+    name = peewee.CharField(max_length=128)
+    path = peewee.CharField(max_length=128)
+    folder = peewee.BooleanField(default=False)
+
+
+class Folder(BaseModel):
+    this = peewee.ForeignKeyField(File, backref="this_folder", on_delete="CASCADE")
+    sub_folders = peewee.ManyToManyField(File, backref="parent_folder", on_delete="CASCADE")
+    files = peewee.ManyToManyField(File, backref="folder", on_delete="CASCADE")
 
 
 class Task(BaseModel):
-    id: str
-    name: str
-    status: Literal["running", "stopped", "finished"]
+    name = peewee.CharField(max_length=128)
+    status = peewee.CharField(choices=["running", "completed", "failed"], default="running")
+    schedule_time = peewee.DateTimeField(default=datetime.now)
+
+
+FolderSubFolder = Folder.sub_folders.get_through_model()
+FolderFile = Folder.files.get_through_model()
 
+db.create_tables([File, FolderSubFolder, FolderFile, Folder, Task])
 
 TaskRequest = TypedDict("TaskRequest", {"path": str, "type": Literal["sync", "upload"]})

gramup/tasks.py

@@ -0,0 +1,50 @@
+from pathlib import Path
+
+from telethon import TelegramClient
+from telethon.tl.types import InputMessagesFilterDocument
+
+from gramup.models import File, Folder, Task
+
+
+async def pull_all_to_db(task: Task):
+    client = TelegramClient("session", 1234, "asdf")
+
+    if not client.is_connected():
+        await client.connect()
+
+    if not await client.is_user_authorized():
+        task.status = "failed"
+        task.save()
+        return
+
+    for message in client.iter_messages("me", filter=InputMessagesFilterDocument):
+        m_id = message.id
+        caption: str = message.caption
+        path = Path(caption if "/" in caption else f"external/{caption}")
+
+        file = File.create(
+            id=m_id,
+            path=str(path),
+            name=path.name
+        )
+
+        file.save()
+        first = True
+
+        for folder in path.parents:
+            folder_new, _ = Folder.get_or_create(
+                path=str(folder),
+                name=folder.name
+            )
+
+            if first:
+                first = False
+                folder_new.files.add(file)
+            else:
+                folder_new.sub_folders.add(file)
+
+            folder_new.save()
+            file = folder_new
+
+    task.status = "completed"
+    task.save()

requirements.txt

@@ -1,5 +1,7 @@
 fastapi==0.96.0
 starlette~=0.27.0
 Telethon~=1.28.5
-
 pydantic~=1.10.8
+shortuuid~=1.0.11
+
+peewee~=3.16.2