From 6b3080ed39404927b3cb9bfc9b80fd0388680683 Mon Sep 17 00:00:00 2001
From: rongzhang <rongzhang@pinterest.com>
Date: Wed, 16 Oct 2024 16:10:48 +0000
Subject: [PATCH] feat: Add oauth flow for querybook github integration

---
 querybook/server/datasources/github.py        |  16 +++
 .../server/lib/github_integration/__init__.py |   0
 .../github_integration/github_integration.py  | 110 ++++++++++++++++++
 .../DataDocGitHub/DataDocGitHubButton.tsx     |  61 ++++++++++
 .../components/DataDocGitHub/GitHub.scss      |   8 ++
 .../components/DataDocGitHub/GitHubAuth.tsx   |  31 +++++
 querybook/webapp/resource/github.ts           |  11 ++
 7 files changed, 237 insertions(+)
 create mode 100644 querybook/server/datasources/github.py
 create mode 100644 querybook/server/lib/github_integration/__init__.py
 create mode 100644 querybook/server/lib/github_integration/github_integration.py
 create mode 100644 querybook/webapp/components/DataDocGitHub/DataDocGitHubButton.tsx
 create mode 100644 querybook/webapp/components/DataDocGitHub/GitHub.scss
 create mode 100644 querybook/webapp/components/DataDocGitHub/GitHubAuth.tsx
 create mode 100644 querybook/webapp/resource/github.ts

