Gitlab support for issues, tickets GUI, UI improvements

Author: Roey7

app/data_source/api/basic_document.py

@@ -9,10 +9,16 @@ class DocumentType(Enum):
     MESSAGE = "message"
     COMMENT = "comment"
     PERSON = "person"
-    GIT_ISSUE = "git_issue"
+    ISSUE = "issue"
     GIT_PR = "git_pr"
 
 
+class DocumentStatus(Enum):
+    OPEN = "open"
+    IN_PROGRESS = "in_progress"
+    CLOSED = "closed"
+
+
 class FileType(Enum):
     GOOGLE_DOC = "doc"
     DOCX = "docx"
@@ -35,8 +41,8 @@ def from_mime_type(cls, mime_type: str):
 
 @dataclass
 class BasicDocument:
-    id: Union[int, str]
-    data_source_id: int
+    id: Union[int, str]  # row id in database
+    data_source_id: int  # data source id in database
     type: DocumentType
     title: str
     content: str
@@ -46,6 +52,7 @@ class BasicDocument:
     location: str
     url: str
     file_type: FileType = None
+    status: DocumentStatus = None
 
     @property
     def id_in_data_source(self):

app/data_source/sources/gitlab/__init__.py

@@ -0,0 +1,112 @@
+import logging
+from datetime import datetime
+
+import requests
+from typing import Dict, List, Optional
+
+from data_source.api.base_data_source import BaseDataSource, BaseDataSourceConfig, ConfigField, HTMLInputType
+from data_source.api.basic_document import BasicDocument, DocumentType, DocumentStatus
+from data_source.api.exception import InvalidDataSourceConfig
+from queues.index_queue import IndexQueue
+
+GITLAB_BASE_URL = "https://gitlab.com/api/v4"
+
+logger = logging.getLogger(__name__)
+
+
+class GitlabConfig(BaseDataSourceConfig):
+    access_token: str
+
+
+def gitlab_status_to_doc_status(status: str) -> Optional[DocumentStatus]:
+    if status == "opened":
+        return DocumentStatus.OPEN
+    elif status == "closed":
+        return DocumentStatus.CLOSED
+    else:
+        logger.warning(f"[!] Unknown status {status}")
+        return None
+
+
+class GitlabDataSource(BaseDataSource):
+
+    def _feed_project_issues(self, project: Dict):
+        project_id = project["id"]
+        project_url = project["web_url"]
+
+        issues_url = f"{GITLAB_BASE_URL}/projects/{project_id}/issues"
+
+        issues_response = self._session.get(issues_url)
+        issues_response.raise_for_status()
+        issues_json = issues_response.json()
+
+        for issue in issues_json:
+            last_modified = datetime.strptime(issue["updated_at"], "%Y-%m-%dT%H:%M:%S.%fZ")
+            if last_modified < self._last_index_time:
+                logger.info(f"Issue {issue['id']} is too old, skipping")
+                continue
+
+            doc = BasicDocument(
+                id=issue["id"],
+                data_source_id=self._data_source_id,
+                type=DocumentType.ISSUE,
+                title=issue['title'],
+                content=issue["description"] if not None else "",
+                author=issue['author']['name'],
+                author_image_url=issue['author']['avatar_url'],
+                location=issue['references']['full'].replace("/", " / "),
+                url=issue['web_url'],
+                timestamp=last_modified,
+                status=gitlab_status_to_doc_status(issue["state"])
+            )
+            IndexQueue.get_instance().put_single(doc=doc)
+
+    @staticmethod
+    def get_config_fields() -> List[ConfigField]:
+        return [
+            ConfigField(label="API Access Token", name="access_token", input_type=HTMLInputType.PASSWORD),
+        ]
+
+    @staticmethod
+    def validate_config(config: Dict) -> None:
+        try:
+            parsed_config = GitlabConfig(**config)
+            session = requests.Session()
+            session.headers.update({"PRIVATE-TOKEN": parsed_config.access_token})
+            projects_response = session.get(f"{GITLAB_BASE_URL}/projects?membership=true")
+            projects_response.raise_for_status()
+        except (KeyError, ValueError) as e:
+            raise InvalidDataSourceConfig from e
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.gitlab_config = GitlabConfig(**self._raw_config)
+        self._session = requests.Session()
+        self._session.headers.update({"PRIVATE-TOKEN": self.gitlab_config.access_token})
+
+    def _list_all_projects(self) -> List[Dict]:
+        projects = []
+        page = 1
+        per_page = 100
+
+        while True:
+            try:
+                projects_response = self._session.get(f"{GITLAB_BASE_URL}/projects?membership=true"
+                                                      f"&per_page={per_page}&page={page}")
+                projects_response.raise_for_status()
+                new_projects: List[Dict] = projects_response.json()
+                projects.extend(new_projects)
+
+                if len(new_projects) < per_page:
+                    break
+
+                page += 1
+            except:
+                logging.exception("Error while fetching projects")
+
+        return projects
+
+    def _feed_new_documents(self) -> None:
+        for project in self._list_all_projects():
+            logger.info(f"Feeding project {project['name']}")
+            self.add_task_to_queue(self._feed_project_issues, project=project)

