Fix Stage directory issues

singlestoredb/management/workspace.py

@@ -215,6 +215,11 @@ def remove(self) -> None:
                 msg='No Stage object is associated with this object.',
             )
 
+        if self.type == 'directory':
+            raise IsADirectoryError(
+                f'path is a directory; use rmdir or removedirs {self.path}',
+            )
+
         self._stage.remove(self.path)
 
     def rmdir(self) -> None:
@@ -224,6 +229,11 @@ def rmdir(self) -> None:
                 msg='No Stage object is associated with this object.',
             )
 
+        if self.type != 'directory':
+            raise NotADirectoryError(
+                f'path is not a directory: {self.path}',
+            )
+
         self._stage.rmdir(self.path)
 
     def removedirs(self) -> None:
@@ -233,6 +243,11 @@ def removedirs(self) -> None:
                 msg='No Stage object is associated with this object.',
             )
 
+        if self.type != 'directory':
+            raise NotADirectoryError(
+                f'path is not a directory: {self.path}',
+            )
+
         self._stage.removedirs(self.path)
 
     def rename(self, new_path: PathLike, *, overwrite: bool = False) -> None:

singlestoredb/tests/test_management.py

@@ -615,7 +615,7 @@ def test_os_directories(self):
         assert 'mkdir_test_2/' not in after
         assert list(sorted(before)) == list(sorted(after + ['mkdir_test_2/']))
 
-        with self.assertRaises(NotADirectoryError):
+        with self.assertRaises(s2.ManagementError):
             st.removedirs('mkdir_test.sql')
 
     def test_os_files(self):