Add support for cross-account drives

jupyter_drives/handlers.py

@@ -101,8 +101,8 @@ async def post(self, drive: str = "", path: str = ""):
         body = self.get_json_body()
         if 'location' in body:
             result = await self._manager.new_drive(drive, **body)
-        elif 'public' in body:
-            result = await self._manager.add_public_drive(drive)
+        elif 'is_public' in body:
+            result = await self._manager.add_external_drive(drive, **body)
         else:
             result = await self._manager.new_file(drive, path, **body)
         self.finish(result)

jupyter_drives/manager.py

@@ -311,7 +311,10 @@ async def mount_drive(self, drive_name, provider):
         """
         try:
             if provider == 's3':
-                region = await self._get_drive_location(drive_name)
+                if drive_name in self._external_drives and self._external_drives[drive_name]["is_public"] is False: 
+                    region = self._external_drives[drive_name]["location"]
+                else:
+                    region = await self._get_drive_location(drive_name)
             self._initialize_content_manager(drive_name, provider, region)
         except Exception as e:
             raise tornado.web.HTTPError(
@@ -731,16 +734,17 @@ async def new_drive(self, new_drive_name, location):
 
         return
 
-    async def add_public_drive(self, drive_name):
+    async def add_external_drive(self, drive_name, is_public, region='us-east-1'):
         """Mount a drive.
 
         Args:
             drive_name: name of public bucket to mount
         """
         try:
             drive = {
-                "is_public": True,
-                "url": drive_name
+                "is_public": is_public,
+                "url": drive_name,
+                "location": region
             };
             self._external_drives[drive_name] = drive;
         except Exception as e:

src/contents.ts

@@ -24,7 +24,8 @@ import {
   createDrive,
   getDrivesList,
   excludeDrive,
-  includeDrive
+  includeDrive,
+  addExternalDrive
 } from './requests';
 
 export class Drive implements Contents.IDrive {
@@ -850,6 +851,42 @@ export class Drive implements Contents.IDrive {
     return data;
   }
 
+  /**
+   * Add external drive.
+   *
+   * @param options: The options used to add the external drive.
+   *
+   * @returns A promise which resolves with the contents model.
+   */
+  async addExternalDrive(
+    driveUrl: string,
+    location: string
+  ): Promise<Contents.IModel> {
+    await addExternalDrive(driveUrl, location);
+
+    const data: Contents.IModel = {
+      name: driveUrl,
+      path: driveUrl,
+      last_modified: '',
+      created: '',
+      content: [],
+      format: 'json',
+      mimetype: '',
+      size: 0,
+      writable: true,
+      type: 'directory'
+    };
+
+    Contents.validateContentsModel(data);
+    this._fileChanged.emit({
+      type: 'new',
+      oldValue: null,
+      newValue: data
+    });
+
+    return data;
+  }
+
   /**
    * Exclude drive from browser.
    *

src/plugins/driveBrowserPlugin.ts

@@ -454,6 +454,38 @@ namespace Private {
       rank: 110
     });
 
+    app.commands.addCommand(CommandIDs.addExternalDrive, {
+      isVisible: () => {
+        return browser.model.path === 's3:';
+      },
+      execute: async () => {
+        return showDialog({
+          title: 'Add External Drive',
+          body: new Private.CreateDriveHandler(drive.name),
+          focusNodeSelector: 'input',
+          buttons: [
+            Dialog.cancelButton(),
+            Dialog.okButton({
+              label: 'Add',
+              ariaLabel: 'Add Drive'
+            })
+          ]
+        }).then(result => {
+          if (result.value) {
+            drive.addExternalDrive(result.value[0], result.value[1]);
+          }
+        });
+      },
+      label: 'Add External Drive',
+      icon: driveBrowserIcon.bindprops({ stylesheet: 'menuItem' })
+    });
+
+    app.contextMenu.addItem({
+      command: CommandIDs.addExternalDrive,
+      selector: '#drive-file-browser.jp-SidePanel .jp-DirListing-content',
+      rank: 110
+    });
+
     app.commands.addCommand(CommandIDs.toggleFileFilter, {
       execute: () => {
         // Update toggled state, then let the toolbar button update

src/plugins/drivelistmanager.tsx

@@ -4,6 +4,7 @@ import { Button, Search } from '@jupyter/react-components';
 import { useState } from 'react';
 import { IDriveInfo } from '../token';
 import {
+  addExternalDrive,
   addPublicDrive,
   excludeDrive,
   getDrivesList,
@@ -19,15 +20,23 @@ interface IProps {
 
 export interface IDriveInputProps {
   isName: boolean;
-  value: string;
+  driveValue: string;
+  regionValue: string;
   setPublicDrive: (value: string) => void;
+  setRegion: (value: string) => void;
   onSubmit: () => void;
+  isPublic: boolean;
+  setIsPublic: (value: boolean) => void;
 }
 
 export function DriveInputComponent({
-  value,
+  driveValue,
+  regionValue,
   setPublicDrive,
-  onSubmit
+  setRegion,
+  onSubmit,
+  isPublic,
+  setIsPublic
 }: IDriveInputProps) {
   return (
     <div>
@@ -38,7 +47,7 @@ export function DriveInputComponent({
             setPublicDrive(event.target.value);
           }}
           placeholder="Enter drive name"
-          value={value}
+          value={driveValue}
         />
         <Button
           className="input-add-drive-button"
@@ -48,6 +57,24 @@ export function DriveInputComponent({
           add
         </Button>
       </div>
+      <div className="add-public-drive-section">
+        {'Public drive?'}
+        <input
+          type="checkbox"
+          checked={isPublic}
+          onChange={event => setIsPublic(event.target.checked)}
+        />
+        {!isPublic && (
+          <input
+            className="drive-region-input"
+            onInput={(event: any) => {
+              setRegion(event.target.value);
+            }}
+            placeholder="Region (e.g.: us-east-1)"
+            value={regionValue}
+          />
+        )}
+      </div>
     </div>
   );
 }
@@ -154,6 +181,8 @@ export function DriveListManagerComponent({ model }: IProps) {
   const [availableDrives, setAvailableDrives] = useState<Partial<IDriveInfo>[]>(
     model.availableDrives
   );
+  const [isPublic, setIsPublic] = useState<boolean>(false);
+  const [driveRegion, setDriveRegion] = useState<string>('');
 
   // Called after mounting.
   React.useEffect(() => {
@@ -168,7 +197,12 @@ export function DriveListManagerComponent({ model }: IProps) {
   }, [model]);
 
   const onAddedPublicDrive = async () => {
-    await addPublicDrive(publicDrive);
+    if (isPublic) {
+      await addPublicDrive(publicDrive);
+    } else {
+      await addExternalDrive(publicDrive, driveRegion);
+      setDriveRegion('');
+    }
     setPublicDrive('');
     await model.refresh();
   };
@@ -195,11 +229,15 @@ export function DriveListManagerComponent({ model }: IProps) {
         </div>
 
         <div className="drives-manager-section">
-          <div className="drives-section-title"> Add public drive</div>
+          <div className="drives-section-title"> Add external drive</div>
           <DriveInputComponent
             isName={false}
-            value={publicDrive}
+            driveValue={publicDrive}
+            regionValue={driveRegion}
             setPublicDrive={setPublicDrive}
+            setRegion={setDriveRegion}
+            isPublic={isPublic}
+            setIsPublic={setIsPublic}
             onSubmit={onAddedPublicDrive}
           />
         </div>

src/requests.ts

@@ -509,7 +509,22 @@ export async function createDrive(
  */
 export async function addPublicDrive(driveUrl: string) {
   return await requestAPI<any>('drives/' + driveUrl + '/', 'POST', {
-    public: true
+    is_public: true
+  });
+}
+
+/**
+ * Add external drive.
+ *
+ * @param driveUrl The drive URL.
+ * @param location The drive region.
+ *
+ * @returns A promise which resolves with the contents model.
+ */
+export async function addExternalDrive(driveUrl: string, location: string) {
+  return await requestAPI<any>('drives/' + driveUrl + '/', 'POST', {
+    is_public: false,
+    region: location
   });
 }
 

src/token.ts

@@ -9,6 +9,7 @@ export namespace CommandIDs {
   export const toggleBrowser = 'drives:toggle-main';
   export const createNewDrive = 'drives:create-new-drive';
   export const addPublicDrive = 'drives:add-public-drive';
+  export const addExternalDrive = 'drives:add-external-drive';
   export const launcher = 'launcher:create';
   export const toggleFileFilter = 'drives:toggle-file-filter';
   export const createNewDirectory = 'drives:create-new-directory';

style/base.css

@@ -50,6 +50,14 @@ li {
   justify-content: center;
 }
 
+.drive-region-input {
+  height: 20px !important;
+  width: 100%;
+  flex: 1;
+  justify-content: center;
+  margin-right: 61px;
+}
+
 .drive-data-grid {
   width: 380px;
   flex-direction: row;