From 7798d92d30abda6e6f55d4ab70c55389e21c338c Mon Sep 17 00:00:00 2001
From: Seldom <HaydenBadgerAZ@gmail.com>
Date: Mon, 11 Oct 2021 17:56:13 -0700
Subject: [PATCH] Support for projects in organizations

---
 README.md                   |  6 ++---
 src/config/config.json.dist |  2 +-
 src/gh/api_wrapper.py       | 52 +++++++++++++++++++++++++++++++++----
 src/main.py                 | 16 ++++++++----
 4 files changed, 62 insertions(+), 14 deletions(-)

diff --git a/README.md b/README.md
index 4cae5c3..7ac7465 100644
--- a/README.md
+++ b/README.md
@@ -11,7 +11,7 @@ An easy to use [burndown chart](https://www.scrum.org/resources/scrum-glossary#:
 
 ## Features
 * Create a **burndown chart for a GitHub Project Board**.
-* Works for **private repositories**.
+* Works for projects in **private repositories** and **organizations**.
 * Includes a **trend line** for the current sprint.
 * Supports custom labels for **tracking points for issues**
 
@@ -66,9 +66,9 @@ pip install -r requirements.txt
 4. Make a copy of `src/config/config.json.dist` without the `.dist` ending.
     - This allows the `.gitignore` to exclude your `config.json` from being accidentally committed.
 5. Fill out all the configuration settings
-    - `repo_owner`: The username of the owner of the repo.
+    - `username`: The username of the organization or owner of the repo.
         - For example, `jhale1805`
-    - `repo_name`: The name of the repo.
+    - `repo_name`: The name of the repo. Leave blank if the project is in an organization instead of a repo.
         - For example, `github-projects-burndown-chart`
     - `project_number`: The id of the project for which you want to generate a burndown chart. This is found in the URL when looking at the project board on GitHub.
         - For example, `1` from [`https://github.com/jhale1805/github-projects-burndown-chart/projects/1`](https://github.com/jhale1805/github-projects-burndown-chart/projects/1)
diff --git a/src/config/config.json.dist b/src/config/config.json.dist
index ae8c190..afc6456 100644
--- a/src/config/config.json.dist
+++ b/src/config/config.json.dist
@@ -1,5 +1,5 @@
 {
-    "repo_owner": "",
+    "username": "",
     "repo_name": "",
     "project_number": 1,
     "sprint_start_date": "",
diff --git a/src/gh/api_wrapper.py b/src/gh/api_wrapper.py
index 316e807..d84b22c 100644
--- a/src/gh/api_wrapper.py
+++ b/src/gh/api_wrapper.py
@@ -12,9 +12,9 @@ __ch.setFormatter(
     logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
 __logger.addHandler(__ch)
 
-project_query = """
+repo_project_query = """
 query {
-  repository(owner: "%(repo_owner)s", name: "%(repo_name)s") {
+  repository(owner: "%(username)s", name: "%(repo_name)s") {
     project(number: %(project_number)d) {
       name
       columns(first: 5) {
@@ -46,15 +46,57 @@ query {
 }
 """  # Heavily inspired by https://github.com/radekstepan/burnchart/issues/129#issuecomment-394469442
 
-def get_project(repo_owner: str, repo_name: str, project_number: int) -> dict:
-    query = project_query % {
-        'repo_owner': repo_owner,
+org_project_query = """
+query {
+  organization(login: "%(username)s") {
+    project(number: %(project_number)d) {
+      name
+      columns(first: 5) {
+        nodes {
+          name
+          cards(first: 50) {
+            nodes {
+              id
+              note
+              state
+              content {
+                ... on Issue {
+                  title
+                  createdAt
+                  closedAt
+                  labels(first: 5) {
+                    nodes {
+                      name
+                    }
+                  }
+                }
+              }
+            }
+          }
+        }
+      }
+    }
+  }
+}
+"""  # Heavily inspired by https://github.com/radekstepan/burnchart/issues/129#issuecomment-394469442
+
+def get_repo_project(username: str, repo_name: str, project_number: int) -> dict:
+    query = repo_project_query % {
+        'username': username,
         'repo_name': repo_name,
         'project_number': project_number}
     query_response = gh_api_query(query)
     project_data = query_response['data']['repository']['project']
     return Project(project_data)
 
+def get_org_project(username: str, project_number: int) -> dict:
+    query = org_project_query % {
+        'username': username,
+        'project_number': project_number}
+    query_response = gh_api_query(query)
+    project_data = query_response['data']['organization']['project']
+    return Project(project_data)
+
 def gh_api_query(query: str) -> dict:
     headers = {'Authorization': 'bearer %s' % secrets['github_token']} \
         if 'github_token' in secrets else {}
diff --git a/src/main.py b/src/main.py
index 032adfd..48b47fb 100644
--- a/src/main.py
+++ b/src/main.py
@@ -1,12 +1,18 @@
 from chart.burndown import BurndownChart
 from config import config
-from gh.api_wrapper import get_project
+from gh.api_wrapper import get_repo_project, get_org_project
 
 if __name__ == '__main__':
-    project = get_project(
-        config['repo_owner'],
-        config['repo_name'],
-        config['project_number'])
+    repo_name = config['repo_name']
+    if repo_name:
+        project = get_project(
+            config['username'],
+            config['repo_name'],
+            config['project_number'])
+    else:
+        project = get_org_project(
+            config['username'],
+            config['project_number'])
     burndown_chart = BurndownChart(project)
     burndown_chart.render()
     print('Done')
\ No newline at end of file
-- 
2.51.2