fs-ops/code/tool-db2word/main.py

import os
import logging
import configparser
import psycopg2
from docx import Document
from docx.shared import Inches, Pt, RGBColor
from docx.enum.text import WD_LINE_SPACING
from docx.enum.table import WD_TABLE_ALIGNMENT, WD_CELL_VERTICAL_ALIGNMENT
from docx.oxml import parse_xml
from docx.oxml.ns import qn, nsdecls

LOG_DIR = 'log'
DOC_DIR = 'output'


def set_global_normal_style(doc):
    style = doc.styles['Normal']
    style.font.name = 'Times New Roman'  # 英文、数字字体
    style._element.rPr.rFonts.set(qn('w:eastAsia'), u'宋体')  # 中文字体
    style.font.size = Pt(10.5)  # 五号
    style.paragraph_format.space_before = 0  # 段前
    style.paragraph_format.space_after = 0  # 段后
    style.paragraph_format.line_spacing_rule = WD_LINE_SPACING.SINGLE  # 单倍行距


def set_paragraph_space(p):
    # 设置间距
    p.paragraph_format.space_before = 0  # 段前
    p.paragraph_format.space_after = 0  # 段后
    p.paragraph_format.line_spacing = Pt(20)  # 设置段落行距为20磅


def set_paragraph_format(p):
    set_paragraph_space(p)
    # 设置首行缩进2个字符
    p.paragraph_format.first_line_indent = Inches(0.3)


def set_heading_format(h, r):
    set_paragraph_space(h)
    r.font.name = 'Times New Roman'
    r._element.rPr.rFonts.set(qn('w:eastAsia'), u'宋体')
    r.font.size = Pt(10.5)  # 五号
    r.font.color.rgb = RGBColor(0, 0, 0)
    r.bold = True


def set_table_header_background_color(table, cols_count, bg_color):
    shading_list = locals()
    for i in range(cols_count):
        shading_list['shading_elm_' + str(i)] = parse_xml(
            r'<w:shd {} w:fill="{bgColor}"/>'.format(nsdecls('w'), bgColor=bg_color))
        table.rows[0].cells[i]._tc.get_or_add_tcPr().append(shading_list['shading_elm_' + str(i)])


def insert_table_to_doc(doc, t_fields):
    # 在文档中添加一个空表格
    table = doc.add_table(rows=0, cols=6, style='Table Grid')

    # 添加表格标题
    heading_cells = table.add_row().cells
    heading_cells[0].text = '序号'
    heading_cells[1].text = '字段名称'
    heading_cells[2].text = '类型'
    heading_cells[3].text = '长度'
    heading_cells[4].text = '必填'
    heading_cells[5].text = '描述'

    # 设置表头背景色（浅灰色）
    set_table_header_background_color(table, 6, '#D9D9D9')

    # 循环添加数据
    for field in t_fields:
        row_cells = table.add_row().cells
        row_cells[0].text = f'{field["field_sn"]}'
        row_cells[1].text = field["field_name"]
        row_cells[2].text = field["field_type"]
        row_cells[3].text = f'{field["field_length"]}' if field["field_length"] is not None else ''
        row_cells[4].text = field["field_required"]
        row_cells[5].text = f'{field["field_comment"]}' if field["field_comment"] is not None else ''

    # 设置表格垂直对齐方式
    table.alignment = WD_TABLE_ALIGNMENT.CENTER

    # 设置单元格对齐方式
    for row in table.rows:
        for cell in row.cells:
            cell.vertical_alignment = WD_CELL_VERTICAL_ALIGNMENT.CENTER  # 垂直居中
            cell.paragraphs[0].paragraph_format.alignment = WD_TABLE_ALIGNMENT.CENTER  # 水平居中


def write_tables_to_word():
    # 创建一个新的Word文档
    doc = Document()
    set_global_normal_style(doc)

    h = doc.add_heading(level=1)
    r = h.add_run('一、数据表结构设计')
    set_heading_format(h, r)

    for t_data in tables_data:
        t_fields = t_data['table_fields']
        p = doc.add_paragraph('表名：{}    字段数目：{}'.format(t_data['table_name'], len(t_fields)))
        set_paragraph_format(p)
        p = doc.add_paragraph('描述：{}'.format(t_data['table_comment'] if t_data['table_comment'] is not None else ''))
        set_paragraph_format(p)
        insert_table_to_doc(doc, t_fields)  # 将当前表的结构说明添加到文档
        doc.add_paragraph()  # 每个表后面加一个空行

    # 保存文档
    report_file_name = '数据库表结构设计.docx'
    report_file = f'{DOC_DIR}/{report_file_name}'
    doc.save(report_file)
    logging.info('数据库表说明文档创建成功：{}'.format(report_file))
    print('数据库表说明文档创建成功：{}'.format(report_file))


