vincbeck commented on code in PR #47869:
URL: https://github.com/apache/airflow/pull/47869#discussion_r2003462453
##########
airflow/api_fastapi/auth/managers/simple/simple_auth_manager.py:
##########
@@ -98,37 +99,46 @@ def get_users() -> list[dict[str, str]]:
@staticmethod
def get_passwords(users: list[dict[str, str]]) -> dict[str, str]:
- user_passwords_from_file = {}
-
- # Read passwords from file
password_file = SimpleAuthManager.get_generated_password_file()
- if os.path.isfile(password_file):
- with open(password_file) as file:
- passwords_str = file.read().strip()
- user_passwords_from_file = json.loads(passwords_str)
-
- usernames = {user["username"] for user in users}
- return {
- username: password
- for username, password in user_passwords_from_file.items()
- if username in usernames
- }
+ with open(password_file, "r+") as file:
+ return SimpleAuthManager._get_passwords(users=users,
stream=file)[0]
def init(self) -> None:
is_simple_auth_manager_all_admins = conf.getboolean("core",
"simple_auth_manager_all_admins")
if is_simple_auth_manager_all_admins:
return
users = self.get_users()
- passwords = self.get_passwords(users)
- for user in users:
- if user["username"] not in passwords:
- # User dot not exist in the file, adding it
- passwords[user["username"]] = self._generate_password()
-
- self._print_output(f"Password for user '{user['username']}':
{passwords[user['username']]}")
-
- with open(self.get_generated_password_file(), "w") as file:
- file.write(json.dumps(passwords))
+ password_file = self.get_generated_password_file()
+
+ try:
+ if not os.path.exists(password_file):
+ with open(password_file, "w") as file:
+ file.write("{}\n")
+ with open(password_file, "r+") as file:
+ # Non-blocking exclusive lock on this file
+ # Fastapi spins up N workers, so this method is called N times
in N different processes
+ # This needs to be called only once so we use the file
``password_file`` as locking mechanism
+ fcntl.flock(file, fcntl.LOCK_EX | fcntl.LOCK_NB)
+ passwords, changed = self._get_passwords(users=users,
stream=file)
+ for user in users:
+ if user["username"] not in passwords:
+ # User dot not exist in the file, adding it
+ passwords[user["username"]] = self._generate_password()
+ self._print_output(
+ f"Password for user '{user['username']}':
{passwords[user['username']]}"
+ )
+ changed = True
+
+ if changed:
+ file.seek(0)
+ file.truncate()
+ file.write(json.dumps(passwords) + "\n")
+
+ # Release lock
+ fcntl.flock(file, fcntl.LOCK_UN)
Review Comment:
Added
--
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]
