diff --git a/examples/axon-git/.gitignore b/examples/axon-git/.gitignore new file mode 100644 index 0000000..46e5b6d --- /dev/null +++ b/examples/axon-git/.gitignore @@ -0,0 +1,3 @@ +.venv +__pycache__ +*.pyc \ No newline at end of file diff --git a/examples/axon-git/Dockerfile b/examples/axon-git/Dockerfile new file mode 100644 index 0000000..4b472d9 --- /dev/null +++ b/examples/axon-git/Dockerfile @@ -0,0 +1,15 @@ +FROM ghcr.io/cortexapps/cortex-axon-agent:latest +WORKDIR /project +COPY requirements.txt . +ENV VIRTUAL_ENV=/project/.venv +RUN python3 -m venv $VIRTUAL_ENV +ENV PATH="$VIRTUAL_ENV/bin:$PATH" + +# Install Python dependencies +RUN pip install --no-cache-dir -r requirements.txt +COPY . . +RUN echo "python main.py" > /project/run.sh +RUN chmod +x /project/run.sh +ENV PYTHONUNBUFFERED=1 + +ENTRYPOINT [ "/app/app_entrypoint.sh", "/project/run.sh" ] diff --git a/examples/axon-git/Makefile b/examples/axon-git/Makefile new file mode 100644 index 0000000..8805570 --- /dev/null +++ b/examples/axon-git/Makefile @@ -0,0 +1,11 @@ +IMAGE_NAME ?= ghcr.io/cortexapps/cortex-axon-agent:latest + +docker: + docker build --build-arg "IMAGE_NAME=$(IMAGE_NAME)" -t axon-git:local . + +dryrun: docker + docker run -e "DRYRUN=1" --rm axon-git:local + + +run: docker + docker run -e "CORTEX_API_TOKEN=$(CORTEX_API_TOKEN)" --rm axon-git:local diff --git a/examples/axon-git/README.md b/examples/axon-git/README.md new file mode 100644 index 0000000..bc95f37 --- /dev/null +++ b/examples/axon-git/README.md @@ -0,0 +1,89 @@ +# Cortex Axon SDK for Python + +This is the official Cortex Axon SDK for Python. It provides a simple way to interact with and extend your Cortex instance. + +## Getting started + +To run the Cortex SDK you need to get the Axon Agent via Docker: + +``` +docker pull ghcr.io/cortexapps/cortex-axon-agent:latest +``` + +Then to scaffold a Python project: + +``` +docker run -it --rm -v "$(pwd):/src" ghcr.io/cortexapps/cortex-axon-agent:latest init --language python my-python-axon project +``` + +This will create a new directory `my-python-axon` with a Go project scaffolded in the current directory. + +## Running locally + +To run your project, first start the agent Docker container like: + +``` +docker run -it --rm -p "50051:50051" -e "DRYRUN=true" ghcr.io/cortexapps/cortex-axon-agent:latest serve +``` + +This is `DRYRUN` mode that prints what it would have called, to run against the Cortex API remove the `DRYRUN` environment variable and add `-e "CORTEX_API_TOKEN=$CORTEX_API_TOKEN`. Be sure to export your token first, e.g. `export CORTEX_API_TOKEN=your-token`. + + +## Adding handlers + +To add a handler, open `main.py` and create a function: + +```python + +@cortex_scheduled(interval="5s") +def my_handler(ctx: HandlerContext): + + payload = { + "values": { + "my-service": [ + { + "key": "exampleKey1", + "value": "exampleValue1", + }, + { + "key": "exampleKey2", + "value": "exampleValue2", + }, + ] + } + } + + json_payload = json.dumps(payload) + + response = ctx.api.Call( + cortex_api_pb2.CallRequest( + method="PUT", + path="/api/v1/catalog/custom-data", + body=json_payload, + ) + ) + + if response.status_code >= 400: + ctx.log(f"SetCustomTags error: {response.body}", level="ERROR") + exit(1) + + ctx.log("CortexApi PUT custom-data called successfully!") + +``` + +Now start the agent in a separate terminal: +``` +make run-agent +``` + +And run your project: +``` +python main.py +``` + +This will begin executing your handler every 5 seconds. + + + + + diff --git a/examples/axon-git/git_manager/access.py b/examples/axon-git/git_manager/access.py new file mode 100644 index 0000000..34a5153 --- /dev/null +++ b/examples/axon-git/git_manager/access.py @@ -0,0 +1,95 @@ +import threading +from queue import Queue +from readerwriterlock import rwlock +from git_manager.manager import GitRepository + + +class TaskManager: + def __init__(self, parallel_limit: int): + """ + Initialize the TaskManager. + + :param parallel_limit: Maximum number of tasks to run in parallel. + """ + self.parallel_limit = parallel_limit + self.task_queue = Queue() + self.locks = {} # Dictionary to store ReaderWriterLocks for each task type + self.lock = threading.Lock() # Lock to protect access to the locks dictionary + self.threads = [] + + def _get_lock(self, task_type: str): + """ + Get or create a ReaderWriterLock for the given task type. + + :param task_type: The type of task. + :return: A ReaderWriterLock object. + """ + with self.lock: + if task_type not in self.locks: + self.locks[task_type] = rwlock.RWLockFair() + return self.locks[task_type] + + def run_task(self, repo: GitRepository, action: callable): + done_event = threading.Event() + task = GitTask(repo, self._get_lock(repo.target_dir), action) + def task_wrapper(): + try: + task.run() + + finally: + done_event.set() # Signal that the task is complete + + self.task_queue.put(task_wrapper) + done_event.wait() + return task.result + + def _worker(self): + """ + Worker thread to process tasks from the queue. + """ + while True: + task = self.task_queue.get() + try: + task() + except Exception as e: + print(f"Error executing task: {e}") + self.task_queue.task_done() + + def run(self): + """ + Start processing tasks with the specified parallelization limit. + """ + for _ in range(self.parallel_limit): + thread = threading.Thread(target=self._worker) + thread.start() + self.threads.append(thread) + +class GitTask: + def __init__(self, repo: GitRepository, rwlock: rwlock, action: callable): + self.repo = repo + self.action = action + self.rwlock = rwlock + self.result = None + + def key(self): + return self.repo.target_dir + + def run(self): + if self.repo.needs_update(): + wlock = self.rwlock.gen_wlock() + with wlock: + try: + self.repo.update(force=True) + except Exception as e: + print(f"Error executing task {task.key()}: {e}") + raise e + + with self.rwlock.gen_rlock(): + try: + self.result = self.action() + return self.result + except Exception as e: + print(f"Error executing task {task.key()}: {e}") + raise e + + \ No newline at end of file diff --git a/examples/axon-git/git_manager/formatter.py b/examples/axon-git/git_manager/formatter.py new file mode 100644 index 0000000..28567f3 --- /dev/null +++ b/examples/axon-git/git_manager/formatter.py @@ -0,0 +1,69 @@ +import jinja2 + + +GIT_REPO_URL_TEMPLATE_KEY = "git_repo_url" +GIT_COMMIT_URL_TEMPLATE_KEY = "git_commit_url" +GIT_BLOB_URL_TEMPLATE_KEY = "git_blob_url" +GIT_BRANCH_URL_TEMPLATE_KEY = "git_branch_url" + +default_templates = { + GIT_REPO_URL_TEMPLATE_KEY: "https://{{ git_host }}/{{ repo_name }}", + GIT_COMMIT_URL_TEMPLATE_KEY: "https://{{ git_host }}/{{ repo_name }}/commit/{{ sha }}", + GIT_BLOB_URL_TEMPLATE_KEY: "https://{{ git_host }}/{{ repo_name }}/blob/{{ sha }}/{{ path }}", + GIT_BRANCH_URL_TEMPLATE_KEY: "https://{{ git_host }}/{{ repo_name }}/tree/{{ ref }}" +} + +class GitFormatter: + _global_formatter = None + @staticmethod + def set_global_formatter(formatter: 'GitFormtter'): + GitFormatter._global_formatter = formatter + + @staticmethod + def get_global_formatter() -> 'GitFormtter': + if GitFormtter._global_formatter is None: + raise ValueError("Global formatter not set") + return GitFormtter._global_formatter + + def __init__(self, git_host: str, templates: dict): + self.git_host = git_host + self.templates = default_templates.copy() + self.templates.update(templates) + + def __render(self, template_key: str, data: dict) -> str: + template = self.templates.get(template_key) + if not template: + raise ValueError(f"Template {template_key} not found") + return jinja2.Template(template).render(data) + + def repo_url(self, repo_name: str) -> str: + data = { + "git_host": self.git_host, + "repo_name": repo_name + } + return self.__render(GIT_REPO_URL_TEMPLATE_KEY, data) + + def commit_url(self, repo_name: str, sha: str) -> str: + data = { + "git_host": self.git_host, + "repo_name": repo_name, + "sha": sha + } + return self.__render(GIT_COMMIT_URL_TEMPLATE_KEY, data) + + def blob_url(self, repo_name: str, sha: str, path: str) -> str: + data = { + "git_host": self.git_host, + "repo_name": repo_name, + "sha": sha, + "path": path + } + return self.__render(GIT_BLOB_URL_TEMPLATE_KEY, data) + + def branch_url(self, repo_name: str, ref: str) -> str: + data = { + "git_host": self.git_host, + "repo_name": repo_name, + "ref": ref + } + return self.__render(GIT_BRANCH_URL_TEMPLATE_KEY, data) diff --git a/examples/axon-git/git_manager/manager.py b/examples/axon-git/git_manager/manager.py new file mode 100644 index 0000000..f7bb906 --- /dev/null +++ b/examples/axon-git/git_manager/manager.py @@ -0,0 +1,264 @@ +import pygit2 +from pathlib import Path +import os +from datetime import datetime +from git_manager.formatter import GitFormatter + +CLONE_DEPTH = 100 +UPDATE_TTL_SECONDS = 300 + +class GitRepositoryManager: + + _credentials_by_type = { + pygit2.enums.CredentialType.USERNAME: None, + pygit2.enums.CredentialType.SSH_KEY: None, + pygit2.enums.CredentialType.USERPASS_PLAINTEXT: None, + } + + def add_credentials(self, creds): + existing = GitRepositoryManager._credentials_by_type.get(creds.credential_type) + if existing: + raise ValueError(f"Credentials for {creds.credential_type} already set.") + GitRepositoryManager._credentials_by_type[creds.credential_type] = creds + + repositories = {} + root_dir: str = None + + def __init__(self, root_dir: str, formatter: GitFormatter): + self.root_dir = root_dir + self.formatter = formatter + os.makedirs(root_dir, exist_ok=True) + + def get(self, repo_name: str, branch: str = None) -> 'GitRepository': + repo_url = self.formatter.repo_url(repo_name) + repo = GitRepository(repo_name, repo_url, self.root_dir, branch) + if repo.target_dir not in self.repositories: + self.repositories[repo.target_dir] = repo + return self.repositories[repo.target_dir] + + +class GitRepository: + + last_updated: datetime = None + + def __init__(self, repo_name: str, repo_url: str, root_dir: str, branch: str = None): + """ + Initialize the GitRepositoryManager. + + :param repo_url: The HTTPS URL of the repository to clone. + :param username: The username for authentication. + :param password: The password or personal access token for authentication. + :param target_dir: The directory where the repository will be cloned. + """ + self.repo_url = repo_url + self.repo_name = repo_name + self.branch = None + self.root_dir = root_dir + self._callbacks = MyRemoteCallbacks(GitRepositoryManager._credentials_by_type) + self.target_dir = self._get_repo_path(repo_url) + + def needs_update(self) -> bool: + if not self.last_updated: + return True + + # Check if the last update was more than 1 hour ago + now = datetime.now() + delta = now - self.last_updated + if delta.total_seconds() > UPDATE_TTL_SECONDS: + return True + + return False + + def exists(self) -> bool: + """ + Check if the repository exists in the target directory. + + :return: True if the repository exists, False otherwise. + """ + if Path(self.target_dir).exists(): + return True + + try: + self.update() + except Exception as e: + self.last_updated = datetime.now() + pass + return Path(self.target_dir).exists() + + + def get_default_branch(self): + + try: + # Open the repository in read-only mode + repo = pygit2.Repository(path=self.target_dir) + # Get the default branch + self.branch = repo.head.shorthand + return self.branch + except Exception as e: + print(f"Failed to get default branch: {e}") + raise + + def get_branches(self) -> list[dict]: + try: + # Open the repository in read-only mode + repo = pygit2.Repository(path=self.target_dir) + # Get the list of branches + branches = [] + for branch in repo.branches.local: + branches.append({"name": branch, "head": str(repo.branches.local[branch].raw_target)}) + return branches + except Exception as e: + print(f"Failed to get branches: {e}") + raise + + def _get_repo_path(self, repo_url: str): + if repo_url.startswith("https://"): + repo_url = repo_url.replace("https://", "") + + # replace : with _ + repo_url = repo_url.replace(":", "/") + branch = self.branch if self.branch else "__default__" + return self.root_dir + "/" + repo_url + "/" + branch + + def _get_file(self, file_path: str) -> Path: + + if file_path.startswith(self.target_dir): + return Path(file_path) + + full_path = self.target_dir + "/" + file_path.lstrip("/") + return Path(full_path) + + def get_file_contents_string(self, file_path: str) -> str: + file = self._get_file(file_path) + if not file.exists(): + raise FileNotFoundError(f"File {file_path} not found in repository {self.repo_name}") + with open(file, "r") as f: + return f.read() + return None + + def get_file_contents_binary(self, file_path: str) -> str: + file = self._get_file(file_path) + if not file.exists(): + raise FileNotFoundError(f"File {file_path} not found in repository {self.repo_name}") + + # Read binary file content then return it as base64 encoded string + with open(file, "rb") as f: + content = f.read() + return content.decode("utf-8") + + return None + + def file_path_exists(self, file_path: str) -> bool: + return self._get_file(file_path).exists() + + def _to_repo_relative(self, file_path: str) -> str: + if file_path.startswith(self.target_dir): + file_path = file_path[len(self.target_dir):].lstrip("/") + return file_path + + def get_file_sha(self, file_path: str) -> str: + try: + repo = pygit2.Repository(path=self.target_dir) + rel_path = self._to_repo_relative(file_path) + blob_sha = repo.get(repo.head.target).tree[rel_path].id + return str(blob_sha) + except Exception as e: + print(f"Failed to get file SHA: {e}") + return None + + def update(self, force: bool = False) -> str: + if not force and not self.needs_update(): + return None + + print(f"Updating repository {self.repo_name}") + try: + if not Path(self.target_dir).exists(): + return self.clone() + + self.fetch_and_reset() + + except Exception as e: + print(f"Failed to update repository, cloning: {e}") + return self.clone() + + def commits(self, limit: int = 10) -> list[pygit2.Commit]: + repo = pygit2.Repository(path=self.target_dir) + last = repo[repo.head.target] + for commit in repo.walk(last.id, pygit2.enums.SortMode.TIME): + if limit <= 0: + break + + if not isinstance(commit, pygit2.Commit): + continue + + limit -= 1 + yield commit + + def sha(self) -> str: + try: + # Open the repository + repo = pygit2.Repository(path=self.target_dir) + # Get the current commit SHA + sha = str(repo.head.raw_target) + return sha + except Exception as e: + print(f"Failed to get SHA: {e}") + return None + + def fetch_and_reset(self) -> str: + try: + # Open the repository + repo = pygit2.Repository(path=self.target_dir) + # Fetch the latest changes + remote = repo.remotes["origin"] + remote.fetch(callbacks=self._callbacks) + # Hard reset to the latest commit + repo.reset(repo.head.target, pygit2.GIT_RESET_HARD) + self.last_updated = datetime.now() + print(f"Repository updated successfully in {self.target_dir} to {repo.head.target}") + except Exception as e: + print(f"Failed to update repository: {e}") + raise + + def clone(self) -> str: + try: + # Clone the repository + repo = pygit2.clone_repository( + self.repo_url, + self.target_dir, + depth=CLONE_DEPTH, + callbacks=self._callbacks, + checkout_branch=self.branch + ) + self.last_updated = datetime.now() + print(f"Repository cloned successfully to {self.target_dir} with SHA {repo.head.target}") + return repo.head.target + except Exception as e: + print(f"Failed to clone repository: {e}") + raise + + +class MyRemoteCallbacks(pygit2.RemoteCallbacks): + + def __init__(self, credentials_by_type): + super().__init__() + self.credentials_by_type = credentials_by_type + + def transfer_progress(self, stats): + #print(f'{stats.indexed_objects}/{stats.total_objects}') + pass + + def credentials(self, url, username_from_url, allowed_types): + if allowed_types & pygit2.enums.CredentialType.USERNAME: + creds = self.credentials_by_type[pygit2.enums.CredentialType.USERNAME] + elif allowed_types & pygit2.enums.CredentialType.SSH_KEY: + creds = self.credentials_by_type[pygit2.enums.CredentialType.SSH_KEY] + elif allowed_types & pygit2.enums.CredentialType.USERPASS_PLAINTEXT: + creds = self.credentials_by_type[pygit2.enums.CredentialType.USERPASS_PLAINTEXT] + else: + return None + + if not creds: + raise ValueError("Credentials not provided for the requested type: {}".format(allowed_types)) + return creds + diff --git a/examples/axon-git/git_manager/models/git.py b/examples/axon-git/git_manager/models/git.py new file mode 100644 index 0000000..dc80e23 --- /dev/null +++ b/examples/axon-git/git_manager/models/git.py @@ -0,0 +1,66 @@ +import pygit2 +from typing import Optional +from datetime import datetime, timedelta, timezone +from dataclasses import dataclass + +@dataclass +class Contributor: + email: str + name: Optional[str] = None + username: Optional[str] = None + image: Optional[str] = None + numCommits: Optional[int] = None + url: Optional[str] = None + alias: Optional[str] = None + + def to_dict(self) -> dict: + return { + "email": self.email, + "name": self.name, + "username": self.username, + "image": self.image, + "numCommits": self.numCommits, + "url": self.url, + "alias": self.alias + } + +@dataclass +class Commit: + url: Optional[str] + sha: str + committer: Contributor + message: str + date: datetime + + @staticmethod + def from_commit(commit: pygit2.Commit, url: str = None) -> 'Commit': + + username_from_email = commit.author.email and commit.author.email.split("@")[0] + + return Commit( + url = url, + sha = str(commit.id), + message = commit.message.strip(), + committer= Contributor( + name = commit.author.name, + email = commit.author.email, + username=username_from_email, + ), + date = _format_time(commit), + ) + + def to_dict(self) -> dict: + return { + "url": self.url, + "sha": self.sha, + "committer": self.committer.to_dict(), + "message": "[AXON] " + self.message, + "date": self.date, + } + + +def _format_time(commit: Commit) -> str: + ts = datetime.fromtimestamp(commit.commit_time) + offset = timedelta(minutes=commit.commit_time_offset) + tz = timezone(offset) + return ts.replace(tzinfo=tz).isoformat() \ No newline at end of file diff --git a/examples/axon-git/main.py b/examples/axon-git/main.py new file mode 100644 index 0000000..2c704bf --- /dev/null +++ b/examples/axon-git/main.py @@ -0,0 +1,349 @@ +import json +from cortex_axon.axon_client import AxonClient, HandlerContext +from cortex_axon.handler import cortex_handler +from git_manager import manager, formatter +import os +import pygit2 +from typing import Optional +from datetime import datetime, timedelta, timezone +import isodate +from dataclasses import dataclass +from git_manager.models.git import Commit +from repository_info import RepositoryInfo +from pathlib import Path +from git_manager.access import TaskManager +from git_manager.manager import GitRepositoryManager +from concurrent.futures import ThreadPoolExecutor + + +root_dir = "/tmp/cortex-axon-git" + +formatter = formatter.GitFormatter( + git_host = "github.com", + templates = { + "git_repo_url": "https://{{ git_host }}/{{ repo_name }}.git", + "git_commit_url": "https://{{ git_host }}/{{ repo_name }}/commit/{{ sha }}", + "git_file_url": "https://{{ git_host }}/{{ repo_name }}/blob/{{ sha }}/{{ path }}", + "git_branch_url": "https://{{ git_host }}/{{ repo_name }}/tree/{{ ref }}", + } +) +formatter.set_global_formatter(formatter) + +repo_manager = GitRepositoryManager(root_dir=root_dir, formatter=formatter) + +task_manager = TaskManager(parallel_limit=20) + +def __get_body_arg(context: HandlerContext, arg: str, default: object = None) -> object: + body = context.args.get("body") + if not body: + raise ValueError("No body provided") + body_dict = json.loads(body) + return body_dict.get(arg, default) + +@cortex_handler() +def get_repo_exists(context: HandlerContext) -> str: + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + result = repo.exists() + return json_dumps( { + "exists": result + }) + +@cortex_handler() +def get_repository_details(context: HandlerContext) -> str: + + # data class GitRepoDetails( + # val name: String, + # val url: String, + # val commits: List, + # val contributors: List, + # val releases: List?, + # val language: String?, + # val provider: GitProvider, + # val missingPermissions: List, + # val basePath: String?, + # val defaultBranch: String, + # val alias: String?, + # ) + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + + commits = _get_commits(context, limit=1000, lookback="P1M") + contributor_summary = {} + for commit in commits: + author = commit.committer + if not author: + continue + if author.email not in contributor_summary: + contributor_summary[author.email] = [] + contributor_summary[author.email].append(commit) + + contributors = [] + for email, user_commits in contributor_summary.items(): + c = user_commits[0].committer + contributors.append({ + "email": c.email, + "username": c.username or c.email, + "name": c.name, + "numCommits": len(user_commits), + # "url": formatter.user_url(c.username or c.email), + "alias": c.alias, + }) + + contributors.sort(key=lambda x: x["numCommits"], reverse=True) + + return json_dumps( { + "name": repo.repo_name, + "url": repo.repo_url, + "commits": commits[:100], + "contributors": contributors, + "defaultBranch": repo.get_default_branch(), + }) + + +def _get_commits(context: HandlerContext, limit: int = 100, lookback: str = None) -> [Commit]: + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + lookback = __get_body_arg(context, "lookback", lookback or "P7D") + + commits = task_manager.run_task(repo=repo, action = lambda: repo.commits(limit=limit)) + + delta = isodate.parse_duration(lookback) + now = datetime.now(timezone.utc) + start_time = now - delta + + # Map the commits to a new shape + mapped_commits = [] + for commit in commits: + if commit.commit_time < start_time.timestamp(): + continue + + url = formatter.commit_url(repo.repo_name, commit.id) + mapped_commit = Commit.from_commit(commit, url=url) + mapped_commits.append(mapped_commit) + return list(mapped_commits) + +@cortex_handler() +def get_commits(context: HandlerContext) -> str: + + + commits = _get_commits(context) + + return json_dumps({ + "commits": commits, + }) + +@cortex_handler() +def get_last_commit(context: HandlerContext) -> str: + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + + commits = task_manager.run_task(repo=repo, action = lambda: repo.commits(limit=1)) + # Map the commits to a new shape + mapped_commits = [] + for commit in commits: + url = formatter.commit_url(repo.repo_name, commit.id) + mapped_commit = Commit.from_commit(commit, url=url) + mapped_commits.append(mapped_commit) + + if not mapped_commits: + return "", + + commit = mapped_commits[0] if len(mapped_commits) > 0 else None + return json_dumps({ + "commit": commit + }) + + +@cortex_handler() +def get_branches(context: HandlerContext) -> str: + + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + + branches = task_manager.run_task(repo=repo, action = lambda: repo.branches()) + + # Map the branches to a new shape + mapped_branches = [] + for branch in branches: + mapped_branch = { + "name": branch.name, + "head": str(branch.head), + "url": formatter.branch_url(repo.repo_name, branch.name), + } + mapped_branches.append(mapped_branch) + + return json_dumps({ + "branches": mapped_branches, + }) + +@cortex_handler() +def file_path_exists(context: HandlerContext) -> str: + + filePath = __get_body_arg(context, "filePath", None) + if not filePath: + raise ValueError("No filePath provided") + + basePath = __get_body_arg(context, "basePath", None) + if basePath: + filePath = os.path.join(basePath, filePath) + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + result = task_manager.run_task(repo, lambda: repo.file_path_exists(filePath)) + return json_dumps( { + "filePath": filePath, + "exists": result + }) + + +# Note: language is not supported +# data class GitSearchParams( +# val query: String? = null, +# val inFile: Boolean? = null, +# val inPath: Boolean? = null, +# val path: String? = null, +# val fileName: String? = null, +# val fileExtension: String? = null, +# ) + +def _searchfile(file: Path, query: str) -> bool: + try: + with open(file, "r") as f: + content = f.read() + return query.casefold() in content.casefold() + except Exception as e: + print(f"Error reading file {file}: {e}") + return False + +@cortex_handler() +def search_code(context: HandlerContext) -> str: + + params = __get_body_arg(context, "params", None) + if not params: + raise ValueError("No params provided") + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + + repo_info = RepositoryInfo.from_context(context) + + + target_dir = repo.target_dir + + basePath = __get_body_arg(context, "basePath", None) + if basePath: + target_dir = os.path.join(target_dir, basePath) + + subPath = params.get("path", None) + if subPath: + target_dir = os.path.join(target_dir, subPath) + + # gather all of the files that match the search + glob_suffix = "*" + + if params.get("fileName"): + glob_suffix = params.get("fileName") + elif params.get("fileExtension"): + glob_suffix = f"*{params.get("fileExtension")}" + + files_in_scope = list(Path(target_dir).glob("**/{}".format(glob_suffix))) + files_in_scope = [f for f in files_in_scope if f.is_file() and not str(f).startswith(".git")] + query = params.get("query", "") + + + + # Trim off the target_dir from the file path + def get_results(): + output = [] + + if params.get("inPath"): + files_selected = [f for f in files_in_scope if query.casefold() in str(f).casefold()] + else: + with ThreadPoolExecutor(max_workers=10) as executor: + files_selected = list(executor.map(lambda f: (f, _searchfile(f, query)), files_in_scope)) + + for file, found in files_selected: + if not found: + continue + file_path = str(file) + if file_path.startswith(target_dir): + file_path = file_path[len(target_dir):] + file_path = file_path.lstrip("/") + output.append({ + "name": file.name, + "path": file_path, + "sha": repo.get_file_sha(str(file)), + "url": formatter.blob_url(repo.repo_name, repo.sha(), file_path), + }) + return output + + result_files = task_manager.run_task(repo, action=get_results) + + return json_dumps( { + "results": result_files or [], + }) + + + + +@cortex_handler() +def read_file(context: HandlerContext) -> str: + + filePath = __get_body_arg(context, "filePath", None) + if not filePath: + raise ValueError("No filePath provided") + + basePath = __get_body_arg(context, "basePath", None) + if basePath: + filePath = os.path.join(basePath, filePath) + + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + + result = repo.get_file_contents_string(filePath) + return json_dumps( { + "content": result, + }) + + + +@cortex_handler() +def read_file_binary(context: HandlerContext) -> str: + + filePath = __get_body_arg(context, "filePath", None) + if not filePath: + raise ValueError("No filePath provided") + + basePath = __get_body_arg(context, "basePath", None) + if basePath: + filePath = os.path.join(basePath, filePath) + + + repo = RepositoryInfo.repo_from_context(context, repo_manager) + + repo_info = RepositoryInfo.from_context(context) + + result = repo.get_file_contents_binary(filePath) + return json_dumps( { + "content": result, + }) + + +def json_dumps(data) -> str: + return json.dumps(data, indent=4, default=lambda x: x.to_dict()) + +def run(): + + repo_manager.add_credentials( + pygit2.credentials.UserPass( + os.getenv("GITHUB_USERNAME"), + os.getenv("GITHUB_PASSWORD"), + ) + ) + task_manager.run() + client = AxonClient(scope=globals()) + client.run() + +if __name__ == '__main__': + run() + + diff --git a/examples/axon-git/repository_info.py b/examples/axon-git/repository_info.py new file mode 100644 index 0000000..e5722ec --- /dev/null +++ b/examples/axon-git/repository_info.py @@ -0,0 +1,54 @@ + +import json +from cortex_axon.axon_client import AxonClient, HandlerContext +from cortex_axon.handler import cortex_handler +from git_manager import manager, formatter +import os +import pygit2 +from typing import Optional +from datetime import datetime, timedelta, timezone +from dataclasses import dataclass +from git_manager.models.git import Commit +from git_manager import formatter + +def _get_root_dir() -> str: + root_dir = os.getenv("GIT_ROOT_DIR") + if not root_dir: + return "/tmp/cortex-axon-git" + return root_dir + +@dataclass +class RepositoryInfo: + name: str + basePath: Optional[str] = None + url: Optional[str] = None + branch: Optional[str] = None + + @staticmethod + def from_json(json_str: str) -> 'RepositoryInfo': + json_data = json.loads(json_str) + info = RepositoryInfo( + name=json_data.get("name"), + basePath=json_data.get("basePath", None), + url=json_data.get("url", None), + branch=json_data.get("branch", None) + ) + return info + + + @staticmethod + def repo_from_context(context: HandlerContext, manager: manager.GitRepositoryManager) -> manager.GitRepository: + info = RepositoryInfo.from_context(context) + return manager.get( + repo_name=info.name, + branch=info.branch + ) + + @staticmethod + def from_context(context: HandlerContext) -> manager.GitRepositoryManager: + body = context.args["body"] + if not body: + raise ValueError("No body provided") + + return RepositoryInfo.from_json(body) + diff --git a/examples/axon-git/requirements.txt b/examples/axon-git/requirements.txt new file mode 100644 index 0000000..c5b3604 --- /dev/null +++ b/examples/axon-git/requirements.txt @@ -0,0 +1,9 @@ +grpcio +grpcio-tools +pygit2 +readerwriterlock +jinja2 +isodate + +--extra-index-url https://test.pypi.org/simple +cortex-axon-sdk