feat: Add oauth flow for querybook github integration

querybook/server/datasources/__init__.py

@@ -18,6 +18,8 @@
 from . import comment
 from . import survey
 from . import query_transform
+from . import github
+
 
 # Keep this at the end of imports to make sure the plugin APIs override the default ones
 try:
@@ -47,3 +49,4 @@
 survey
 query_transform
 api_plugin
+github

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}

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()

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}
+                />
+            )}
+        </>
+    );
+};

querybook/webapp/components/DataDocGitHub/GitHub.scss

@@ -0,0 +1,8 @@
+.GitHubAuth {
+    text-align: center;
+    padding: 20px;
+}
+
+.GitHubAuth-icon {
+    margin-bottom: 20px;
+}

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>
+);

querybook/webapp/components/DataDocRightSidebar/DataDocRightSidebar.tsx

@@ -3,6 +3,7 @@ import { useDispatch, useSelector } from 'react-redux';
 
 import { DataDocBoardsButton } from 'components/DataDocBoardsButton/DataDocBoardsButton';
 import { DataDocDAGExporterButton } from 'components/DataDocDAGExporter/DataDocDAGExporterButton';
+import { DataDocGitHubButton } from 'components/DataDocGitHub/DataDocGitHubButton';
 import { DataDocTemplateButton } from 'components/DataDocTemplateButton/DataDocTemplateButton';
 import { DataDocUIGuide } from 'components/UIGuide/DataDocUIGuide';
 import { ComponentType, ElementType } from 'const/analytics';
@@ -83,6 +84,8 @@ export const DataDocRightSidebar: React.FunctionComponent<IProps> = ({
         <DataDocRunAllButton docId={dataDoc.id} />
     );
 
+    const githubButtonDOM = <DataDocGitHubButton docId={dataDoc.id} />;
+
     const buttonSection = (
         <div className="DataDocRightSidebar-button-section vertical-space-between">
             <div className="DataDocRightSidebar-button-section-top flex-column">
@@ -131,6 +134,7 @@ export const DataDocRightSidebar: React.FunctionComponent<IProps> = ({
             </div>
             <div className="DataDocRightSidebar-button-section-bottom flex-column mb8">
                 {runAllButtonDOM}
+                {githubButtonDOM}
                 {isEditable && exporterExists && (
                     <DataDocDAGExporterButton docId={dataDoc.id} />
                 )}

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/'),
+};

querybook/webapp/ui/Icon/LucideIcons.ts

@@ -52,6 +52,7 @@ import {
     FileText,
     Filter,
     FormInput,
+    Github,
     GripVertical,
     Hash,
     HelpCircle,
@@ -167,6 +168,7 @@ const AllLucideIcons = {
     FileText,
     Filter,
     FormInput,
+    Github,
     GripVertical,
     Hash,
     HelpCircle,

webpack.config.js

@@ -48,6 +48,10 @@ function getDevServerSettings(env) {
                 target: QUERYBOOK_UPSTREAM,
                 changeOrigin: true,
             },
+            '/github/oauth2callback': {
+                target: QUERYBOOK_UPSTREAM,
+                changeOrigin: true,
+            },
         },
         publicPath: '/build/',
         onListening: (server) => {