def create_postgres_connection():
    # 获取 postgres 配置参数
    ch_cfg = config['postgres']
    host = os.environ.get('PG_HOST', ch_cfg['host'])
    port = int(os.environ.get('PG_PORT', ch_cfg.getint('port')))
    username = os.environ.get('PG_USERNAME', ch_cfg['username'])
    password = os.environ.get('PG_PASSWORD', ch_cfg['password'])
    database = os.environ.get('PG_DATABASE', ch_cfg['database'])
    # 建立连接
    try:
        conn = psycopg2.connect(
            host=host,
            port=port,
            user=username,
            password=password,
            dbname=database
        )
        logging.info("成功连接到数据库")
        return conn
    except psycopg2.Error as e:
        logging.error("数据库连接错误:{}".format(e))
        raise RuntimeError('数据库连接错误:', error)


def querystring_all_tables():
    return '''
select relname as table_name,(select description from pg_description where objoid=oid and objsubid=0) 
as table_comment from pg_class where relkind ='r' and relname NOT LIKE 'pg%' AND relname NOT LIKE 'sql_%'
order by table_name;
    '''


def querystring_fields_in_table(t_name):
    return '''
SELECT
    ROW_NUMBER() OVER () AS 序号,
    a.attname as 字段名称,
    CASE
        WHEN format_type(a.atttypid, a.atttypmod) ~ '^character varying\(' THEN 'varchar'
        WHEN format_type(a.atttypid, a.atttypmod) ~ '^numeric\(' THEN 'numeric'
        WHEN format_type(a.atttypid, a.atttypmod) ~ '^timestamp' THEN 'timestamp'
        ELSE format_type(a.atttypid, a.atttypmod)
    END as 类型,
    CASE
        WHEN a.atttypmod - 4 > 0 THEN
            CASE
                WHEN format_type(a.atttypid, a.atttypmod) ~ '^numeric\(' THEN
                    substring(format_type(a.atttypid, a.atttypmod) from '\((\d+,\d+)\)')
                WHEN format_type(a.atttypid, a.atttypmod) ~ '^timestamp' THEN NULL
                ELSE '' || a.atttypmod - 4
            END
        ELSE NULL
    END as 长度,
    (CASE WHEN a.attnotnull = true THEN '是' ELSE '否' END) as 必填,
    col_description(a.attrelid, a.attnum) as 描述
FROM pg_attribute a
WHERE attstattarget = -1 AND attrelid = (SELECT oid FROM pg_class WHERE relname = '{}');
    '''.format(t_name)


def try_create_dir(dir_name):
    # 判断目录是否存在，如果不存在，则创建目录
    if not os.path.exists(dir_name):
        print(f"目录'{dir_name}'不存在，即将创建...")
        os.makedirs(dir_name)
        print(f"目录'{dir_name}'创建成功！")


if __name__ == '__main__':
    try:
        try_create_dir(DOC_DIR)

        try_create_dir(LOG_DIR)
        logging.basicConfig(filename='{}/runtime.log'.format(LOG_DIR), level=logging.INFO,
                            format='%(asctime)s.%(msecs)03d - %(levelname)s: %(message)s',
                            datefmt='%Y-%m-%d %H:%M:%S')

        config = configparser.ConfigParser()
        config.read('config.ini')

        connection = create_postgres_connection()

        cursor = connection.cursor()
        cursor.execute(querystring_all_tables())
        table_rows = cursor.fetchall()

        tables_data = []
        for (table_name, table_comment) in table_rows:
            cursor.execute(querystring_fields_in_table(table_name))
            field_rows = cursor.fetchall()
            fields_data = []
            for field in field_rows:
                fields_data.append({
                    'field_sn': field[0],  # 序号
                    'field_name': field[1],  # 字段名称
                    'field_type': field[2],  # 类型
                    'field_length': field[3],  # 长度
                    'field_required': field[4],  # 必填
                    'field_comment': field[5]  # 描述
                })
            tables_data.append({
                'table_name': table_name,
                'table_comment': table_comment,
                'table_fields': fields_data
            })

        cursor.close()

        write_tables_to_word()  # 生成 word 文档
    except Exception as error:
        print('程序运行出错：', error)
        logging.error('程序运行出错：{}'.format(error))

    # 关闭数据库连接
    finally:
        if connection:
            connection.close()
            logging.info("数据库连接已关闭")