【小工具-模拟数据】Python实现自动针对Postgres SQL docker容器执行SQL脚本文件

1.工具简要说明

自动生成SQL脚本文件，自动针对docker运行容器形式的PG数据库进行数据模拟插入。可以以此为基础进行脚本改造，达到模拟大数据的用途。源码的写法比较土，后面有空可以优化一下代码。

2.工具源码

python文件名：init_db.py

#!/usr/bin/python
#create time: 2022/03/12  by epic
import os;

def get_db_container_id():
    print('*** get_db_container_id');
    docker_ps = os.popen('docker ps|grep postgres').readlines();
    container_info = docker_ps[0].split();
    return container_info[0];

def generate_sql_script_file(sql_script_file, number):
    os.system('rm -f ' + sql_script_file);
    os.system('touch ' + sql_script_file);
    for i in range(number):
        sno = str(i + 1) + '';
        sname = 'stu' + sno;
        sex = 1;
        os.system("echo delete from student where sno=\\'" + sno + "\\'\\; >> " + sql_script_file);
        os.system("echo insert into student values\\(\\'" + sno +"\\',\\'"+sname+"\\',"+str(sex)+"\\)\\; >> " + sql_script_file);

    return;

def copy_script_file_to_container(src_script_file,
                                  dest_script_file,
                                  container_id):

    print('*** copy_script_file_to_container: ' + src_script_file + ',' + dest_script_file + ',' + container_id);
    return os.system('docker cp ' + src_script_file + ' ' + container_id + ':' + dest_script_file);
#
#chown postgres:postgres /media/init_db_script.sql
#su - postgres -c "psql -d db1 -f /media/init_db_script.sql"
#
def generate_host_init_db_shell(host_init_db_shell,
                                container_init_db_shell,
                                container_init_db_script):

    print('*** generate_host_init_db_shell: ' + host_init_db_shell + ',' + container_init_db_shell + ',' + container_init_db_script);
    os.system('mv ' + host_init_db_shell + ' ' + host_init_db_shell + '.bak');
    os.system('touch ' + host_init_db_shell);
    os.system('echo chown postgres:postgres ' + container_init_db_shell + '>> ' + host_init_db_shell);
    os.system('echo chown postgres:postgres ' + container_init_db_script + '>> ' + host_init_db_shell);
    os.system('echo su - postgres -c \\"psql -d db1 -f ' + container_init_db_script + '\\" >> ' + host_init_db_shell);
    return;
#
#docker exec -it 1deab09a9907 /bin/sh -c 'sh /media/init_db.sh'
#
def exec_script_file(container_id,
                     container_init_db_shell):

    print('*** exec_script_file: ' + container_id + ',' + container_init_db_shell);
    os.system('docker exec -it ' + container_id + " /bin/sh -c ' sh " + container_init_db_shell + "'");
    return;

def main():
    print("^*********************--------------------***********************^");
    container_id = get_db_container_id();
    host_init_db_shell = "/root/init_db.sh";
    host_init_db_script = "/root/init_db_script.sql";
    container_init_db_shell = "/media/init_db.sh";
    container_init_db_script = "/media/init_db_script.sql";
    number = 2;
    generate_host_init_db_shell(host_init_db_shell, container_init_db_shell, container_init_db_script);
    copy_script_file_to_container(host_init_db_shell, container_init_db_shell, container_id);
    generate_sql_script_file(host_init_db_script, number);
    copy_script_file_to_container(host_init_db_script, container_init_db_script, container_id);
    exec_script_file(container_id, container_init_db_shell);
    print(">>>>>>insert number: " + str(number));

    print("^*********************--------------------***********************^");
    return;

main();

运行python脚本文件进行数据模拟，最后打印了模拟了多少条数据：

[root@localhost ~]# ./init_db.py
^*********************--------------------***********************^
*** get_db_container_id
*** generate_host_init_db_shell: /root/init_db.sh,/media/init_db.sh,/media/init_db_script.sql
*** copy_script_file_to_container: /root/init_db.sh,/media/init_db.sh,6825fcf46076
*** copy_script_file_to_container: /root/init_db_script.sql,/media/init_db_script.sql,6825fcf46076
*** exec_script_file: 6825fcf46076,/media/init_db.sh
DELETE 0
INSERT 0 1
DELETE 0
INSERT 0 1
>>>>>>insert number: 2
^*********************--------------------***********************^

-----------

运行python脚本后生成SQL脚本文件和shell脚本文件，并且自动拷贝到指定的容器中，接着自动在容器中执行SQL文件。下面2个文件是生成的中间文件：

文件1：生成shell脚本文件init_db.sh内容为：

chown postgres:postgres /media/init_db.sh
chown postgres:postgres /media/init_db_script.sql
su - postgres -c "psql -d db1 -f /media/init_db_script.sql"

文件2：生成的SQL文件init_db_script.sql内容为：

[root@localhost ~]# cat init_db_script.sql
delete from student where sno='1';
insert into student values('1','stu1',1);
delete from student where sno='2';
insert into student values('2','stu2',1);

3. 运行后的效果

[root@localhost ~]# docker exec -it 6825fcf46076 /bin/sh
# su - postgres
postgres@6825fcf46076:~$ psql
psql (10.17 (Debian 10.17-1.pgdg90+1))
Type "help" for help.

postgres=# \c db1
You are now connected to database "db1" as user "postgres".
db1=# truncate student;
TRUNCATE TABLE
db1=# select * from student;
 sno | sname | sex
-----+-------+-----
 1   | stu1  |   1
 2   | stu2  |   1
(2 rows)