diff --git a/querybook/server/datasources/github.py b/querybook/server/datasources/github.py
new file mode 100644
index 000000000..82420604c
--- /dev/null
+++ b/querybook/server/datasources/github.py
@@ -0,0 +1,16 @@
+from app.datasource import register
+from lib.github_integration.github_integration import get_github_manager
+from typing import Dict
+
+
+@register("/github/auth/", methods=["GET"])
+def connect_github() -> Dict[str, str]:
+    github_manager = get_github_manager()
+    return github_manager.initiate_github_integration()
+
+
+@register("/github/is_authenticated/", methods=["GET"])
+def is_github_authenticated() -> str:
+    github_manager = get_github_manager()
+    is_authenticated = github_manager.get_github_token() is not None
+    return {"is_authenticated": is_authenticated}
diff --git a/querybook/server/lib/github_integration/__init__.py b/querybook/server/lib/github_integration/__init__.py
new file mode 100644
index 000000000..e69de29bb
diff --git a/querybook/server/lib/github_integration/github_integration.py b/querybook/server/lib/github_integration/github_integration.py
new file mode 100644
index 000000000..4cced5956
--- /dev/null
+++ b/querybook/server/lib/github_integration/github_integration.py
@@ -0,0 +1,110 @@
+import certifi
+from flask import session as flask_session, request
+from app.auth.github_auth import GitHubLoginManager
+from env import QuerybookSettings
+from lib.logger import get_logger
+from app.flask_app import flask_app
+from typing import Optional, Dict, Any
+
+LOG = get_logger(__file__)
+
+
+GITHUB_OAUTH_CALLBACK = "/github/oauth2callback"
+
+
+class GitHubIntegrationManager(GitHubLoginManager):
+    def __init__(self, additional_scopes: Optional[list] = None):
+        self.additional_scopes = additional_scopes or []
+        super().__init__()
+
+    @property
+    def oauth_config(self) -> Dict[str, Any]:
+        config = super().oauth_config
+        config["scope"] = "user email " + " ".join(self.additional_scopes)
+        config[
+            "callback_url"
+        ] = f"{QuerybookSettings.PUBLIC_URL}{GITHUB_OAUTH_CALLBACK}"
+        return config
+
+    def save_github_token(self, token: str) -> None:
+        flask_session["github_access_token"] = token
+        LOG.debug("Saved GitHub token to session")
+
+    def get_github_token(self) -> Optional[str]:
+        return flask_session.get("github_access_token")
+
+    def initiate_github_integration(self) -> Dict[str, str]:
+        github = self.oauth_session
+        authorization_url, state = github.authorization_url(
+            self.oauth_config["authorization_url"]
+        )
+        flask_session["oauth_state"] = state
+        return {"url": authorization_url}
+
+    def github_integration_callback(self) -> str:
+        try:
+            github = self.oauth_session
+            access_token = github.fetch_token(
+                self.oauth_config["token_url"],
+                client_secret=self.oauth_config["client_secret"],
+                authorization_response=request.url,
+                cert=certifi.where(),
+            )
+            self.save_github_token(access_token["access_token"])
+            return self.success_response()
+        except Exception as e:
+            LOG.error(f"Failed to obtain credentials: {e}")
+            return self.error_response(str(e))
+
+    def success_response(self) -> str:
+        return """
+            <p>Success! Please close the tab.</p>
+            <script>
+                window.opener.receiveChildMessage()
+            </script>
+        """
+
+    def error_response(self, error_message: str) -> str:
+        return f"""
+            <p>Failed to obtain credentials, reason: {error_message}</p>
+        """
+
+
+def get_github_manager() -> GitHubIntegrationManager:
+    return GitHubIntegrationManager(additional_scopes=["repo"])
+
+
+@flask_app.route(GITHUB_OAUTH_CALLBACK)
+def github_callback() -> str:
+    github_manager = get_github_manager()
+    return github_manager.github_integration_callback()
+
+
+# Test GitHub OAuth Flow
+def main():
+    github_manager = GitHubIntegrationManager()
+    oauth_config = github_manager.oauth_config
+    client_id = oauth_config["client_id"]
+    client_secret = oauth_config["client_secret"]
+
+    from requests_oauthlib import OAuth2Session
+
+    github = OAuth2Session(client_id)
+    authorization_url, state = github.authorization_url(
+        oauth_config["authorization_url"]
+    )
+    print("Please go here and authorize,", authorization_url)
+
+    redirect_response = input("Paste the full redirect URL here:")
+    github.fetch_token(
+        oauth_config["token_url"],
+        client_secret=client_secret,
+        authorization_response=redirect_response,
+    )
+
+    user_profile = github.get(oauth_config["profile_url"]).json()
+    print(user_profile)
+
+
+if __name__ == "__main__":
+    main()
diff --git a/querybook/webapp/components/DataDocGitHub/DataDocGitHubButton.tsx b/querybook/webapp/components/DataDocGitHub/DataDocGitHubButton.tsx
new file mode 100644
index 000000000..a34f4626e
--- /dev/null
+++ b/querybook/webapp/components/DataDocGitHub/DataDocGitHubButton.tsx
@@ -0,0 +1,61 @@
+import React, { useCallback, useEffect, useState } from 'react';
+
+import { GitHubResource } from 'resource/github';
+import { IconButton } from 'ui/Button/IconButton';
+
+import { GitHubModal } from './GitHubModal';
+
+interface IProps {
+    docId: number;
+}
+
+export const DataDocGitHubButton: React.FunctionComponent<IProps> = ({
+    docId,
+}) => {
+    const [isModalOpen, setIsModalOpen] = useState(false);
+    const [isAuthenticated, setIsAuthenticated] = useState(false);
+
+    useEffect(() => {
+        const checkAuthentication = async () => {
+            try {
+                const { data } = await GitHubResource.isAuthenticated();
+                setIsAuthenticated(data.is_authenticated);
+            } catch (error) {
+                console.error(
+                    'Failed to check GitHub authentication status:',
+                    error
+                );
+            }
+        };
+
+        checkAuthentication();
+    }, []);
+
+    const handleOpenModal = useCallback(() => {
+        setIsModalOpen(true);
+    }, []);
+
+    const handleCloseModal = useCallback(() => {
+        setIsModalOpen(false);
+    }, []);
+
+    return (
+        <>
+            <IconButton
+                icon="Github"
+                onClick={handleOpenModal}
+                tooltip="Connect to GitHub"
+                tooltipPos="left"
+                title="GitHub"
+            />
+            {isModalOpen && (
+                <GitHubModal
+                    docId={docId}
+                    isAuthenticated={isAuthenticated}
+                    setIsAuthenticated={setIsAuthenticated}
+                    onClose={handleCloseModal}
+                />
+            )}
+        </>
+    );
+};
diff --git a/querybook/webapp/components/DataDocGitHub/GitHub.scss b/querybook/webapp/components/DataDocGitHub/GitHub.scss
new file mode 100644
index 000000000..f7b911929
--- /dev/null
+++ b/querybook/webapp/components/DataDocGitHub/GitHub.scss
@@ -0,0 +1,8 @@
+.GitHubAuth {
+    text-align: center;
+    padding: 20px;
+}
+
+.GitHubAuth-icon {
+    margin-bottom: 20px;
+}
diff --git a/querybook/webapp/components/DataDocGitHub/GitHubAuth.tsx b/querybook/webapp/components/DataDocGitHub/GitHubAuth.tsx
new file mode 100644
index 000000000..0f3caebe6
--- /dev/null
+++ b/querybook/webapp/components/DataDocGitHub/GitHubAuth.tsx
@@ -0,0 +1,31 @@
+import React from 'react';
+
+import { Button } from 'ui/Button/Button';
+import { Icon } from 'ui/Icon/Icon';
+import { Message } from 'ui/Message/Message';
+
+import './GitHub.scss';
+
+interface IProps {
+    onAuthenticate: () => void;
+}
+
+export const GitHubAuth: React.FunctionComponent<IProps> = ({
+    onAuthenticate,
+}) => (
+    <div className="GitHubAuth">
+        <Icon name="Github" size={64} className="GitHubAuth-icon" />
+        <Message
+            title="Connect to GitHub"
+            message="You currently do not have a GitHub provider linked to your account. Please authenticate to enable GitHub features on Querybook."
+            type="info"
+            iconSize={32}
+        />
+        <Button
+            onClick={onAuthenticate}
+            title="Connect Now"
+            color="accent"
+            theme="fill"
+        />
+    </div>
+);
diff --git a/querybook/webapp/resource/github.ts b/querybook/webapp/resource/github.ts
new file mode 100644
index 000000000..678a0f816
--- /dev/null
+++ b/querybook/webapp/resource/github.ts
@@ -0,0 +1,11 @@
+import ds from 'lib/datasource';
+
+export interface IGitHubAuthResponse {
+    url: string;
+}
+
+export const GitHubResource = {
+    connectGithub: () => ds.fetch<IGitHubAuthResponse>('/github/auth/'),
+    isAuthenticated: () =>
+        ds.fetch<{ is_authenticated: boolean }>('/github/is_authenticated/'),
+};