Modularize the test.sh and separate DBMSs configuration into new folders

Dockerfile.template

@@ -3,20 +3,16 @@ FROM ubuntu:20.04
 ENV DEBIAN_FRONTEND=noninteractive
 
 RUN apt-get update && \
-    apt-get install -y default-jdk maven git jq {{DBMS_PACKAGE}} && \
+    apt-get install -y default-jdk maven git jq curl unzip python3 python3-pip && \
     rm -rf /var/lib/apt/lists/*
 
-RUN mkdir -p /var/run/mysqld && chown -R mysql:mysql /var/run/mysqld || true
-RUN mkdir -p /var/run/postgresql && chown -R postgres:postgres /var/run/postgresql || true
-
-WORKDIR /root
+# [DBMS_INSTALL_HERE]
 
 COPY entrypoint.sh /root/entrypoint.sh
-COPY sqlancer_runner.sh /root/sqlancer_runner.sh
-COPY init_schema.sql /root/init_schema.sql
 COPY config.json /root/config.json
-COPY db_init/ /root/db_init/
 
-RUN chmod +x /root/entrypoint.sh /root/sqlancer_runner.sh /root/db_init/*.sh
+WORKDIR /root
+RUN mkdir -p /root/sqlancer
 
+RUN chmod +x /root/entrypoint.sh
 CMD ["/root/entrypoint.sh"]

config.json

@@ -1,10 +1,12 @@
 {
   "dbms": "mysql",
+  "version": "8.0",
+  "num_threads": 4,
+  "timeout_seconds": 60,
   "oracle": "FUZZER",
-  "username": "root",
   "password": "12345678",
-  "host": "127.0.0.1",
-  "port": 3306,
-  "num_threads": 4,
-  "timeout_seconds": 60
+  "versions": {
+    "mysql": "8.0",
+    "postgres": "13"
+  }
 }

entrypoint.sh

@@ -1,22 +1,39 @@
 #!/bin/bash
 set -e
 
-CONFIG_PATH="/root/config.json"
-DBMS=$(jq -r '.dbms' "$CONFIG_PATH")
+DBMS="${DBMS:-mysql}"
+VERSION="${VERSION:-8.0}"
+SQLANCER_THREADS="${SQLANCER_THREADS:-4}"
+SQLANCER_TIMEOUT="${SQLANCER_TIMEOUT:-60}"
+SQLANCER_ORACLE="${SQLANCER_ORACLE:-FUZZER}"
+SQLANCER_USERNAME="${SQLANCER_USERNAME:-root}"
+SQLANCER_PASSWORD="${SQLANCER_PASSWORD:-12345678}"
 
-echo "[BOOT] Selected DBMS: $DBMS"
+echo "=== [BOOT] Selected DBMS: $DBMS ==="
 
-case "$DBMS" in
-  mysql)
-    bash /root/db_init/start_mysql.sh
-    ;;
-  postgres)
-    bash /root/db_init/start_postgres.sh
-    ;;
-  *)
-    echo "[ERROR] Unsupported DBMS: $DBMS"
-    exit 1
-    ;;
-esac
+export PYTHONPATH=/root
 
-bash /root/sqlancer_runner.sh
+# init db
+python3 - <<PYCODE
+import os
+dbms = os.environ["DBMS"]
+password = os.environ["SQLANCER_PASSWORD"]
+mod = __import__(f"{dbms}.docker_ops", fromlist=["init"])
+mod.init(password)
+PYCODE
+
+echo "[SQLANCER] Cloning SQLancer..."
+git clone https://github.com/sqlancer/sqlancer.git
+cd sqlancer
+mvn package -DskipTests
+
+echo "[SQLancer] Running test with command:"
+echo "java -jar target/sqlancer-*.jar --num-threads $SQLANCER_THREADS --timeout-seconds $SQLANCER_TIMEOUT --username $SQLANCER_USERNAME --password $SQLANCER_PASSWORD $DBMS --oracle $SQLANCER_ORACLE"
+
+java -jar target/sqlancer-*.jar \
+  --num-threads "$SQLANCER_THREADS" \
+  --timeout-seconds "$SQLANCER_TIMEOUT" \
+  --username "$SQLANCER_USERNAME" \
+  --password "$SQLANCER_PASSWORD" \
+  "$DBMS" \
+  --oracle "$SQLANCER_ORACLE"

generate_dockerfile.py

@@ -1,24 +1,67 @@
+import sys
+import os
 import json
 
-with open("config.json") as f:
-    config = json.load(f)
+CONFIG_PATH = "config.json"
 
-dbms = config.get("dbms", "").lower()
+def load_config():
+    if not os.path.exists(CONFIG_PATH):
+        raise FileNotFoundError("config.json not found")
+    with open(CONFIG_PATH) as f:
+        return json.load(f)
 
-dbms_packages = {
-    "mysql": "mysql-server",
-    "postgres": "postgresql",
-}
+def generate_dockerfile(dbms=None, version=None):
+    config = load_config()
+    dbms = dbms or config.get("dbms")
+    version = version or config.get("version")
+    dbms_dir = os.path.join(dbms)
+    init_sql_path = os.path.join(dbms_dir, "init.sql")
 
-if dbms not in dbms_packages:
-    raise ValueError(f"Unsupported DBMS: {dbms}")
+    dockerfile = f"""FROM ubuntu:20.04
 
-with open("Dockerfile.template") as f:
-    template = f.read()
+ENV DEBIAN_FRONTEND=noninteractive
 
-final_dockerfile = template.replace("{{DBMS_PACKAGE}}", dbms_packages[dbms])
+RUN apt-get update && \\
+    apt-get install -y default-jdk maven git jq curl unzip python3 python3-pip && \\
+    rm -rf /var/lib/apt/lists/*
 
-with open("Dockerfile", "w") as f:
-    f.write(final_dockerfile)
+"""
 
-print(f"[✔] Dockerfile generated for DBMS: {dbms}")
+    install_script = ""
+    if dbms == "mysql":
+        from mysql.docker_ops import get_install_script
+        install_script = get_install_script(version)
+    elif dbms == "postgres":
+        from postgres.docker_ops import get_install_script
+        install_script = get_install_script(version)
+    else:
+        raise ValueError(f"Unsupported DBMS: {dbms}")
+
+    dockerfile += install_script
+
+
+    dockerfile += f"""
+COPY entrypoint.sh /root/entrypoint.sh
+COPY config.json /root/config.json
+COPY {dbms} /root/{dbms}
+WORKDIR /root
+RUN mkdir -p /root/sqlancer && chmod +x /root/entrypoint.sh
+"""
+
+    if os.path.exists(init_sql_path):
+        dockerfile += f"""
+COPY {init_sql_path} /db_init/init.sql
+RUN chmod 755 /db_init && chmod 644 /db_init/init.sql
+"""
+
+    dockerfile += 'CMD ["/root/entrypoint.sh"]\n'
+
+    with open("Dockerfile", "w") as f:
+        f.write(dockerfile)
+
+    print(f"✅ Dockerfile generated，including {dbms}:{version} ")
+
+if __name__ == "__main__":
+    dbms_arg = sys.argv[1] if len(sys.argv) > 1 else None
+    version_arg = sys.argv[2] if len(sys.argv) > 2 else None
+    generate_dockerfile(dbms_arg, version_arg)

mysql/docker_ops.py

@@ -0,0 +1,34 @@
+# mysql/docker_ops.py
+
+import subprocess
+import os
+
+
+def get_install_script(version: str) -> str:
+    """
+    return MySQL - Dockerfile 
+    """
+    return f"""
+# install MySQL {version}
+RUN apt-get update && \\
+    apt-get install -y mysql-server && \\
+    rm -rf /var/lib/apt/lists/*
+"""
+
+def pull_docker_image(version: str):
+    """
+    pull MySQL 
+    """
+    image = f"mysql:{version}"
+    print(f"[MySQL]  {image} ...")
+    subprocess.run(["sudo", "docker", "pull", image], check=True)
+
+def init(sql_password: str):
+    print("[MYSQL] Starting MySQL...")
+    subprocess.run("service mysql start", shell=True, check=True)
+    print("[MYSQL] Setting root password...")
+
+    init_sql = "/db_init/init.sql"
+    if os.path.exists(init_sql):
+        print("[MYSQL] Executing init.sql...")
+        subprocess.run(f"mysql -u root -p'{sql_password}' < {init_sql}", shell=True, check=True)

mysql/init.sql

@@ -0,0 +1,6 @@
+-- mysql/init.sql
+
+CREATE DATABASE IF NOT EXISTS test_db;
+ALTER USER 'root'@'localhost' IDENTIFIED WITH mysql_native_password BY '12345678';
+SELECT host, user FROM mysql.user;
+FLUSH PRIVILEGES;

postgres/docker_ops.py

@@ -0,0 +1,33 @@
+# postgres/docker_ops.py
+
+import subprocess
+import os
+
+
+def get_install_script(version: str) -> str:
+    """
+    return postgresql dockerfile
+    """
+    return f"""
+# install PostgreSQL {version}
+RUN apt-get update && \\
+    apt-get install -y postgresql postgresql-contrib && \\
+    rm -rf /var/lib/apt/lists/*
+"""
+
+def pull_docker_image(version: str):
+    """
+    pull PostgreSQL 
+    """
+    image = f"postgres:{version}"
+    print(f"[PostgreSQL] pull {image} ...")
+    subprocess.run(["sudo", "docker", "pull", image], check=True)
+
+def init(sql_password: str):
+    print("[POSTGRES] Starting PostgreSQL...")
+    subprocess.run("service postgresql start", shell=True, check=True)
+    init_sql = "/db_init/init.sql"
+    if os.path.exists(init_sql):
+        print("[POSTGRES] Executing init.sql...")
+
+        subprocess.run(f"su - postgres -c \"psql < {init_sql}\"", shell=True, check=True)

postgres/init.sql

@@ -0,0 +1,6 @@
+-- postgres/init.sql
+
+CREATE DATABASE test;
+ALTER USER postgres WITH PASSWORD '12345678';
+ALTER USER postgres WITH SUPERUSER;
+GRANT ALL PRIVILEGES ON DATABASE test_db TO postgres;