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/DataDocGitHub/GitHubModal.tsx

@@ -0,0 +1,82 @@
+import React, { useCallback, useState } from 'react';
+
+import { ComponentType, ElementType } from 'const/analytics';
+import { trackClick } from 'lib/analytics';
+import { GitHubResource, IGitHubAuthResponse } from 'resource/github';
+import { Message } from 'ui/Message/Message';
+import { Modal } from 'ui/Modal/Modal';
+
+import { GitHubAuth } from './GitHubAuth';
+
+interface IProps {
+    docId: number;
+    isAuthenticated: boolean;
+    setIsAuthenticated: (isAuthenticated: boolean) => void;
+    onClose: () => void;
+}
+
+export const GitHubModal: React.FunctionComponent<IProps> = ({
+    docId,
+    isAuthenticated,
+    setIsAuthenticated,
+    onClose,
+}) => {
+    const [errorMessage, setErrorMessage] = useState<string>(null);
+
+    const handleConnectGitHub = useCallback(async () => {
+        trackClick({
+            component: ComponentType.DATADOC_PAGE,
+            element: ElementType.GITHUB_CONNECT_BUTTON,
+        });
+
+        try {
+            const { data }: { data: IGitHubAuthResponse } =
+                await GitHubResource.connectGithub();
+            const url = data.url;
+            if (!url) {
+                throw new Error('Failed to get GitHub authentication URL');
+            }
+            const authWindow = window.open(url);
+
+            const receiveMessage = () => {
+                authWindow.close();
+                delete window.receiveChildMessage;
+                window.removeEventListener('message', receiveMessage, false);
+                setIsAuthenticated(true);
+            };
+            window.receiveChildMessage = receiveMessage;
+
+            // If the user closes the authentication window manually, clean up
+            const timer = setInterval(() => {
+                if (authWindow.closed) {
+                    clearInterval(timer);
+                    window.removeEventListener(
+                        'message',
+                        receiveMessage,
+                        false
+                    );
+                    throw new Error('Authentication process failed');
+                }
+            }, 1000);
+        } catch (error) {
+            console.error('GitHub authentication failed:', error);
+            setErrorMessage('GitHub authentication failed. Please try again.');
+        }
+    }, [setIsAuthenticated]);
+
+    return (
+        <Modal onHide={onClose} title="GitHub Integration">
+            <div className="GitHubModal-content">
+                {isAuthenticated ? (
+                    <Message message="Connected to GitHub!" type="success" />
+                ) : (
+                    <GitHubAuth onAuthenticate={handleConnectGitHub} />
+                )}
+                {errorMessage && (
+                    <Message message={errorMessage} type="error" />
+                )}
+                <button onClick={onClose}>Close</button>
+            </div>
+        </Modal>
+    );
+};

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/const/analytics.ts

@@ -115,6 +115,10 @@ export enum ElementType {
     QUERY_GENERATION_REJECT_BUTTON = 'QUERY_GENERATION_REJECT_BUTTON',
     QUERY_GENERATION_APPLY_BUTTON = 'QUERY_GENERATION_APPLY_BUTTON',
     QUERY_GENERATION_APPLY_AND_RUN_BUTTON = 'QUERY_GENERATION_APPLY_AND_RUN_BUTTON',
+
+    // Github Integration
+    GITHUB_CONNECT_BUTTON = 'GITHUB_CONNECT_BUTTON',
+    GITHUB_LINK_BUTTON = 'GITHUB_LINK_BUTTON',
 }
 
 export interface EventData {

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,