From cafc9eb8a80d017af97c9b7224bd0f00edd30f1f Mon Sep 17 00:00:00 2001
From: Saurab-Shrestha <shresthasaurab030@gmail.com>
Date: Sun, 11 Feb 2024 09:15:32 +0545
Subject: [PATCH] Minor changes in delete filename

---
 private_gpt/server/ingest/ingest_router.py | 12 +++++++++++-
 private_gpt/users/db/session.py            |  7 ++++++-
 private_gpt/users/models/documents.py      | 13 +++++++------
 private_gpt/users/schemas/__init__.py      |  1 +
 4 files changed, 25 insertions(+), 8 deletions(-)

diff --git a/private_gpt/server/ingest/ingest_router.py b/private_gpt/server/ingest/ingest_router.py
index dab59cae..c591acbd 100644
--- a/private_gpt/server/ingest/ingest_router.py
+++ b/private_gpt/server/ingest/ingest_router.py
@@ -10,6 +10,7 @@ from pydantic import BaseModel, Field
 from private_gpt.home import Home
 from private_gpt.users import crud, models, schemas
 from private_gpt.users.api import deps
+from private_gpt.users.constants.role import Role
 
 from private_gpt.server.ingest.ingest_service import IngestService
 from private_gpt.server.ingest.model import IngestedDoc
@@ -128,8 +129,11 @@ def delete_ingested(request: Request, doc_id: str) -> None:
 def delete_file(
         request: Request,
         delete_input: DeleteFilename,
+        db: Session = Depends(deps.get_db),
         current_user: models.User = Security(
             deps.get_current_user,
+            scopes=[Role.ADMIN["name"], Role.SUPER_ADMIN["name"]],
+
         )) -> dict:
     """Delete the specified filename.
 
@@ -146,7 +150,8 @@ def delete_file(
 
         for doc_id in doc_ids:
             service.delete(doc_id)
-
+        document = crud.documents.get_by_filename(db,file_name=filename)
+        crud.documents.remove(db=db, id=document.id)
         return {"status": "SUCCESS", "message": f"{filename}' successfully deleted."}
     except Exception as e:
         logger.error(
@@ -162,6 +167,7 @@ def ingest_file(
         file: UploadFile = File(...),
         current_user: models.User = Security(
             deps.get_current_user,
+            scopes=[Role.ADMIN["name"], Role.SUPER_ADMIN["name"]],
         )) -> IngestResponse:
     """Ingests and processes a file, storing its chunks to be used as context."""
     service = request.state.injector.get(IngestService)
@@ -198,8 +204,12 @@ def ingest_file(
         logger.info(f"{file.filename} is uploaded by the {current_user.fullname}.")
 
         return IngestResponse(object="list", model="private-gpt", data=ingested_documents)
+    except HTTPException:
+        raise
+
     except Exception as e:
         logger.error(f"There was an error uploading the file(s): {str(e)}")
+        print("ERROR: ", e)
         raise HTTPException(
             status_code=status.HTTP_500_INTERNAL_SERVER_ERROR,
             detail="Internal Server Error: Unable to ingest file.",
diff --git a/private_gpt/users/db/session.py b/private_gpt/users/db/session.py
index 259a8614..383da320 100644
--- a/private_gpt/users/db/session.py
+++ b/private_gpt/users/db/session.py
@@ -1,6 +1,11 @@
 from private_gpt.users.core.config import SQLALCHEMY_DATABASE_URI
 from sqlalchemy import create_engine
 from sqlalchemy.orm import sessionmaker
+import logging
 
-engine = create_engine(SQLALCHEMY_DATABASE_URI, echo=True, future=True, pool_pre_ping=True)
+logging.basicConfig()
+logging.getLogger("sqlalchemy.engine").setLevel(logging.INFO)
+logging.getLogger("sqlalchemy.pool").setLevel(logging.DEBUG)
+engine = create_engine(SQLALCHEMY_DATABASE_URI, echo=True,
+                       future=True, pool_pre_ping=True, logging_name="myengine")
 SessionLocal = sessionmaker(autocommit=False, autoflush=False, bind=engine)
diff --git a/private_gpt/users/models/documents.py b/private_gpt/users/models/documents.py
index e0476c36..5970fd55 100644
--- a/private_gpt/users/models/documents.py
+++ b/private_gpt/users/models/documents.py
@@ -1,22 +1,23 @@
 from private_gpt.users.db.base_class import Base
-from datetime import datetime, timedelta
+from datetime import datetime
 from sqlalchemy.orm import relationship
-from sqlalchemy import Column, Integer, String, Boolean, Float, ForeignKey, DateTime
+from sqlalchemy import Column, Integer, String, ForeignKey, DateTime
 
 
 class Documents(Base):
-    """Models a user table"""
-    _tablename_ = "document"
+    """Models a document table"""
+    __tablename__ = "document"  
     id = Column(Integer, primary_key=True, index=True)
     filename = Column(String(225), nullable=False, unique=True)
     uploaded_by = Column(
         Integer,
         ForeignKey("users.id"),
-        primary_key=True,
         nullable=False,
     )
     uploaded_at = Column(
         DateTime,
         default=datetime.utcnow,
+        nullable=False,
     )
-    uploaded_by_user = relationship("User", back_populates="uploaded_documents")
\ No newline at end of file
+    uploaded_by_user = relationship(
+        "User", back_populates="uploaded_documents")
diff --git a/private_gpt/users/schemas/__init__.py b/private_gpt/users/schemas/__init__.py
index cc14f1d1..6763d498 100644
--- a/private_gpt/users/schemas/__init__.py
+++ b/private_gpt/users/schemas/__init__.py
@@ -4,3 +4,4 @@ from .user import User, UserCreate, UserInDB, UserUpdate, UserBaseSchema, Profil
 from .user_role import UserRole, UserRoleCreate, UserRoleInDB, UserRoleUpdate
 from .subscription import Subscription, SubscriptionBase, SubscriptionCreate, SubscriptionUpdate
 from .company import Company, CompanyBase, CompanyCreate, CompanyUpdate
+from .documents import Document, DocumentCreate, DocumentsBase, DocumentUpdate
\ No newline at end of file