update scheduler init db process & db_user priority

Author: wacmkxiaoyi

scheduler/app/database/__init__.py

@@ -11,15 +11,12 @@
 DB_PASSWD = os.getenv("DB_PASSWD")
 db_user = ""
 db_passwd = ""
-if DB_ROOT_USER and DB_ROOT_PASSWD:
-    db_user = "root"
-    db_passwd = DB_ROOT_PASSWD
-elif DB_USER and DB_PASSWD:
+if DB_USER:
     db_user = DB_USER
     db_passwd = DB_PASSWD
 else:
     db_user = "root"
-    db_passwd = ""
+    db_passwd = DB_ROOT_PASSWD
 
 DB_HOST = os.getenv("DB_HOST", "seatable-mysql")
 DB_PORT = os.getenv("DB_PORT", "3306")

scheduler/app/database/init_db.py

@@ -9,15 +9,12 @@
 DB_PASSWD = os.getenv("DB_PASSWD")
 db_user = ""
 db_passwd = ""
-if DB_ROOT_USER and DB_ROOT_PASSWD:
-    db_user = "root"
-    db_passwd = DB_ROOT_PASSWD
-elif DB_USER and DB_PASSWD:
+if DB_USER:
     db_user = DB_USER
     db_passwd = DB_PASSWD
 else:
     db_user = "root"
-    db_passwd = ""
+    db_passwd = DB_ROOT_PASSWD
 
 DB_HOST = os.getenv("DB_HOST", "seatable-mysql")
 DB_PORT = int(os.getenv("DB_PORT", "3306"))
@@ -38,21 +35,49 @@ def wait_for_mysql():
         connection.close()
         return
 
+def check_and_create_mysql_user():
+    connection = pymysql.connect(host=DB_HOST, port=DB_PORT, user=DB_ROOT_USER, passwd=DB_ROOT_PASSWD)
+
+    try:
+        with connection.cursor() as cursor:
+            if db_user != 'root':
+                query_user_exists = f"SELECT EXISTS(SELECT 1 FROM mysql.user WHERE User='{db_user}' AND Host='%') AS user_exists;"
+                cursor.execute(query_user_exists)
+                result = cursor.fetchone()
+                user_exists = result[0] == 1
+                if not user_exists:
+                    create_user_sql = f"CREATE USER '{db_user}'@'%' IDENTIFIED BY '{db_passwd}';"
+                    cursor.execute(create_user_sql)
+                    print(f"Create user '{db_user}'@'%' sucessfully.")
+    
+                grant_privileges_sql = (
+                    f"GRANT ALL PRIVILEGES ON {DATABASE_NAME}.* TO '{db_user}'@'%';"
+                )
+                for stmt in grant_privileges_sql.split(';'):
+                    if stmt.strip():
+                        cursor.execute(stmt)
+                print(f"Granted user '{db_user}'@'%' privileges sucessfully.")
+    
+                cursor.execute("FLUSH PRIVILEGES;")
+    finally:
+        connection.close()
 
 wait_for_mysql()
 
-if db_user == "root":
-    sql = 'mysql -h %s -u%s -p%s -e "CREATE DATABASE IF NOT EXISTS %s;"' % (
-        shlex.quote(DB_HOST),
-        shlex.quote(db_user),
-        shlex.quote(db_passwd),
-        DATABASE_NAME,
-    )
-    os.system(sql)
+sql = 'mysql -h %s -u%s -p%s -e "CREATE DATABASE IF NOT EXISTS %s;"' % (
+    shlex.quote(DB_HOST),
+    shlex.quote(DB_ROOT_USER),
+    shlex.quote(DB_ROOT_PASSWD),
+    DATABASE_NAME,
+)
+os.system(sql)
+
+check_and_create_mysql_user()
+
 sql = "mysql -h %s -u%s -p%s %s </opt/scheduler/database/initial_tables.sql" % (
     shlex.quote(DB_HOST),
-    shlex.quote(db_user),
-    shlex.quote(db_passwd),
+    shlex.quote(DB_ROOT_USER),
+    shlex.quote(DB_ROOT_PASSWD),
     DATABASE_NAME,
 )
 os.system(sql)

scheduler/app/upgrade/upgrade.py

@@ -23,15 +23,12 @@
 DB_PASSWD = os.getenv("DB_PASSWD")
 db_user = ""
 db_passwd = ""
-if DB_ROOT_USER and DB_ROOT_PASSWD:
-    db_user = "root"
-    db_passwd = DB_ROOT_PASSWD
-elif DB_USER and DB_PASSWD:
+if DB_USER:
     db_user = DB_USER
     db_passwd = DB_PASSWD
 else:
     db_user = "root"
-    db_passwd = ""
+    db_passwd = DB_ROOT_PASSWD
 
 DB_HOST = os.getenv("DB_HOST", "seatable-mysql")
 DB_PORT = os.getenv("DB_PORT", "3306")