app/data_source/sources/gitlab/gitlab.py

@@ -180,11 +180,14 @@ def _feed_file(self, file):
             author = first_owner.get('displayName')
             author_image_url = first_owner.get('photoLink')
 
+        # title is file name without extension
+        title = file['name'].split('.')[0]
+
         doc = BasicDocument(
             id=file_id,
             data_source_id=self._data_source_id,
             type=DocumentType.DOCUMENT,
-            title=file['name'],
+            title=title,
             content=content,
             author=author,
             author_image_url=author_image_url,

app/data_source/sources/google_drive/google_drive.py

@@ -126,7 +126,7 @@ def _feed_conversation(self, conv: SlackConversation):
             message_id = message['client_msg_id']
             readable_timestamp = datetime.datetime.fromtimestamp(float(timestamp))
             message_url = f"https://slack.com/app_redirect?channel={conv.id}&message_ts={timestamp}"
-            last_msg = BasicDocument(title=conv.name, content=text, author=author.name,
+            last_msg = BasicDocument(title=author.name, content=text, author=author.name,
                                      timestamp=readable_timestamp, id=message_id,
                                      data_source_id=self._data_source_id, location=conv.name,
                                      url=message_url, author_image_url=author.image_url,

app/data_source/sources/slack/slack.py

@@ -1,19 +1,26 @@
 import logging
 import re
-from typing import List
+from enum import Enum
+from typing import List, Optional
 
 from data_source.api.basic_document import BasicDocument
+from db_engine import Session
+from indexing.bm25_index import Bm25Index
+from indexing.faiss_index import FaissIndex
+from models import bi_encoder
 from paths import IS_IN_DOCKER
 from schemas import Document, Paragraph
-from models import bi_encoder
-from indexing.faiss_index import FaissIndex
-from indexing.bm25_index import Bm25Index
-from db_engine import Session
-
 
 logger = logging.getLogger(__name__)
 
 
+def get_enum_value_or_none(enum: Optional[Enum]) -> Optional[str]:
+    if enum is None:
+        return None
+
+    return enum.value
+
+
 class Indexer:
 
     @staticmethod
@@ -23,7 +30,8 @@ def index_documents(documents: List[BasicDocument]):
         ids_in_data_source = [document.id_in_data_source for document in documents]
 
         with Session() as session:
-            documents_to_delete = session.query(Document).filter(Document.id_in_data_source.in_(ids_in_data_source)).all()
+            documents_to_delete = session.query(Document).filter(
+                Document.id_in_data_source.in_(ids_in_data_source)).all()
             if documents_to_delete:
                 logging.info(f'removing documents that were updated and need to be re-indexed.')
                 Indexer.remove_documents(documents_to_delete, session)
@@ -43,7 +51,8 @@ def index_documents(documents: List[BasicDocument]):
                     data_source_id=document.data_source_id,
                     id_in_data_source=document.id_in_data_source,
                     type=document.type.value,
-                    file_type=document.file_type.value if document.file_type is not None else None,
+                    file_type=get_enum_value_or_none(document.file_type),
+                    status=get_enum_value_or_none(document.status),
                     title=document.title,
                     author=document.author,
                     author_image_url=document.author_image_url,
@@ -120,7 +129,7 @@ def _add_metadata_for_indexing(paragraph: Paragraph) -> str:
         return result
 
     @staticmethod
-    def remove_documents(documents: List[Document], session = None):
+    def remove_documents(documents: List[Document], session=None):
         logger.info(f"Removing {len(documents)} documents")
 
         # Get the paragraphs from the documents
@@ -135,4 +144,4 @@ def remove_documents(documents: List[Document], session = None):
         logger.info(f"Removing documents from BM25 index...")
         Bm25Index.get().update(session=session)
 
-        logger.info(f"Finished removing {len(documents)} documents => {len(db_paragraphs)} paragraphs")
+        logger.info(f"Finished removing {len(documents)} documents => {len(db_paragraphs)} paragraphs")

app/data_sources/gitlab.py

@@ -1,7 +1,7 @@
 from typing import Optional
 
 from schemas.base import Base
-from sqlalchemy import String, DateTime, ForeignKey, Column, Integer
+from sqlalchemy import String, DateTime, ForeignKey, Column, Integer, Boolean
 from sqlalchemy.orm import Mapped, mapped_column, relationship
 
 
@@ -14,6 +14,7 @@ class Document(Base):
     data_source = relationship("DataSource", back_populates="documents")
     type: Mapped[Optional[str]] = mapped_column(String(32))
     file_type: Mapped[Optional[str]] = mapped_column(String(32))
+    status: Mapped[Optional[bool]] = mapped_column(String(32))
     title: Mapped[Optional[str]] = mapped_column(String(128))
     author: Mapped[Optional[str]] = mapped_column(String(64))
     author_image_url: Mapped[Optional[str]] = mapped_column(String(512))