From 9f1fabe27aa84ac8400b4cdad863583a056f1db8 Mon Sep 17 00:00:00 2001 From: Rajesh Jonnalagadda Date: Tue, 8 Oct 2024 22:20:28 +0530 Subject: [PATCH] feat: add runbook data into db. to do: incident link needs to be done. --- keep-ui/app/runbooks/runbook-table.tsx | 1 + keep-ui/utils/hooks/useRunbook.ts | 27 +++- keep/api/models/db/migrations/env.py | 1 + .../versions/2024-10-08-16-09_8902e1a17f66.py | 55 +++++++ keep/api/models/db/runbook.py | 143 ++++++++---------- keep/api/routes/runbooks.py | 47 +++++- keep/providers/base/base_provider.py | 2 +- .../github_provider/github_provider.py | 24 +-- .../gitlab_provider/gitlab_provider.py | 24 +-- keep/runbooks/runbooks_service.py | 51 +++++++ 10 files changed, 265 insertions(+), 110 deletions(-) create mode 100644 keep/api/models/db/migrations/versions/2024-10-08-16-09_8902e1a17f66.py create mode 100644 keep/runbooks/runbooks_service.py diff --git a/keep-ui/app/runbooks/runbook-table.tsx b/keep-ui/app/runbooks/runbook-table.tsx index cee7353b6..b640663dd 100644 --- a/keep-ui/app/runbooks/runbook-table.tsx +++ b/keep-ui/app/runbooks/runbook-table.tsx @@ -230,6 +230,7 @@ function SettingsPage() { diff --git a/keep-ui/utils/hooks/useRunbook.ts b/keep-ui/utils/hooks/useRunbook.ts index 37e0e1939..14675a860 100644 --- a/keep-ui/utils/hooks/useRunbook.ts +++ b/keep-ui/utils/hooks/useRunbook.ts @@ -14,7 +14,7 @@ export const useRunBookTriggers = (values: any, refresh: number) => { const [synced, setSynced] = useState(false); const [fileData, setFileData] = useState({}); const [reposData, setRepoData] = useState([]); - const { pathToMdFile, repoName, userName, providerId, domain } = values || {}; + const { pathToMdFile, repoName, userName, providerId, domain, runBookTitle } = values || {}; const { data: session } = useSession(); const { installed_providers, providers } = (providersData?.data || {}) as ProvidersResponse; @@ -67,18 +67,31 @@ export const useRunBookTriggers = (values: any, refresh: number) => { if (repoName) { params.append("repo", repoName); } + if(runBookTitle){ + params.append("title", runBookTitle); + + } //TO DO backend runbook records needs to be created. - const response = await fetcher( - `${baseApiurl}/runbooks/${provider?.type}/${ + const response = await fetch(`${baseApiurl}/runbooks/${provider?.type}/${ provider?.id - }/runbook?${params.toString()}`, - session?.accessToken - ); + }?${params.toString()}`, { + method: "POST", + headers: { + "Content-Type": "application/json", + Authorization: `Bearer ${session?.accessToken}`, + }, + }); if (!response) { return setError("Something went wrong. try agian after some time"); } - setFileData(response); + + if(!response.ok) { + return setError("Something went wrong. try agian after some time"); + + } + const result = await response.json(); + setFileData(result); setSynced(false); } catch (err) { return setError("Something went wrong. try agian after some time"); diff --git a/keep/api/models/db/migrations/env.py b/keep/api/models/db/migrations/env.py index 6aab4afad..912a921b9 100644 --- a/keep/api/models/db/migrations/env.py +++ b/keep/api/models/db/migrations/env.py @@ -21,6 +21,7 @@ from keep.api.models.db.topology import * from keep.api.models.db.user import * from keep.api.models.db.workflow import * +from keep.api.models.db.runbook import * target_metadata = SQLModel.metadata diff --git a/keep/api/models/db/migrations/versions/2024-10-08-16-09_8902e1a17f66.py b/keep/api/models/db/migrations/versions/2024-10-08-16-09_8902e1a17f66.py new file mode 100644 index 000000000..8612283ee --- /dev/null +++ b/keep/api/models/db/migrations/versions/2024-10-08-16-09_8902e1a17f66.py @@ -0,0 +1,55 @@ +"""add runbook table + +Revision ID: 8902e1a17f66 +Revises: 01ebe17218c0 +Create Date: 2024-10-08 16:09:28.158034 + +""" + +import sqlalchemy as sa +import sqlalchemy_utils +import sqlmodel +from alembic import op +from sqlalchemy.dialects import postgresql + +# revision identifiers, used by Alembic. +revision = "8902e1a17f66" +down_revision = "01ebe17218c0" +branch_labels = None +depends_on = None + + +def upgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.create_table( + "runbook", + sa.Column("tenant_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("relative_path", sa.Text(), nullable=True), + sa.Column("title", sa.Text(), nullable=True), + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("repo_id", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("provider_type", sqlmodel.sql.sqltypes.AutoString(), nullable=False), + sa.Column("provider_id", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(["tenant_id"], ["tenant.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + op.create_table( + "runbookcontent", + sa.Column("runbook_id", sqlmodel.sql.sqltypes.GUID(), nullable=True), + sa.Column("content", sa.Text(), nullable=True), + sa.Column("link", sa.Text(), nullable=True), + sa.Column("id", sqlmodel.sql.sqltypes.GUID(), nullable=False), + sa.Column("encoding", sqlmodel.sql.sqltypes.AutoString(), nullable=True), + sa.Column("created_at", sa.DateTime(), nullable=False), + sa.ForeignKeyConstraint(["runbook_id"], ["runbook.id"], ondelete="CASCADE"), + sa.PrimaryKeyConstraint("id"), + ) + # ### end Alembic commands ### + + +def downgrade() -> None: + # ### commands auto generated by Alembic - please adjust! ### + op.drop_table("runbookcontent") + op.drop_table("runbook") + # ### end Alembic commands ### diff --git a/keep/api/models/db/runbook.py b/keep/api/models/db/runbook.py index b0f577210..3d32eadae 100644 --- a/keep/api/models/db/runbook.py +++ b/keep/api/models/db/runbook.py @@ -1,93 +1,82 @@ +from uuid import UUID, uuid4 from datetime import datetime from typing import List, Optional -from uuid import UUID, uuid4 - -from sqlalchemy import ForeignKey, Column, TEXT, JSON from sqlmodel import Field, Relationship, SQLModel -from keep.api.models.db.tenant import Tenant +from sqlalchemy import Column, ForeignKey, Text +from pydantic import BaseModel -# Link Model between Runbook and Incident -class RunbookToIncident(SQLModel, table=True): - tenant_id: str = Field(foreign_key="tenant.id") - runbook_id: UUID = Field(foreign_key="runbook.id", primary_key=True) - incident_id: UUID = Field(foreign_key="incident.id", primary_key=True) - incident_id: UUID = Field( - sa_column=Column( - UUID(binary=False), - ForeignKey("incident.id", ondelete="CASCADE"), - primary_key=True, - ) +# RunbookContent Model +class RunbookContent(SQLModel, table=True): + id: UUID = Field(default_factory=uuid4, primary_key=True) + runbook_id: UUID = Field( + sa_column=Column(ForeignKey("runbook.id", ondelete="CASCADE")) # Foreign key with CASCADE delete ) + runbook: Optional["Runbook"] = Relationship(back_populates="contents") + content: str = Field(sa_column=Column(Text), nullable=False) # Using SQLAlchemy's Text type + link: str = Field(sa_column=Column(Text), nullable=False) # Using SQLAlchemy's Text type + encoding: Optional[str] = None + created_at: datetime = Field(default_factory=datetime.utcnow) # Timestamp for creation + class Config: + orm_mode = True # Runbook Model class Runbook(SQLModel, table=True): id: UUID = Field(default_factory=uuid4, primary_key=True) - tenant_id: str = Field(foreign_key="tenant.id") - tenant: Tenant = Relationship() - repo_id: str = Field(nullable=False) # Github repo id - relative_path: str = Field(nullable=False) # Relative path to the .md file - title: str = Field(nullable=False) # Title of the runbook - link: str = Field(nullable=False) # Link to the .md file - - incidents: List["Incident"] = Relationship( - back_populates="runbooks", link_model=RunbookToIncident + tenant_id: str = Field( + sa_column=Column(ForeignKey("tenant.id", ondelete="CASCADE")) # Foreign key with CASCADE delete ) - provider_type: str - provider_id: Optional[str] = None - created_at: datetime = Field(default_factory=datetime.utcnow) + repo_id: str # Repository ID + relative_path: str = Field(sa_column=Column(Text), nullable=False) # Path in the repo, must be set + title: str = Field(sa_column=Column(Text), nullable=False) # Title of the runbook, must be set + contents: List["RunbookContent"] = Relationship(back_populates="runbook") # Relationship to RunbookContent + provider_type: str # Type of the provider + provider_id: Optional[str] = None # Optional provider ID + created_at: datetime = Field(default_factory=datetime.utcnow) # Timestamp for creation class Config: - arbitrary_types_allowed = True - - -# Incident Model -class Incident(SQLModel, table=True): - id: UUID = Field(default_factory=uuid4, primary_key=True) - tenant_id: str = Field(foreign_key="tenant.id") - tenant: Tenant = Relationship() - - user_generated_name: Optional[str] = None - ai_generated_name: Optional[str] = None - - user_summary: Optional[str] = Field(sa_column=Column(TEXT), nullable=True) - generated_summary: Optional[str] = Field(sa_column=Column(TEXT), nullable=True) - - assignee: Optional[str] = None - # severity: int = Field(default=IncidentSeverity.CRITICAL.order) - - creation_time: datetime = Field(default_factory=datetime.utcnow) - - start_time: Optional[datetime] = None - end_time: Optional[datetime] = None - last_seen_time: Optional[datetime] = None - - runbooks: List["Runbook"] = Relationship( - back_populates="incidents", link_model=RunbookToIncident - ) - - is_predicted: bool = Field(default=False) - is_confirmed: bool = Field(default=False) - - alerts_count: int = Field(default=0) - affected_services: List = Field(sa_column=Column(JSON), default_factory=list) - sources: List = Field(sa_column=Column(JSON), default_factory=list) - - rule_id: Optional[UUID] = Field( - sa_column=Column( - UUID(binary=False), - ForeignKey("rule.id", ondelete="CASCADE"), - nullable=True, - ), - ) - - rule_fingerprint: str = Field(default="", sa_column=Column(TEXT)) + orm_mode = True # Enable ORM mode for compatibility with Pydantic models + + +class RunbookDto(BaseModel, extra="ignore"): + id: UUID + tenant_id: str + repo_id: str + relative_path: str + title: str + contents: List["RunbookContent"] = [] + provider_type: str + provider_id: Optional[str] = None - def __init__(self, **kwargs): - super().__init__(**kwargs) - if "runbooks" not in kwargs: - self.runbooks = [] +class RunbookContentDto(BaseModel, extra="ignore"): + id: UUID + content: str + link: str + encoding: Optional[str] = None + + @classmethod + def from_orm(cls, content: "RunbookContent") -> "RunbookContentDto": + return cls( + id=content.id, + content=content.content, + link=content.link, + encoding=content.encoding + ) - class Config: - arbitrary_types_allowed = True +class RunbookDtoOut(RunbookDto): + contents: List[RunbookContentDto] = [] + @classmethod + def from_orm( + cls, runbook: "Runbook" + ) -> "RunbookDtoOut": + return cls( + id=runbook.id, + title=runbook.title, + tenant_id=runbook.tenant_id, + repo_id=runbook.repo_id, + relative_path=runbook.relative_path, + provider_type=runbook.provider_type, + provider_id=runbook.provider_id, + contents=[RunbookContentDto.from_orm(content) for content in runbook.contents] + ) diff --git a/keep/api/routes/runbooks.py b/keep/api/routes/runbooks.py index b94ee3054..82ec5292f 100644 --- a/keep/api/routes/runbooks.py +++ b/keep/api/routes/runbooks.py @@ -7,6 +7,13 @@ from keep.identitymanager.identitymanagerfactory import IdentityManagerFactory from keep.providers.providers_factory import ProvidersFactory from keep.secretmanager.secretmanagerfactory import SecretManagerFactory +from keep.runbooks.runbooks_service import ( + RunbookService + ) +from keep.api.models.db.runbook import ( + RunbookDtoOut + ) + logger = logging.getLogger(__name__) @@ -45,8 +52,8 @@ def get_repositories( return provider.pull_repositories() -@router.get("/{provider_type}/{provider_id}/runbook") -def get_repositories( +@router.get("/{provider_type}/{provider_id}") +def get_runbook( provider_type: str, provider_id: str, authenticated_entity: AuthenticatedEntity = Depends( @@ -56,6 +63,7 @@ def get_repositories( repo: str = Query(None), branch: str = Query(None), md_path: str = Query(None), + title: str = Query(None), ): tenant_id = authenticated_entity.tenant_id logger.info("Getting runbook", extra={"provider_type": provider_type, "provider_id": provider_id}) @@ -76,4 +84,37 @@ def get_repositories( context_manager, provider_id, provider_type, provider_config ) - return provider.pull_runbook(repo=repo, branch=branch, md_path=md_path) + return provider.pull_runbook(repo=repo, branch=branch, md_path=md_path, title=title) + + +@router.post( + "/{provider_type}/{provider_id}", + description="Create a new Runbook", + # response_model=RunbookDtoOut, +) +def create_runbook( + provider_type: str, + provider_id: str, + authenticated_entity: AuthenticatedEntity = Depends( + IdentityManagerFactory.get_auth_verifier(["write:runbook"]) + ), + session: Session = Depends(get_session), + repo: str = Query(None), + branch: str = Query(None), + md_path: str = Query(None), + title: str = Query(None), +): + tenant_id = authenticated_entity.tenant_id + logger.info("Creating Runbook", extra={tenant_id: tenant_id}) + context_manager = ContextManager(tenant_id=tenant_id) + secret_manager = SecretManagerFactory.get_secret_manager(context_manager) + provider_config = secret_manager.read_secret( + f"{tenant_id}_{provider_type}_{provider_id}", is_json=True + ) + provider = ProvidersFactory.get_provider( + context_manager, provider_id, provider_type, provider_config + ) + + runbook_dto= provider.pull_runbook(repo=repo, branch=branch, md_path=md_path, title=title) + return RunbookService.create_runbook(session, tenant_id, runbook_dto) + diff --git a/keep/providers/base/base_provider.py b/keep/providers/base/base_provider.py index 09721b973..c20b6636c 100644 --- a/keep/providers/base/base_provider.py +++ b/keep/providers/base/base_provider.py @@ -666,7 +666,7 @@ def pull_topology(self) -> list[TopologyServiceInDto]: class BaseRunBookProvider(BaseProvider): - def pull_runbook(self, repo=None, branch=None, md_path=None): + def pull_runbook(self, repo=None, branch=None, md_path=None, title=None): raise NotImplementedError("get_runbook() method not implemented") def pull_repositories(self): diff --git a/keep/providers/github_provider/github_provider.py b/keep/providers/github_provider/github_provider.py index f89db02a1..47ccdda61 100644 --- a/keep/providers/github_provider/github_provider.py +++ b/keep/providers/github_provider/github_provider.py @@ -122,13 +122,12 @@ def pull_repositories(self, project_id=None): repos_list = self._format_repos(repos) return repos_list - def _format_runbook(self, runbook, repo): + def _format_runbook(self, runbook, repo, title): """ Format the runbook data into a dictionary. """ - if runbook is None: - return {} - + + # TO DO. currently we are handling the one file only. we user give folder path. then we might get multiple files as input(runbook) return { "file_name": runbook.name, "file_path": runbook.path, @@ -138,15 +137,18 @@ def _format_runbook(self, runbook, repo): "repo_name": repo.get("name"), "repo_display_name": repo.get("display_name"), "provider_type": "github", - "provider_id": self.config.authentication.get("provider_id"), - "link": f"https://api.github.com/{repo.get('full_name')}/blob/{repo.get('default_branch')}/{runbook.path}", - "content": runbook.content, - "encoding": runbook.encoding, + "provider_id": self.provider_id, + "contents": [{ + "content":runbook.content, + "link": f"https://api.github.com/{repo.get('full_name')}/blob/{repo.get('default_branch')}/{runbook.path}", + "encoding": runbook.encoding + }], + "title": title, } - def pull_runbook(self, repo=None, branch=None, md_path=None): + def pull_runbook(self, repo=None, branch=None, md_path=None, title=None): """Retrieve markdown files from the GitHub repository using the GitHub client.""" repo_name = repo if repo else self.authentication_config.repository @@ -167,7 +169,7 @@ def pull_runbook(self, repo=None, branch=None, md_path=None): raise Exception(f"Repository {repo_name} not found") runbook = repo.get_contents(md_path, branch) - response = self._format_runbook(runbook, self._format_repo(repo)) + response = self._format_runbook(runbook, self._format_repo(repo), title) return response except GithubException as e: @@ -250,7 +252,7 @@ def _query( }, ) provider = GithubProvider(context_manager, provider_id="github", config=config) - result = provider.pull_runbook() + result = provider.pull_runbook(title="test") result = provider.pull_repositories() print(result) diff --git a/keep/providers/gitlab_provider/gitlab_provider.py b/keep/providers/gitlab_provider/gitlab_provider.py index 480f5dbff..114f27b6f 100644 --- a/keep/providers/gitlab_provider/gitlab_provider.py +++ b/keep/providers/gitlab_provider/gitlab_provider.py @@ -228,13 +228,12 @@ def pull_repositories(self, project_id=None): raise Exception("Failed to get repositories: personal_access_token not set") - def _format_runbook(self, runbook, repo): + def _format_runbook(self, runbook, repo, title): """ Format the runbook data into a dictionary. """ - if runbook is None: - return {} - + + # TO DO. currently we are handling the one file only. we user give folder path. then we might get multiple files as input(runbook) return { "file_name": runbook.get("file_name"), "file_path": runbook.get("file_path"), @@ -244,13 +243,16 @@ def _format_runbook(self, runbook, repo): "repo_name": repo.get("name"), "repo_display_name": repo.get("display_name"), "provider_type": "gitlab", - "config": self.config.authentication.get("provider_id"), - "link": f"{self.gitlab_host}/api/v4/projects/{repo.get('id')}/repository/files/{runbook.get('file_path')}/raw", - "content": runbook.get("content"), - "encoding": runbook.get("encoding"), + "config": self.provider_id, + "contents": [{ + "content": runbook.get("content"), + "link": f"{self.gitlab_host}/api/v4/projects/{repo.get('id')}/repository/files/{runbook.get('file_path')}/raw", + "encoding": runbook.get("encoding"), + }], + "title": title, } - def pull_runbook(self, repo=None, branch=None, md_path=None): + def pull_runbook(self, repo=None, branch=None, md_path=None, title=None): """Retrieve markdown files from the GitLab repository.""" repo = repo if repo else self.authentication_config.repository branch = branch if branch else "main" @@ -270,7 +272,7 @@ def pull_runbook(self, repo=None, branch=None, md_path=None): except HTTPError as e: raise Exception(f"Failed to get runbook: {e}") - return self._format_runbook(resp.json(), repo_meta) + return self._format_runbook(resp.json(), repo_met, title) raise Exception("Failed to get runbook: repository or md_path not set") @@ -326,6 +328,6 @@ def _notify(self, id: str, title: str, description: str = "", labels: str = "", description="Test Alert Description", ) - result = provider.pull_runbook() + result = provider.pull_runbook(title="test") result = provider.pull_repositories() print(result) diff --git a/keep/runbooks/runbooks_service.py b/keep/runbooks/runbooks_service.py new file mode 100644 index 000000000..05d13c4a0 --- /dev/null +++ b/keep/runbooks/runbooks_service.py @@ -0,0 +1,51 @@ +import logging +from typing import List, Optional +from pydantic import ValidationError +from sqlalchemy.orm import joinedload, selectinload +from uuid import UUID +import json + +from sqlmodel import Session, select +from keep.api.models.db.runbook import ( + Runbook, + RunbookContent, + RunbookDtoOut +) +logger = logging.getLogger(__name__) + + +class RunbookService: + @staticmethod + def create_runbook(session: Session, tenant_id: str, runbook_dto: dict): + try: + new_runbook = Runbook( + tenant_id=tenant_id, + title=runbook_dto["title"], + repo_id=runbook_dto["repo_id"], + relative_path=runbook_dto["file_path"], + provider_type=runbook_dto["provider_type"], + provider_id=runbook_dto["provider_id"] + ) + + session.add(new_runbook) + session.flush() + contents = runbook_dto["contents"] if runbook_dto["contents"] else [] + + new_contents = [ + RunbookContent( + runbook_id=new_runbook.id, + content=content["content"], + link=content["link"], + encoding=content["encoding"] + ) + for content in contents + ] + + session.add_all(new_contents) + session.commit() + session.expire(new_runbook, ["contents"]) + session.refresh(new_runbook) # Refresh the runbook instance + result = RunbookDtoOut.from_orm(new_runbook) + return result + except ValidationError as e: + logger.exception(f"Failed to create runbook {e}") \ No newline at end of file