Merge pull request #887 from Kaggle/add-session-client

Add new client to retrieve exportable .ipynb of current session.

Author: rosbo

Dockerfile

@@ -488,6 +488,7 @@ RUN echo "deb [signed-by=/usr/share/keyrings/cloud.google.gpg] http://packages.c
 ENV PYTHONUSERBASE "/root/.local"
 ADD patches/kaggle_gcp.py /root/.local/lib/python3.7/site-packages/kaggle_gcp.py
 ADD patches/kaggle_secrets.py /root/.local/lib/python3.7/site-packages/kaggle_secrets.py
+ADD patches/kaggle_session.py /root/.local/lib/python3.7/site-packages/kaggle_session.py
 ADD patches/kaggle_web_client.py /root/.local/lib/python3.7/site-packages/kaggle_web_client.py
 ADD patches/kaggle_datasets.py /root/.local/lib/python3.7/site-packages/kaggle_datasets.py
 ADD patches/log.py /root/.local/lib/python3.7/site-packages/log.py

patches/kaggle_session.py

@@ -0,0 +1,26 @@
+"""
+This library adds support for retrieving data related to the current user session.
+"""
+
+import os
+
+from kaggle_web_client import KaggleWebClient
+
+
+class UserSessionClient():
+    GET_SOURCE_ENDPOINT = '/requests/GetKernelRunSourceForCaipRequest'
+
+    def __init__(self):
+        self.web_client = KaggleWebClient()
+
+    def get_exportable_ipynb(self):
+        """Fetch the .ipynb source of the current notebook session.
+        
+        If Kaggle datasets are attached to the notebook, the source will
+        include an additonnal cell with logic to download the datasets
+        outside the Kaggle platform.
+        """
+        request_body = {
+            'UseDraft': True,        
+        }
+        return self.web_client.make_post_request(request_body, self.GET_SOURCE_ENDPOINT)

tests/test_user_session.py

@@ -0,0 +1,104 @@
+import json
+import os
+import threading
+import unittest
+
+from http.server import BaseHTTPRequestHandler, HTTPServer
+from test.support import EnvironmentVarGuard
+from urllib.parse import urlparse
+from kaggle_session import UserSessionClient
+from kaggle_web_client import (_KAGGLE_URL_BASE_ENV_VAR_NAME,
+                            _KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME,
+                            CredentialError, BackendError)
+
+class UserSessionHTTPHandler(BaseHTTPRequestHandler):
+
+    def set_request(self):
+        raise NotImplementedError()
+
+    def get_response(self):
+        raise NotImplementedError()
+
+    def do_HEAD(s):
+        s.send_response(200)
+
+    def do_POST(s):
+        s.set_request()
+        s.send_response(200)
+        s.send_header("Content-type", "application/json")
+        s.end_headers()
+        s.wfile.write(json.dumps(s.get_response()).encode("utf-8"))
+
+class TestUserSessionClient(unittest.TestCase):
+    SERVER_ADDRESS = urlparse(os.getenv(_KAGGLE_URL_BASE_ENV_VAR_NAME, default="http://127.0.0.1:8001"))
+    TEST_JWT = 'test-secrets-key'
+
+    def _test_client(self, client_func, expected_path, expected_body, source=None, success=True):
+        _request = {}
+
+        class GetKernelRunSourceForCaipHandler(UserSessionHTTPHandler):
+            def set_request(self):
+                _request['path'] = self.path
+                content_len = int(self.headers.get('Content-Length'))
+                _request['body'] = json.loads(self.rfile.read(content_len))
+                _request['headers'] = self.headers
+            
+            def get_response(self):
+                if success:
+                    return {'result': {'source': source}, 'wasSuccessful': 'true'}
+                return {'wasSuccessful': 'false'}
+        
+        env = EnvironmentVarGuard()
+        env.set(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME, self.TEST_JWT)
+        with env:
+            with HTTPServer((self.SERVER_ADDRESS.hostname, self.SERVER_ADDRESS.port), GetKernelRunSourceForCaipHandler) as httpd:
+                threading.Thread(target=httpd.serve_forever).start()
+
+                try:
+                    client_func()
+                finally:
+                    httpd.shutdown()
+                
+                path, headers, body = _request['path'], _request['headers'], _request['body']
+                self.assertEqual(
+                    path,
+                    expected_path,
+                    msg="Fake server did not receive the right request from UserSessionClient.")
+                self.assertEqual(
+                    body,
+                    expected_body,
+                    msg="Fake server did not receive the right body from UserSessionClient.")
+
+    def test_no_token_fails(self):
+        env = EnvironmentVarGuard()
+        env.unset(_KAGGLE_USER_SECRETS_TOKEN_ENV_VAR_NAME)
+        with env:
+            with self.assertRaises(CredentialError):
+                client = UserSessionClient()
+    
+    def test_get_exportable_ipynb_succeeds(self):
+        source = "import foo"
+
+        def call_get_ipynb():
+            client = UserSessionClient()
+            response = client.get_exportable_ipynb()
+            self.assertEqual(source, response['source'])
+
+        self._test_client(
+            call_get_ipynb,
+            '/requests/GetKernelRunSourceForCaipRequest',
+            {'UseDraft': True},
+            source=source,
+            success=True)
+
+    def test_get_exportable_ipynb_fails(self):
+        def call_get_ipynb():
+            client = UserSessionClient()
+            with self.assertRaises(BackendError):
+                client.get_exportable_ipynb()
+
+        self._test_client(
+            call_get_ipynb,
+            '/requests/GetKernelRunSourceForCaipRequest',
+            {'UseDraft': True},
+            success=False)