uranusjr commented on a change in pull request #18916:
URL: https://github.com/apache/airflow/pull/18916#discussion_r733280253



##########
File path: airflow/cli/commands/role_command.py
##########
@@ -42,3 +45,45 @@ def roles_create(args):
     for role_name in args.role:
         appbuilder.sm.add_role(role_name)
     print(f"Added {len(args.role)} role(s)")
+
+
+@suppress_logs_and_warning
+def roles_export(args):
+    """
+    Exports all the rules from the data base to a file.
+    Note, this function does not export the permissions associated for each 
role.
+    Strictly, it exports the role names into the passed role json file.
+    """
+    appbuilder = cached_app().appbuilder
+    roles = appbuilder.sm.get_all_roles()
+    exporting_roles = [role.name for role in roles if role.name not in 
EXISTING_ROLES]
+    filename = os.path.expanduser(args.file)
+    with open(filename, 'w') as file:
+        file.write(json.dumps(exporting_roles, sort_keys=True, indent=4))

Review comment:
       ```suggestion
       with open(filename, 'w') as f:
           json.dump(f, exporting_roles, sort_keys=True, indent=4)
   ```
   
   I feel it may be worthwhile to add some formatting options to the command, 
since both `sort_keys` and `indent` have performance considerations and a user 
may reasonably want to not set those if there's a large amount of data. Maybe 
we should even default to _not_ pretty-format the output.

##########
File path: tests/cli/commands/test_role_command.py
##########
@@ -17,11 +17,13 @@
 # under the License.
 #
 import io
+import tempfile
 from contextlib import redirect_stdout
 
 import pytest
 
 from airflow.cli.commands import role_command
+from airflow.settings import json

Review comment:
       Just `import json`?

##########
File path: airflow/cli/commands/role_command.py
##########
@@ -42,3 +44,43 @@ def roles_create(args):
     for role_name in args.role:
         appbuilder.sm.add_role(role_name)
     print(f"Added {len(args.role)} role(s)")
+
+
+@suppress_logs_and_warning
+def roles_export(args):
+    """
+    Exports all the rules from the data base to a file.
+    """
+    from airflow.www.security import EXISTING_ROLES
+    appbuilder = cached_app().appbuilder
+    roles = appbuilder.sm.get_all_roles()
+    exporting_roles = [role.name for role in roles if role.name not in 
EXISTING_ROLES]
+    with open(os.path.expanduser(args.export), 'w') as file:
+        file.write(json.dumps(exporting_roles, sort_keys=True, indent=4))
+        print(f"{len(exporting_roles)} roles successfully exported to 
{file.name}")
+
+
+@cli_utils.action_logging
+@suppress_logs_and_warning
+def roles_import(args):
+    """
+    Import all the roles into the db from the given json file.
+    """
+    json_file = getattr(args, 'import')
+    if not os.path.exists(json_file):
+        print(f"File '{json_file}' does not exist")
+        exit(1)
+
+    role_list = None
+    try:
+        with open(json_file, 'r') as file:
+            role_list = json.loads(file.read())
+    except ValueError as e:
+        print(f"File '{json_file}' is not a valid JSON file. Error: {e}")
+        exit(1)
+    appbuilder = cached_app().appbuilder
+    existing_roles = [role.name for role in appbuilder.sm.get_all_roles()]
+    roles_to_import = [role for role in role_list if role not in 
existing_roles]
+    for role_name in roles_to_import:
+        appbuilder.sm.add_role(role_name)

Review comment:
       That makes sense, thanks. What's the reasoning behind not importing the 
permissions? Would it be a good idea to add this functionality? (Not 
necessarily in this PR but whether this is a good idea in general.)




-- 
This is an automated message from the Apache Git Service.
To respond to the message, please log on to GitHub and use the
URL above to go to the specific comment.

To unsubscribe, e-mail: [email protected]

For queries about this service, please contact Infrastructure at:
[email protected]


Reply via email to