diff --git a/flask_appbuilder/cli.py b/flask_appbuilder/cli.py index 062f8775ef..f993f93881 100644 --- a/flask_appbuilder/cli.py +++ b/flask_appbuilder/cli.py @@ -1,7 +1,7 @@ from io import BytesIO import os import shutil -from typing import Optional +from typing import Optional, Union from urllib.request import urlopen from zipfile import ZipFile @@ -150,9 +150,22 @@ def create_db(): @fab.command("export-roles") @with_appcontext @click.option("--path", "-path", help="Specify filepath to export roles to") -def export_roles(path: Optional[str] = None) -> None: - """ Exports roles with permissions and view menus to JSON file """ - current_app.appbuilder.sm.export_roles(path) +@click.option("--indent", help="Specify indent of generated JSON file") +def export_roles( + path: Optional[str] = None, indent: Optional[Union[int, str]] = None +) -> None: + """Exports roles with permissions and view menus to JSON file""" + # Cast negative numbers to int (as they're passed as str from CLI) + try: + indent = int(indent) + except TypeError: + # Don't cast None + pass + except ValueError: + # Don't cast non-int-like strings + pass + + current_app.appbuilder.sm.export_roles(path=path, indent=indent) @fab.command("import-roles") diff --git a/flask_appbuilder/security/manager.py b/flask_appbuilder/security/manager.py index 186215b745..bd173a6813 100644 --- a/flask_appbuilder/security/manager.py +++ b/flask_appbuilder/security/manager.py @@ -3,7 +3,7 @@ import json import logging import re -from typing import Any, Dict, List, Optional, Set, Tuple +from typing import Any, Dict, List, Optional, Set, Tuple, Union from flask import g, session, url_for from flask_babel import lazy_gettext as _ @@ -2023,7 +2023,9 @@ def del_permission_role(self, role, perm_view): """ raise NotImplementedError - def export_roles(self, path: Optional[str] = None) -> None: + def export_roles( + self, path: Optional[str] = None, indent: Optional[Union[int, str]] = None + ) -> None: """ Exports roles to JSON file. """ raise NotImplementedError diff --git a/flask_appbuilder/security/mongoengine/manager.py b/flask_appbuilder/security/mongoengine/manager.py index 78af62d32c..bf5cc9f2e0 100644 --- a/flask_appbuilder/security/mongoengine/manager.py +++ b/flask_appbuilder/security/mongoengine/manager.py @@ -1,7 +1,7 @@ from datetime import datetime import json import logging -from typing import List, Optional +from typing import List, Optional, Union import uuid from werkzeug.security import generate_password_hash @@ -408,8 +408,10 @@ def del_permission_role(self, role, perm_view): except Exception as e: log.error(c.LOGMSG_ERR_SEC_DEL_PERMROLE.format(str(e))) - def export_roles(self, path: Optional[str] = None) -> None: - """ Exports roles to JSON file. """ + def export_roles( + self, path: Optional[str] = None, indent: Optional[Union[int, str]] = None + ) -> None: + """Exports roles to JSON file.""" timestamp = datetime.now().strftime("%Y%m%dT%H%M%S") filename = path or f"roles_export_{timestamp}.json" @@ -430,7 +432,7 @@ def export_roles(self, path: Optional[str] = None) -> None: serialized_roles.append(serialized_role) with open(filename, "w") as fd: - fd.write(json.dumps(serialized_roles)) + fd.write(json.dumps(serialized_roles, indent=indent)) def import_roles(self, path: str) -> None: """ Imports roles from JSON file. """ diff --git a/flask_appbuilder/security/sqla/manager.py b/flask_appbuilder/security/sqla/manager.py index a64f4068bb..0db0b2941b 100755 --- a/flask_appbuilder/security/sqla/manager.py +++ b/flask_appbuilder/security/sqla/manager.py @@ -1,7 +1,7 @@ from datetime import datetime import json import logging -from typing import List, Optional +from typing import List, Optional, Union import uuid from sqlalchemy import and_, func, literal, update @@ -660,7 +660,9 @@ def del_permission_role(self, role, perm_view): log.error(c.LOGMSG_ERR_SEC_DEL_PERMROLE.format(str(e))) self.get_session.rollback() - def export_roles(self, path: Optional[str] = None) -> None: + def export_roles( + self, path: Optional[str] = None, indent: Optional[Union[int, str]] = None + ) -> None: """ Exports roles to JSON file. """ timestamp = datetime.now().strftime("%Y%m%dT%H%M%S") filename = path or f"roles_export_{timestamp}.json" @@ -680,7 +682,7 @@ def export_roles(self, path: Optional[str] = None) -> None: serialized_roles.append(serialized_role) with open(filename, "w") as fd: - fd.write(json.dumps(serialized_roles)) + fd.write(json.dumps(serialized_roles, indent=indent)) def import_roles(self, path: str) -> None: """ Imports roles from JSON file. """ diff --git a/flask_appbuilder/tests/test_fab_cli.py b/flask_appbuilder/tests/test_fab_cli.py index 8eeedb08a9..e06f5f583e 100644 --- a/flask_appbuilder/tests/test_fab_cli.py +++ b/flask_appbuilder/tests/test_fab_cli.py @@ -3,6 +3,7 @@ import logging import os import tempfile +from unittest.mock import ANY, patch from click.testing import CliRunner from flask import Flask @@ -146,6 +147,29 @@ def test_export_roles_filename(self): len(glob.glob(os.path.join(tmp_dir, "roles_export_*"))), 0 ) + @patch("json.dumps") + def test_export_roles_indent(self, mock_json_dumps): + """Test that json.dumps is called with the correct argument passed from CLI.""" + with tempfile.TemporaryDirectory() as tmp_dir: + app = Flask("src_app") + app.config.from_object("flask_appbuilder.tests.config_security") + app.config[ + "SQLALCHEMY_DATABASE_URI" + ] = f"sqlite:///{os.path.join(tmp_dir, 'src.db')}" + db = SQLA(app) + app_builder = AppBuilder(app, db.session) # noqa: F841 + cli_runner = app.test_cli_runner() + + cli_runner.invoke(export_roles) + mock_json_dumps.assert_called_with(ANY, indent=None) + mock_json_dumps.reset_mock() + + example_cli_args = ["", "foo", -1, 0, 1] + for arg in example_cli_args: + cli_runner.invoke(export_roles, [f"--indent={arg}"]) + mock_json_dumps.assert_called_with(ANY, indent=arg) + mock_json_dumps.reset_mock() + def test_import_roles(self): with tempfile.TemporaryDirectory() as tmp_dir: app = Flask("dst_app")