singlestoredb 1.16.1__py3-none-any.whl

singlestoredb/tests/test_fusion.py

@@ -0,0 +1,1527 @@
+#!/usr/bin/env python
+# type: ignore
+"""SingleStoreDB Fusion testing."""
+import os
+import random
+import secrets
+import tempfile
+import time
+import unittest
+from typing import Any
+from typing import List
+
+import pytest
+
+import singlestoredb as s2
+from singlestoredb.tests import utils
+
+
+class TestFusion(unittest.TestCase):
+
+    dbname: str = ''
+    dbexisted: bool = False
+
+    @classmethod
+    def setUpClass(cls):
+        sql_file = os.path.join(os.path.dirname(__file__), 'test.sql')
+        cls.dbname, cls.dbexisted = utils.load_sql(sql_file)
+
+    @classmethod
+    def tearDownClass(cls):
+        if not cls.dbexisted:
+            utils.drop_database(cls.dbname)
+
+    def setUp(self):
+        self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED')
+        os.environ['SINGLESTOREDB_FUSION_ENABLED'] = '1'
+        self.conn = s2.connect(database=type(self).dbname, local_infile=True)
+        self.cur = self.conn.cursor()
+
+    def tearDown(self):
+        if self.enabled:
+            os.environ['SINGLESTOREDB_FUSION_ENABLED'] = self.enabled
+        else:
+            del os.environ['SINGLESTOREDB_FUSION_ENABLED']
+
+        try:
+            if self.cur is not None:
+                self.cur.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+        try:
+            if self.conn is not None:
+                self.conn.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+    def test_env_var(self):
+        os.environ['SINGLESTOREDB_FUSION_ENABLED'] = '0'
+
+        with self.assertRaises(s2.ProgrammingError):
+            self.cur.execute('show fusion commands')
+
+        del os.environ['SINGLESTOREDB_FUSION_ENABLED']
+
+        with self.assertRaises(s2.ProgrammingError):
+            self.cur.execute('show fusion commands')
+
+        os.environ['SINGLESTOREDB_FUSION_ENABLED'] = 'yes'
+
+        self.cur.execute('show fusion commands')
+        assert list(self.cur)
+
+    def test_show_commands(self):
+        self.cur.execute('show fusion commands')
+        cmds = [x[0] for x in self.cur.fetchall()]
+        assert cmds
+        assert [x for x in cmds if x.strip().startswith('SHOW FUSION GRAMMAR')], cmds
+
+        self.cur.execute('show fusion commands like "create%"')
+        cmds = [x[0] for x in self.cur.fetchall()]
+        assert cmds
+        assert [x for x in cmds if x.strip().startswith('CREATE')] == cmds, cmds
+
+    def test_show_grammar(self):
+        self.cur.execute('show fusion grammar for "create workspace"')
+        cmds = [x[0] for x in self.cur.fetchall()]
+        assert cmds
+        assert [x for x in cmds if x.strip().startswith('CREATE WORKSPACE')], cmds
+
+
+@pytest.mark.management
+class TestWorkspaceFusion(unittest.TestCase):
+
+    id: str = secrets.token_hex(8)
+    dbname: str = ''
+    dbexisted: bool = False
+    workspace_groups: List[Any] = []
+
+    @classmethod
+    def setUpClass(cls):
+        sql_file = os.path.join(os.path.dirname(__file__), 'test.sql')
+        cls.dbname, cls.dbexisted = utils.load_sql(sql_file)
+        mgr = s2.manage_workspaces()
+        us_regions = [x for x in mgr.regions if x.name.startswith('US')]
+        non_us_regions = [x for x in mgr.regions if not x.name.startswith('US')]
+        wg = mgr.create_workspace_group(
+            f'A Fusion Testing {cls.id}',
+            region=random.choice(us_regions),
+            firewall_ranges=[],
+        )
+        cls.workspace_groups.append(wg)
+        wg = mgr.create_workspace_group(
+            f'B Fusion Testing {cls.id}',
+            region=random.choice(us_regions),
+            firewall_ranges=[],
+        )
+        cls.workspace_groups.append(wg)
+        wg = mgr.create_workspace_group(
+            f'C Fusion Testing {cls.id}',
+            region=random.choice(non_us_regions),
+            firewall_ranges=[],
+        )
+        cls.workspace_groups.append(wg)
+
+    @classmethod
+    def tearDownClass(cls):
+        if not cls.dbexisted:
+            utils.drop_database(cls.dbname)
+        while cls.workspace_groups:
+            cls.workspace_groups.pop().terminate(force=True)
+
+    def setUp(self):
+        self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED')
+        os.environ['SINGLESTOREDB_FUSION_ENABLED'] = '1'
+        self.conn = s2.connect(database=type(self).dbname, local_infile=True)
+        self.cur = self.conn.cursor()
+
+    def tearDown(self):
+        if self.enabled:
+            os.environ['SINGLESTOREDB_FUSION_ENABLED'] = self.enabled
+        else:
+            del os.environ['SINGLESTOREDB_FUSION_ENABLED']
+
+        try:
+            if self.cur is not None:
+                self.cur.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+        try:
+            if self.conn is not None:
+                self.conn.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+    def test_show_regions(self):
+        self.cur.execute('show regions')
+        regs = list(self.cur)
+        desc = self.cur.description
+
+        us_regs = [x for x in regs if x[0].startswith('US')]
+
+        assert len(desc) == 3
+        assert len(regs) > 5
+        assert len(us_regs) > 5
+
+        # LIKE
+        self.cur.execute('show regions like "US%"')
+        regs = list(self.cur)
+        assert regs == us_regs
+
+        # LIMIT
+        self.cur.execute('show regions like "US%" limit 3')
+        regs = list(self.cur)
+        assert len(regs) == 3
+
+        # ORDER BY
+        self.cur.execute('show regions like "US%" limit 3 order by name')
+        regs = list(self.cur)
+        assert len(regs) == 3
+        assert regs == list(sorted(regs, key=lambda x: x[0]))
+
+        # Wrong column
+        with self.assertRaises(KeyError):
+            self.cur.execute('show regions like "US%" limit 3 order by foo')
+
+    def test_show_workspace_groups(self):
+        self.cur.execute('show workspace groups')
+        wgs = list(self.cur)
+        desc = self.cur.description
+
+        assert len(desc) == 4
+        assert desc[0].name == 'Name'
+        assert desc[1].name == 'ID'
+        assert desc[2].name == 'Region'
+        assert desc[3].name == 'FirewallRanges'
+        assert len(wgs) >= 3
+
+        names = [x[0] for x in wgs]
+        assert f'A Fusion Testing {self.id}' in names
+        assert f'B Fusion Testing {self.id}' in names
+        assert f'C Fusion Testing {self.id}' in names
+
+        # LIKE clause
+        self.cur.execute(f'show workspace groups like "A%sion Testing {self.id}"')
+        wgs = list(self.cur)
+
+        names = [x[0] for x in wgs]
+        assert f'A Fusion Testing {self.id}' in names
+        assert f'B Fusion Testing {self.id}' not in names
+        assert f'C Fusion Testing {self.id}' not in names
+
+        # LIMIT clause
+        self.cur.execute('show workspace groups limit 2')
+        wgs = list(self.cur)
+        assert len(wgs) == 2
+
+        # EXTENDED attributes
+        self.cur.execute('show workspace groups extended')
+        wgs = list(self.cur)
+        desc = self.cur.description
+
+        assert len(desc) == 6
+        assert desc[4].name == 'CreatedAt'
+        assert desc[5].name == 'TerminatedAt'
+
+        # ORDER BY
+        self.cur.execute(
+            f'show workspace groups like "% Fusion Testing {self.id}" order by name desc',
+        )
+        wgs = list(self.cur)
+
+        names = [x[0] for x in wgs]
+        assert names == [
+            f'C Fusion Testing {self.id}',
+            f'B Fusion Testing {self.id}',
+            f'A Fusion Testing {self.id}',
+        ]
+
+        # All options
+        self.cur.execute(
+            f'show workspace groups like "% Fusion Testing {self.id}" '
+            'extended order by name desc limit 2',
+        )
+        wgs = list(self.cur)
+        desc = self.cur.description
+        names = [x[0] for x in wgs]
+
+        assert len(desc) == 6
+        assert names == [f'C Fusion Testing {self.id}', f'B Fusion Testing {self.id}']
+
+    def test_show_workspaces(self):
+        mgr = s2.manage_workspaces()
+        wg = mgr.workspace_groups[f'B Fusion Testing {self.id}']
+
+        self.cur.execute(
+            'create workspace show-ws-1 in group '
+            f'"B Fusion Testing {self.id}" with size S-00',
+        )
+        self.cur.execute(
+            'create workspace show-ws-2 in group '
+            f'"B Fusion Testing {self.id}" with size S-00',
+        )
+        self.cur.execute(
+            'create workspace show-ws-3 in group '
+            f'"B Fusion Testing {self.id}" with size S-00',
+        )
+
+        time.sleep(30)
+        iterations = 20
+        while True:
+            wgs = wg.workspaces
+            states = [
+                x.state for x in wgs
+                if x.name in ('show-ws-1', 'show-ws-2', 'show-ws-3')
+            ]
+            if len(states) == 3 and states.count('ACTIVE') == 3:
+                break
+            iterations -= 1
+            if not iterations:
+                raise RuntimeError('timed out waiting for workspaces to start')
+            time.sleep(30)
+
+        # SHOW
+        self.cur.execute(f'show workspaces in group "B Fusion Testing {self.id}"')
+        desc = self.cur.description
+        out = list(self.cur)
+        names = [x[0] for x in out]
+        assert len(desc) == 4
+        assert [x[0] for x in desc] == ['Name', 'ID', 'Size', 'State']
+        assert len(out) >= 3
+        assert 'show-ws-1' in names
+        assert 'show-ws-2' in names
+        assert 'show-ws-3' in names
+
+        # SHOW ID
+        self.cur.execute(f'show workspaces in group id {wg.id}')
+        desc = self.cur.description
+        out = list(self.cur)
+        names = [x[0] for x in out]
+        assert len(desc) == 4
+        assert [x[0] for x in desc] == ['Name', 'ID', 'Size', 'State']
+        assert len(out) >= 3
+        assert 'show-ws-1' in names
+        assert 'show-ws-2' in names
+        assert 'show-ws-3' in names
+
+        # LIKE clause
+        self.cur.execute(
+            'show workspaces in group '
+            f'"B Fusion Testing {self.id}" like "%2"',
+        )
+        out = list(self.cur)
+        names = [x[0] for x in out]
+        assert len(out) >= 1
+        assert [x for x in names if x.endswith('2')]
+        assert 'show-ws-1' not in names
+        assert 'show-ws-2' in names
+        assert 'show-ws-3' not in names
+
+        # Extended attributes
+        self.cur.execute(
+            'show workspaces in group '
+            f'"B Fusion Testing {self.id}" extended',
+        )
+        desc = self.cur.description
+        out = list(self.cur)
+        assert len(desc) == 7
+        assert [x[0] for x in desc] == [
+            'Name', 'ID', 'Size', 'State',
+            'Endpoint', 'CreatedAt', 'TerminatedAt',
+        ]
+
+        # ORDER BY
+        self.cur.execute(
+            'show workspaces in group '
+            f'"B Fusion Testing {self.id}" order by name desc',
+        )
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 4
+        names = [x[0] for x in out]
+        assert names == ['show-ws-3', 'show-ws-2', 'show-ws-1']
+
+        # LIMIT clause
+        self.cur.execute(
+            'show workspaces in group '
+            f'"B Fusion Testing {self.id}" order by name desc limit 2',
+        )
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 4
+        names = [x[0] for x in out]
+        assert names == ['show-ws-3', 'show-ws-2']
+
+        # All options
+        self.cur.execute(
+            f'show workspaces in group "B Fusion Testing {self.id}" '
+            'like "show-ws%" extended order by name desc limit 2',
+        )
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 7
+        names = [x[0] for x in out]
+        assert names == ['show-ws-3', 'show-ws-2']
+
+    def test_create_drop_workspace(self):
+        mgr = s2.manage_workspaces()
+        wg = mgr.workspace_groups[f'A Fusion Testing {self.id}']
+
+        self.cur.execute(
+            f'create workspace foobar-1 in group "A Fusion Testing {self.id}" '
+            'with size S-00 wait on active',
+        )
+        foobar_1 = [x for x in wg.workspaces if x.name == 'foobar-1']
+        assert len(foobar_1) == 1
+
+        self.cur.execute(
+            f'create workspace foobar-2 in group "A Fusion Testing {self.id}" '
+            'with size S-00 wait on active',
+        )
+        foobar_2 = [x for x in wg.workspaces if x.name == 'foobar-2']
+        assert len(foobar_2) == 1
+
+        # Drop by name
+        self.cur.execute(
+            f'drop workspace "foobar-1" in group "A Fusion Testing {self.id}" '
+            'wait on terminated',
+        )
+        foobar_1 = [x for x in wg.workspaces if x.name == 'foobar-1']
+        assert len(foobar_1) == 0
+
+        # Drop by ID
+        foobar_2_id = foobar_2[0].id
+        self.cur.execute(
+            f'drop workspace id {foobar_2_id} in group '
+            f'"A Fusion Testing {self.id}" wait on terminated',
+        )
+        foobar_2 = [x for x in wg.workspaces if x.name == 'foobar-2']
+        assert len(foobar_2) == 0
+
+        # Drop non-existent by ID
+        with self.assertRaises(KeyError):
+            self.cur.execute(
+                f'drop workspace id {foobar_2_id} '
+                f'in group "A Fusion Testing {self.id}"',
+            )
+
+        # Drop non-existent by ID with IF EXISTS
+        self.cur.execute(
+            f'drop workspace IF EXISTS id {foobar_2_id} '
+            f'in group "A Fusion Testing {self.id}"',
+        )
+
+    def test_create_drop_workspace_group(self):
+        mgr = s2.manage_workspaces()
+
+        reg = [x for x in mgr.regions if x.name.startswith('US')][0]
+        wg_name = f'Create WG Test {id(self)}'
+
+        try:
+            self.cur.execute(
+                f'create workspace group "{wg_name}" '
+                f'in region "{reg.name}"',
+            )
+            wg = [x for x in mgr.workspace_groups if x.name == wg_name]
+            assert len(wg) == 1
+
+            # Drop it by name
+            self.cur.execute(
+                f'drop workspace group "{wg_name}" '
+                'wait on terminated',
+            )
+            wg = [x for x in mgr.workspace_groups if x.name == wg_name]
+            assert len(wg) == 0
+
+            # Create it again
+            self.cur.execute(
+                f'create workspace group "{wg_name}" in region "{reg.name}"',
+            )
+            wg = [x for x in mgr.workspace_groups if x.name == wg_name]
+            assert len(wg) == 1
+
+            # Drop it by ID
+            wg_id = wg[0].id
+            self.cur.execute(f'drop workspace group id {wg_id} wait on terminated')
+            wg = [x for x in mgr.workspace_groups if x.name == wg_name]
+            assert len(wg) == 0
+
+            # Drop non-existent
+            with self.assertRaises(KeyError):
+                self.cur.execute(f'drop workspace group id {wg_id}')
+
+            # Drop non-existent with IF EXISTS
+            self.cur.execute(f'drop workspace group if exists id {wg_id}')
+
+        finally:
+            try:
+                mgr.workspace_groups[wg_name].terminate(force=True)
+            except Exception:
+                pass
+
+
+@pytest.mark.management
+class TestJobsFusion(unittest.TestCase):
+
+    id: str = secrets.token_hex(8)
+    notebook_name: str = 'Scheduling Test.ipynb'
+    dbname: str = ''
+    dbexisted: bool = False
+    manager: None
+    workspace_group: None
+    workspace: None
+    job_ids = []
+
+    @classmethod
+    def setUpClass(cls):
+        sql_file = os.path.join(os.path.dirname(__file__), 'test.sql')
+        cls.dbname, cls.dbexisted = utils.load_sql(sql_file)
+        cls.manager = s2.manage_workspaces()
+        us_regions = [x for x in cls.manager.regions if x.name.startswith('US')]
+        cls.workspace_group = cls.manager.create_workspace_group(
+            f'Jobs Fusion Testing {cls.id}',
+            region=random.choice(us_regions),
+            firewall_ranges=[],
+        )
+        cls.workspace = cls.workspace_group.create_workspace(
+                f'jobs-test-{cls.id}',
+                wait_on_active=True,
+        )
+        os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] = cls.dbname
+        os.environ['SINGLESTOREDB_WORKSPACE'] = cls.workspace.id
+
+    @classmethod
+    def tearDownClass(cls):
+        for job_id in cls.job_ids:
+            try:
+                cls.manager.organizations.current.jobs.delete(job_id)
+            except Exception:
+                pass
+        if cls.workspace_group is not None:
+            cls.workspace_group.terminate(force=True)
+        cls.manager = None
+        cls.workspace_group = None
+        cls.workspace = None
+        if os.environ.get('SINGLESTOREDB_WORKSPACE', None) is not None:
+            del os.environ['SINGLESTOREDB_WORKSPACE']
+        if os.environ.get('SINGLESTOREDB_DEFAULT_DATABASE', None) is not None:
+            del os.environ['SINGLESTOREDB_DEFAULT_DATABASE']
+
+    def setUp(self):
+        self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED')
+        os.environ['SINGLESTOREDB_FUSION_ENABLED'] = '1'
+        self.conn = s2.connect(database=type(self).dbname, local_infile=True)
+        self.cur = self.conn.cursor()
+
+    def tearDown(self):
+        if self.enabled:
+            os.environ['SINGLESTOREDB_FUSION_ENABLED'] = self.enabled
+        else:
+            del os.environ['SINGLESTOREDB_FUSION_ENABLED']
+
+        try:
+            if self.cur is not None:
+                self.cur.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+        try:
+            if self.conn is not None:
+                self.conn.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+    def test_schedule_drop_job(self):
+        # schedule recurring job
+        self.cur.execute(
+                f'schedule job using notebook "{self.notebook_name}" '
+                'with mode "recurring" '
+                'execute every 5 minutes '
+                'with name "recurring-job" '
+                'create snapshot '
+                'resume target '
+                'with runtime "notebooks-cpu-small" '
+                'with parameters '
+                '{"strParam": "string", "intParam": 1, '
+                '"floatParam": 1.0, "boolParam": true}',
+        )
+        out = list(self.cur)
+        job_id = out[0][0]
+        self.job_ids.append(job_id)
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0][0] == 'JobID'
+        assert len(out) == 1
+        assert out[0][0] == job_id
+
+        # drop job
+        self.cur.execute(f'drop jobs {job_id}')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 2
+        assert [x[0] for x in desc] == [
+            'JobID', 'Success',
+        ]
+        assert len(out) == 1
+        res = out[0]
+        assert res[0] == job_id
+        assert res[1] == 1
+
+    def test_run_wait_drop_job(self):
+        # run job
+        self.cur.execute(
+            f'run job using notebook "{self.notebook_name}" '
+            'with runtime "notebooks-cpu-small" '
+            'with parameters '
+            '{"strParam": "string", "intParam": 1, '
+            '"floatParam": 1.0, "boolParam": true}',
+        )
+        out = list(self.cur)
+        job_id = out[0][0]
+        self.job_ids.append(job_id)
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0][0] == 'JobID'
+        assert len(out) == 1
+        assert out[0][0] == job_id
+
+        # wait on job
+        self.cur.execute(f'wait on jobs {job_id}')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0][0] == 'Success'
+        assert out[0][0] == 1
+
+        # drop job
+        self.cur.execute(f'drop jobs {job_id}')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 2
+        assert [x[0] for x in desc] == [
+            'JobID', 'Success',
+        ]
+        assert len(out) == 1
+        res = out[0]
+        assert res[0] == job_id
+        assert res[1] == 1
+
+    def test_show_jobs_and_executions(self):
+        # schedule recurring job
+        self.cur.execute(
+                f'schedule job using notebook "{self.notebook_name}" '
+                'with mode "recurring" '
+                'execute every 5 minutes '
+                'with name "show-job" '
+                'with runtime "notebooks-cpu-small" '
+                'with parameters '
+                '{"strParam": "string", "intParam": 1, '
+                '"floatParam": 1.0, "boolParam": true}',
+        )
+        out = list(self.cur)
+        job_id = out[0][0]
+        self.job_ids.append(job_id)
+        desc = self.cur.description
+        assert len(desc) == 1
+        assert desc[0][0] == 'JobID'
+        assert len(out) == 1
+        assert out[0][0] == job_id
+
+        # show jobs with name like "show-job"
+        self.cur.execute(f'show jobs {job_id} like "show-job"')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 9
+        assert [x[0] for x in desc] == [
+            'JobID', 'Name', 'CreatedAt', 'EnqueuedBy',
+            'CompletedExecutions', 'NotebookPath', 'DatabaseName', 'TargetID',
+            'TargetType',
+        ]
+        assert len(out) == 1
+        job = out[0]
+        assert job[0] == job_id
+        assert job[1] == 'show-job'
+        assert job[5] == self.notebook_name
+        assert job[6] == self.dbname
+        assert job[7] == self.workspace.id
+        assert job[8] == 'Workspace'
+
+        # show jobs with name like "show-job" extended
+        self.cur.execute(f'show jobs {job_id} like "show-job" extended')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 17
+        assert [x[0] for x in desc] == [
+            'JobID', 'Name', 'CreatedAt', 'EnqueuedBy',
+            'CompletedExecutions', 'NotebookPath', 'DatabaseName', 'TargetID',
+            'TargetType', 'Description', 'TerminatedAt', 'CreateSnapshot',
+            'MaxDurationInMins', 'ExecutionIntervalInMins', 'Mode', 'StartAt',
+            'ResumeTarget',
+        ]
+        assert len(out) == 1
+        job = out[0]
+        assert job[0] == job_id
+        assert job[1] == 'show-job'
+        assert job[5] == self.notebook_name
+        assert job[6] == self.dbname
+        assert job[7] == self.workspace.id
+        assert job[8] == 'Workspace'
+        assert not job[11]
+        assert job[13] == 5
+        assert job[14] == 'Recurring'
+        assert not job[16]
+
+        # show executions for job with id job_id from 1 to 5
+        self.cur.execute(f'show job executions for {job_id} from 1 to 5')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 7
+        assert [x[0] for x in desc] == [
+            'ExecutionID', 'ExecutionNumber', 'JobID',
+            'Status', 'ScheduledStartTime', 'StartedAt', 'FinishedAt',
+        ]
+        exec_job_ids = [x[2] for x in out]
+        for x in exec_job_ids:
+            assert x == job_id
+
+        # show executions for job with id job_id from 1 to 5 extended
+        self.cur.execute(f'show job executions for {job_id} from 1 to 5 extended')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 8
+        assert [x[0] for x in desc] == [
+            'ExecutionID', 'ExecutionNumber', 'JobID',
+            'Status', 'ScheduledStartTime', 'StartedAt', 'FinishedAt',
+            'SnapshotNotebookPath',
+        ]
+        exec_job_ids = [x[2] for x in out]
+        for x in exec_job_ids:
+            assert x == job_id
+
+        # drop job
+        self.cur.execute(f'drop jobs {job_id}')
+        out = list(self.cur)
+        desc = self.cur.description
+        assert len(desc) == 2
+        assert [x[0] for x in desc] == [
+            'JobID', 'Success',
+        ]
+        assert len(out) == 1
+        res = out[0]
+        assert res[0] == job_id
+        assert res[1] == 1
+
+
+@pytest.mark.management
+class TestStageFusion(unittest.TestCase):
+
+    id: str = secrets.token_hex(8)
+    dbname: str = 'information_schema'
+    manager: None
+    workspace_group: None
+    workspace_group_2: None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.manager = s2.manage_workspaces()
+        us_regions = [x for x in cls.manager.regions if x.name.startswith('US')]
+        cls.workspace_group = cls.manager.create_workspace_group(
+            f'Stage Fusion Testing 1 {cls.id}',
+            region=random.choice(us_regions),
+            firewall_ranges=[],
+        )
+        cls.workspace_group_2 = cls.manager.create_workspace_group(
+            f'Stage Fusion Testing 2 {cls.id}',
+            region=random.choice(us_regions),
+            firewall_ranges=[],
+        )
+        # Wait for both workspace groups to start
+        time.sleep(5)
+
+        os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] = 'information_schema'
+        os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] = cls.workspace_group.id
+
+    @classmethod
+    def tearDownClass(cls):
+        if cls.workspace_group is not None:
+            cls.workspace_group.terminate(force=True)
+        if cls.workspace_group_2 is not None:
+            cls.workspace_group_2.terminate(force=True)
+        cls.manager = None
+        cls.workspace_group = None
+        cls.workspace_group_2 = None
+        cls.workspace = None
+        cls.workspace_2 = None
+        if os.environ.get('SINGLESTOREDB_WORKSPACE', None) is not None:
+            del os.environ['SINGLESTOREDB_WORKSPACE']
+        if os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP', None) is not None:
+            del os.environ['SINGLESTOREDB_WORKSPACE_GROUP']
+        if os.environ.get('SINGLESTOREDB_DEFAULT_DATABASE', None) is not None:
+            del os.environ['SINGLESTOREDB_DEFAULT_DATABASE']
+
+    def setUp(self):
+        self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED')
+        os.environ['SINGLESTOREDB_FUSION_ENABLED'] = '1'
+        self.conn = s2.connect(database=type(self).dbname, local_infile=True)
+        self.cur = self.conn.cursor()
+
+    def tearDown(self):
+        self._clear_stage()
+
+        if self.enabled:
+            os.environ['SINGLESTOREDB_FUSION_ENABLED'] = self.enabled
+        else:
+            del os.environ['SINGLESTOREDB_FUSION_ENABLED']
+
+        try:
+            if self.cur is not None:
+                self.cur.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+        try:
+            if self.conn is not None:
+                self.conn.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+    def _clear_stage(self):
+        if self.workspace_group is not None:
+            self.cur.execute(f'''
+                show stage files
+                    in group id '{self.workspace_group.id}' recursive
+            ''')
+            files = list(self.cur)
+            folders = []
+            for file in files:
+                if file[0].endswith('/'):
+                    folders.append(file)
+                    continue
+                self.cur.execute(f'''
+                    drop stage file '{file[0]}'
+                        in group id '{self.workspace_group.id}'
+                ''')
+            for folder in folders:
+                self.cur.execute(f'''
+                    drop stage folder '{folder[0]}'
+                        in group id '{self.workspace_group.id}'
+                ''')
+
+        if self.workspace_group_2 is not None:
+            self.cur.execute(f'''
+                show stage files
+                    in group id '{self.workspace_group_2.id}' recursive
+            ''')
+            files = list(self.cur)
+            folders = []
+            for file in files:
+                if file[0].endswith('/'):
+                    folders.append(file)
+                    continue
+                self.cur.execute(f'''
+                    drop stage file '{file[0]}'
+                        in group id '{self.workspace_group_2.id}'
+                ''')
+            for folder in folders:
+                self.cur.execute(f'''
+                    drop stage folder '{folder[0]}'
+                        in group id '{self.workspace_group_2.id}'
+                ''')
+
+    def test_show_stage(self):
+        test2_sql = os.path.join(os.path.dirname(__file__), 'test2.sql')
+
+        # Should be empty
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Copy files to stage
+        self.cur.execute(
+            f'upload file to stage "new_test_1.sql" from "{test2_sql}"',
+        )
+        self.cur.execute('create stage folder "subdir1"')
+        self.cur.execute(
+            f'upload file to stage "subdir1/new_test_2.sql" from "{test2_sql}"',
+        )
+        self.cur.execute(
+            f'upload file to stage "subdir1/new_test_3.sql" from "{test2_sql}"',
+        )
+        self.cur.execute('create stage folder "subdir2"')
+        self.cur.execute(
+            f'upload file to stage "subdir2/new_test_4.sql" from "{test2_sql}"',
+        )
+        self.cur.execute(
+            f'upload file to stage "subdir2/new_test_5.sql" from "{test2_sql}"',
+        )
+
+        # Make sure files are there
+        self.cur.execute('''
+            show stage files recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 7
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_1.sql',
+            'subdir1/',
+            'subdir1/new_test_2.sql',
+            'subdir1/new_test_3.sql',
+            'subdir2/',
+            'subdir2/new_test_4.sql',
+            'subdir2/new_test_5.sql',
+        ]
+
+        # Do non-recursive listing
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_1.sql',
+            'subdir1/',
+            'subdir2/',
+        ]
+
+        # List files in specific workspace group
+        self.cur.execute(f'''
+            show stage files in group id '{self.workspace_group.id}'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_1.sql',
+            'subdir1/',
+            'subdir2/',
+        ]
+
+        self.cur.execute(f'''
+            show stage files in id '{self.workspace_group.id}'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_1.sql',
+            'subdir1/',
+            'subdir2/',
+        ]
+
+        self.cur.execute(f'''
+            show stage files in group '{self.workspace_group.name}'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_1.sql',
+            'subdir1/',
+            'subdir2/',
+        ]
+
+        self.cur.execute(f'''
+            show stage files in '{self.workspace_group.name}'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_1.sql',
+            'subdir1/',
+            'subdir2/',
+        ]
+
+        # Check other workspace group
+        self.cur.execute(f'''
+            show stage files in group '{self.workspace_group_2.name}'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Limit results
+        self.cur.execute('''
+            show stage files recursive limit 5
+        ''')
+        files = list(self.cur)
+        assert len(files) == 5
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_1.sql',
+            'subdir1/',
+            'subdir1/new_test_2.sql',
+            'subdir1/new_test_3.sql',
+            'subdir2/',
+        ]
+
+        # Order by type and name
+        self.cur.execute('''
+            show stage files order by type, name recursive extended
+        ''')
+        files = list(self.cur)
+        assert len(files) == 7
+        assert list(x[0] for x in files) == [
+            'subdir1/',
+            'subdir2/',
+            'new_test_1.sql',
+            'subdir1/new_test_2.sql',
+            'subdir1/new_test_3.sql',
+            'subdir2/new_test_4.sql',
+            'subdir2/new_test_5.sql',
+        ]
+
+        # Order by type and name descending
+        self.cur.execute('''
+            show stage files order by type desc, name desc recursive extended
+        ''')
+        files = list(self.cur)
+        assert len(files) == 7
+        assert list(x[0] for x in files) == [
+            'subdir2/new_test_5.sql',
+            'subdir2/new_test_4.sql',
+            'subdir1/new_test_3.sql',
+            'subdir1/new_test_2.sql',
+            'new_test_1.sql',
+            'subdir2/',
+            'subdir1/',
+        ]
+
+        # List at specific path
+        self.cur.execute('''
+            show stage files at 'subdir2/' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 2
+        assert list(sorted(x[0] for x in files)) == [
+            'new_test_4.sql',
+            'new_test_5.sql',
+        ]
+
+        # LIKE clause
+        self.cur.execute('''
+            show stage files like '%_4.%' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(sorted(x[0] for x in files)) == [
+            'subdir2/new_test_4.sql',
+        ]
+
+    def test_download_stage(self):
+        test2_sql = os.path.join(os.path.dirname(__file__), 'test2.sql')
+
+        # Should be empty
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Copy file to stage 1
+        self.cur.execute(f'''
+            upload file to stage 'dl_test.sql' from '{test2_sql}'
+        ''')
+
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(sorted(x[0] for x in files)) == ['dl_test.sql']
+
+        # Copy file to stage 2
+        self.cur.execute(f'''
+            upload file to stage 'dl_test2.sql'
+                in group '{self.workspace_group_2.name}'
+                from '{test2_sql}'
+        ''')
+
+        # Make sure only one file in stage 2
+        self.cur.execute(f'''
+            show stage files in group '{self.workspace_group_2.name}'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(sorted(x[0] for x in files)) == ['dl_test2.sql']
+
+        # Download file from stage 1
+        with tempfile.TemporaryDirectory() as tmpdir:
+            self.cur.execute(f'''
+                download stage file 'dl_test.sql' to '{tmpdir}/dl_test.sql'
+            ''')
+            with open(os.path.join(tmpdir, 'dl_test.sql'), 'r') as dl_file:
+                assert dl_file.read() == open(test2_sql, 'r').read()
+
+        # Download file from stage 2
+        with tempfile.TemporaryDirectory() as tmpdir:
+            self.cur.execute(f'''
+                download stage file 'dl_test2.sql'
+                    in group '{self.workspace_group_2.name}'
+                    to '{tmpdir}/dl_test2.sql'
+            ''')
+            with open(os.path.join(tmpdir, 'dl_test2.sql'), 'r') as dl_file:
+                assert dl_file.read() == open(test2_sql, 'r').read()
+
+    def test_stage_multi_wg_operations(self):
+        test_sql = os.path.join(os.path.dirname(__file__), 'test.sql')
+        test2_sql = os.path.join(os.path.dirname(__file__), 'test2.sql')
+
+        # Should be empty
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Copy file to stage 1
+        self.cur.execute(f'''
+            upload file to stage 'new_test.sql' from '{test_sql}'
+        ''')
+
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+
+        # Copy file to stage 2
+        self.cur.execute(f'''
+            upload file to stage 'new_test2.sql'
+                in group '{self.workspace_group_2.name}'
+                from '{test2_sql}'
+        ''')
+
+        # Make sure only one file in stage 1
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert files[0][0] == 'new_test.sql'
+
+        # Make sure only one file in stage 2
+        self.cur.execute(f'''
+            show stage files in group '{self.workspace_group_2.name}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(sorted(x[0] for x in files)) == ['new_test2.sql']
+
+        # Make sure only one file in stage 2 (using IN)
+        self.cur.execute(f'''
+            show stage files in '{self.workspace_group_2.name}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(sorted(x[0] for x in files)) == ['new_test2.sql']
+
+        # Make subdir
+        self.cur.execute(f'''
+            create stage folder 'data' in group '{self.workspace_group_2.name}'
+        ''')
+
+        # Upload file using workspace ID
+        self.cur.execute(f'''
+            upload file to stage 'data/new_test2_sub.sql'
+                in group id '{self.workspace_group_2.id}'
+                from '{test2_sql}'
+        ''')
+
+        # Make sure only one file in stage 1
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert files[0][0] == 'new_test.sql'
+
+        # Make sure two files in stage 2
+        self.cur.execute(f'''
+            show stage files in group id '{self.workspace_group_2.id}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == \
+            ['data/', 'data/new_test2_sub.sql', 'new_test2.sql']
+
+        # Test overwrite
+        with self.assertRaises(OSError):
+            self.cur.execute(f'''
+                upload file to stage 'data/new_test2_sub.sql'
+                    in group id '{self.workspace_group_2.id}'
+                    from '{test2_sql}'
+            ''')
+
+        self.cur.execute(f'''
+            upload file to stage 'data/new_test2_sub.sql'
+                in group id '{self.workspace_group_2.id}'
+                from '{test2_sql}' overwrite
+        ''')
+
+        # Make sure two files in stage 2
+        self.cur.execute(f'''
+            show stage files in group id '{self.workspace_group_2.id}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == \
+            ['data/', 'data/new_test2_sub.sql', 'new_test2.sql']
+
+        # Test LIKE clause
+        self.cur.execute(f'''
+            show stage files
+                in group id '{self.workspace_group_2.id}'
+                like '%_sub%' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(sorted(x[0] for x in files)) == ['data/new_test2_sub.sql']
+
+        # Drop file from default stage
+        self.cur.execute('''
+            drop stage file 'new_test.sql'
+        ''')
+
+        # Make sure no files in stage 1
+        self.cur.execute('''
+            show stage files
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Make sure two files in stage 2
+        self.cur.execute(f'''
+            show stage files in group id '{self.workspace_group_2.id}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(sorted(x[0] for x in files)) == \
+            ['data/', 'data/new_test2_sub.sql', 'new_test2.sql']
+
+        # Attempt to drop directory from stage 2
+        with self.assertRaises(OSError):
+            self.cur.execute(f'''
+                drop stage folder 'data'
+                    in group id '{self.workspace_group_2.id}'
+            ''')
+
+        self.cur.execute(f'''
+            drop stage file 'data/new_test2_sub.sql'
+                in group id '{self.workspace_group_2.id}'
+        ''')
+
+        # Make sure one file and one directory in stage 2
+        self.cur.execute(f'''
+            show stage files in group id '{self.workspace_group_2.id}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 2
+        assert list(sorted(x[0] for x in files)) == ['data/', 'new_test2.sql']
+
+        # Drop stage folder from stage 2
+        self.cur.execute(f'''
+            drop stage folder 'data'
+                in group id '{self.workspace_group_2.id}'
+        ''')
+
+        # Make sure one file in stage 2
+        self.cur.execute(f'''
+            show stage files in group id '{self.workspace_group_2.id}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(sorted(x[0] for x in files)) == ['new_test2.sql']
+
+        # Drop last file
+        self.cur.execute(f'''
+            drop stage file 'new_test2.sql'
+                in group id '{self.workspace_group_2.id}'
+        ''')
+
+        # Make sure no files in stage 2
+        self.cur.execute(f'''
+            show stage files in group id '{self.workspace_group_2.id}' recursive
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+
+@pytest.mark.management
+class TestFilesFusion(unittest.TestCase):
+
+    id: str = secrets.token_hex(8)
+    dbname: str = 'information_schema'
+    manager: None
+    workspace_group: None
+
+    @classmethod
+    def setUpClass(cls):
+        cls.manager = s2.manage_workspaces()
+        us_regions = [x for x in cls.manager.regions if x.name.startswith('US')]
+        cls.workspace_group = cls.manager.create_workspace_group(
+            f'Files Fusion Testing {cls.id}',
+            region=random.choice(us_regions),
+            firewall_ranges=[],
+        )
+        # Wait for both workspace groups to start
+        time.sleep(5)
+
+        os.environ['SINGLESTOREDB_DEFAULT_DATABASE'] = 'information_schema'
+        os.environ['SINGLESTOREDB_WORKSPACE_GROUP'] = cls.workspace_group.id
+
+    @classmethod
+    def tearDownClass(cls):
+        if cls.workspace_group is not None:
+            cls.workspace_group.terminate(force=True)
+        cls.manager = None
+        cls.workspace_group = None
+        cls.workspace = None
+        if os.environ.get('SINGLESTOREDB_WORKSPACE', None) is not None:
+            del os.environ['SINGLESTOREDB_WORKSPACE']
+        if os.environ.get('SINGLESTOREDB_WORKSPACE_GROUP', None) is not None:
+            del os.environ['SINGLESTOREDB_WORKSPACE_GROUP']
+        if os.environ.get('SINGLESTOREDB_DEFAULT_DATABASE', None) is not None:
+            del os.environ['SINGLESTOREDB_DEFAULT_DATABASE']
+
+    def setUp(self):
+        self.enabled = os.environ.get('SINGLESTOREDB_FUSION_ENABLED')
+        os.environ['SINGLESTOREDB_FUSION_ENABLED'] = '1'
+        self.conn = s2.connect(database=type(self).dbname, local_infile=True)
+        self.cur = self.conn.cursor()
+
+    def tearDown(self):
+        self._clear_files()
+
+        if self.enabled:
+            os.environ['SINGLESTOREDB_FUSION_ENABLED'] = self.enabled
+        else:
+            del os.environ['SINGLESTOREDB_FUSION_ENABLED']
+
+        try:
+            if self.cur is not None:
+                self.cur.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+        try:
+            if self.conn is not None:
+                self.conn.close()
+        except Exception:
+            # traceback.print_exc()
+            pass
+
+    def _clear_files(self):
+        cls = type(self)
+        for prefix in ['show', 'dl', 'drop']:
+            for i in range(1, 6):
+                try:
+                    self.cur.execute(
+                        f'''drop personal file "{prefix}_test_{i}_{cls.id}.ipynb"''',
+                    )
+                except (OSError, s2.ManagementError):
+                    pass
+            for i in range(1, 6):
+                try:
+                    self.cur.execute(
+                        f'''drop shared file "{prefix}_test_{i}_{cls.id}.ipynb"''',
+                    )
+                except (OSError, s2.ManagementError):
+                    pass
+
+    def test_show_personal_files(self):
+        return self._test_show_files('personal')
+
+    def test_show_shared_files(self):
+        return self._test_show_files('shared')
+
+    def _test_show_files(self, ftype):
+        cls = type(self)
+        nb = os.path.join(os.path.dirname(__file__), 'test.ipynb')
+
+        # Should be empty
+        self.cur.execute(f'''
+            show {ftype} files like 'show_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Upload files
+        self.cur.execute(
+            f'upload {ftype} file to "show_test_1_{cls.id}.ipynb" from "{nb}"',
+        )
+        self.cur.execute(
+            f'upload {ftype} file to "show_test_2_{cls.id}.ipynb" from "{nb}"',
+        )
+        self.cur.execute(
+            f'upload {ftype} file to "show_test_3_{cls.id}.ipynb" from "{nb}"',
+        )
+        self.cur.execute(
+            f'upload {ftype} file to "show_test_4_{cls.id}.ipynb" from "{nb}"',
+        )
+        self.cur.execute(
+            f'upload {ftype} file to "show_test_5_{cls.id}.ipynb" from "{nb}"',
+        )
+
+        # Make sure files are there
+        self.cur.execute(f'''
+            show {ftype} files like 'show_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 5
+        assert list(sorted(x[0] for x in files)) == [
+            f'show_test_1_{cls.id}.ipynb',
+            f'show_test_2_{cls.id}.ipynb',
+            f'show_test_3_{cls.id}.ipynb',
+            f'show_test_4_{cls.id}.ipynb',
+            f'show_test_5_{cls.id}.ipynb',
+        ]
+
+        # Test ORDER BY
+        self.cur.execute(f'''
+            show {ftype} files like 'show_%{cls.id}%' order by name desc
+        ''')
+        files = list(self.cur)
+        assert len(files) == 5
+        assert list(x[0] for x in files) == [
+            f'show_test_5_{cls.id}.ipynb',
+            f'show_test_4_{cls.id}.ipynb',
+            f'show_test_3_{cls.id}.ipynb',
+            f'show_test_2_{cls.id}.ipynb',
+            f'show_test_1_{cls.id}.ipynb',
+        ]
+
+        # Test LIMIT
+        self.cur.execute(f'''
+            show {ftype} files like 'show_%{cls.id}%' order by name desc limit 3
+        ''')
+        files = list(self.cur)
+        assert len(files) == 3
+        assert list(x[0] for x in files) == [
+            f'show_test_5_{cls.id}.ipynb',
+            f'show_test_4_{cls.id}.ipynb',
+            f'show_test_3_{cls.id}.ipynb',
+        ]
+
+        # Test EXTENDED
+        self.cur.execute(f'''
+            show {ftype} files like 'show_%{cls.id}%' extended
+        ''')
+        assert [x[0] for x in self.cur.description] == \
+            ['Name', 'Type', 'Size', 'Writable', 'CreatedAt', 'LastModifiedAt']
+
+    def test_download_personal_files(self):
+        return self._test_download_files('personal')
+
+    def test_download_shared_files(self):
+        return self._test_download_files('shared')
+
+    def _test_download_files(self, ftype):
+        cls = type(self)
+        nb = os.path.join(os.path.dirname(__file__), 'test.ipynb')
+
+        # Should be empty
+        self.cur.execute(f'''
+            show {ftype} files like 'dl_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Upload files
+        self.cur.execute(f'upload {ftype} file to "dl_test_1_{cls.id}.ipynb" from "{nb}"')
+        self.cur.execute(f'upload {ftype} file to "dl_test_2_{cls.id}.ipynb" from "{nb}"')
+
+        # Make sure files are there
+        self.cur.execute(f'''
+            show {ftype} files like 'dl_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 2
+        assert list(sorted(x[0] for x in files)) == [
+            f'dl_test_1_{cls.id}.ipynb',
+            f'dl_test_2_{cls.id}.ipynb',
+        ]
+
+        # Download files
+        with tempfile.TemporaryDirectory() as tmpdir:
+            self.cur.execute(f'''
+                download {ftype} file 'dl_test_1_{cls.id}.ipynb'
+                    to '{tmpdir}/dl_test_1.ipynb'
+            ''')
+            with open(os.path.join(tmpdir, 'dl_test_1.ipynb'), 'r') as dl_file:
+                assert dl_file.read() == open(nb, 'r').read()
+
+            self.cur.execute(f'''
+                download {ftype} file 'dl_test_2_{cls.id}.ipynb'
+                    to '{tmpdir}/dl_test_2.ipynb'
+            ''')
+            with open(os.path.join(tmpdir, 'dl_test_2.ipynb'), 'r') as dl_file:
+                assert dl_file.read() == open(nb, 'r').read()
+
+    def test_drop_personal_files(self):
+        return self._test_drop_files('personal')
+
+    def test_drop_shared_files(self):
+        return self._test_drop_files('shared')
+
+    def _test_drop_files(self, ftype):
+        cls = type(self)
+        nb = os.path.join(os.path.dirname(__file__), 'test.ipynb')
+
+        # Should be empty
+        self.cur.execute(f'''
+            show {ftype} files like 'drop_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0
+
+        # Upload files
+        self.cur.execute(
+            f'upload {ftype} file to "drop_test_1_{cls.id}.ipynb" from "{nb}"',
+        )
+        self.cur.execute(
+            f'upload {ftype} file to "drop_test_2_{cls.id}.ipynb" from "{nb}"',
+        )
+
+        # Make sure files are there
+        self.cur.execute(f'''
+            show {ftype} files like 'drop_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 2
+        assert list(sorted(x[0] for x in files)) == [
+            f'drop_test_1_{cls.id}.ipynb',
+            f'drop_test_2_{cls.id}.ipynb',
+        ]
+
+        # Drop 1 file
+        self.cur.execute(f'''
+            drop {ftype} file 'drop_test_1_{cls.id}.ipynb'
+        ''')
+
+        # Make sure 1 file is there
+        self.cur.execute(f'''
+            show {ftype} files like 'drop_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 1
+        assert list(x[0] for x in files) == [f'drop_test_2_{cls.id}.ipynb']
+
+        # Drop 2nd file
+        self.cur.execute(f'''
+            drop {ftype} file 'drop_test_2_{cls.id}.ipynb'
+        ''')
+
+        # Make sure no files are there
+        self.cur.execute(f'''
+            show {ftype} files like 'drop_%{cls.id}%'
+        ''')
+        files = list(self.cur)
+        assert len(files) == 0