211 lines
7.5 KiB
Python
211 lines
7.5 KiB
Python
import argparse
|
|
import configparser
|
|
import datetime
|
|
import decimal
|
|
import json
|
|
import logging
|
|
import socket
|
|
import time
|
|
from logging.handlers import RotatingFileHandler
|
|
from pathlib import Path
|
|
from typing import Dict, List
|
|
|
|
import psycopg2
|
|
from psycopg2 import sql
|
|
|
|
script_dir = Path(__file__).resolve().parent
|
|
|
|
ENV_FILE = script_dir / '.env'
|
|
SEND_INTERVAL = 10
|
|
ITEM_SINGLE_RUN = 1000
|
|
|
|
LOG_FILE = script_dir / 'send_pg_data.log'
|
|
logging.basicConfig(
|
|
level=logging.INFO,
|
|
format='%(asctime)s - %(levelname)s - %(message)s',
|
|
handlers=[
|
|
RotatingFileHandler(LOG_FILE, maxBytes=30*1024*1024, backupCount=10),
|
|
logging.StreamHandler()
|
|
]
|
|
)
|
|
logger = logging.getLogger()
|
|
|
|
|
|
def read_config(config_file: Path):
|
|
config = configparser.ConfigParser()
|
|
config.read(config_file)
|
|
return config
|
|
|
|
|
|
def fetch_data(conn: psycopg2.connect, production_line: str, last_time: float, debug: bool = False):
|
|
last_time_str = datetime.datetime.fromtimestamp(last_time).astimezone().isoformat()
|
|
|
|
if production_line == 'gzjc_data':
|
|
query = f"""
|
|
SELECT * FROM {production_line}
|
|
WHERE time > %s
|
|
ORDER BY time ASC
|
|
LIMIT {ITEM_SINGLE_RUN}
|
|
"""
|
|
else:
|
|
query = f"""
|
|
SELECT * FROM {production_line}
|
|
WHERE created_at > %s
|
|
ORDER BY created_at ASC
|
|
LIMIT {ITEM_SINGLE_RUN}
|
|
"""
|
|
|
|
with conn.cursor() as cur:
|
|
cur.execute(query, (last_time_str,))
|
|
columns = [desc[0] for desc in cur.description]
|
|
data = []
|
|
for row in cur.fetchall():
|
|
row_data = {}
|
|
for col, value in zip(columns, row):
|
|
if col in ['id'] and production_line == 'gzjc_data':
|
|
continue
|
|
if col in ['span_begin', 'span_end']:
|
|
continue
|
|
if isinstance(value, decimal.Decimal):
|
|
row_data[col] = float(round(value, 2))
|
|
# elif isinstance(value, datetime.datetime):
|
|
# row_data[col] = value.astimezone().isoformat()
|
|
else:
|
|
row_data[col] = value
|
|
data.append(row_data)
|
|
|
|
# if not data:
|
|
# logger.warning(f"No data found for {production_line} from {last_time_str}, sending empty data")
|
|
# if debug:
|
|
# data = [{col: 0.0 for col in columns if col not in ['id', 'loraId', 'g4Id']}]
|
|
# data[0]["time"] = datetime.datetime.now().astimezone().isoformat()
|
|
# logger.info(f"[DEBUG] Sending empty data for {production_line}")
|
|
|
|
return data
|
|
|
|
|
|
def send_data(data: List[Dict], conn: psycopg2.connect, production_line: str, debug: bool = False):
|
|
try:
|
|
if not data:
|
|
logger.warning(f"[WARNING] No data found for {production_line}, skipping insert.")
|
|
return
|
|
|
|
cur = conn.cursor()
|
|
|
|
columns = data[0].keys()
|
|
column_names = ', '.join(columns)
|
|
placeholders = ', '.join(['%s'] * len(columns))
|
|
|
|
insert_query = sql.SQL("INSERT INTO {table} ({columns}) VALUES ({placeholders})").format(
|
|
table=sql.Identifier(production_line),
|
|
columns=sql.SQL(column_names),
|
|
placeholders=sql.SQL(placeholders)
|
|
)
|
|
# if debug:
|
|
# logger.info(f"[DEBUG] SQL query: {insert_query}")
|
|
|
|
values = [tuple(item.values()) for item in data]
|
|
cur.executemany(insert_query, values)
|
|
|
|
conn.commit()
|
|
except Exception as e:
|
|
logger.error(f"[ERROR] {e}")
|
|
if conn:
|
|
conn.rollback()
|
|
finally:
|
|
if cur:
|
|
cur.close()
|
|
|
|
|
|
def get_last_times(conn: psycopg2.connect, production_lines: List[str]) -> Dict[str, float]:
|
|
last_times = {}
|
|
for production_line in production_lines:
|
|
with conn.cursor() as cur:
|
|
if production_line == 'gzjc_data':
|
|
query = sql.SQL("SELECT MAX(time) FROM {table};").format(table=sql.Identifier(production_line))
|
|
else:
|
|
query = sql.SQL("SELECT MAX(created_at) FROM {table};").format(table=sql.Identifier(production_line))
|
|
cur.execute(query)
|
|
result = cur.fetchone()
|
|
if result and result[0] is not None:
|
|
last_times[production_line] = float(result[0].timestamp())
|
|
else:
|
|
last_times[production_line] = 0.0 # 设置默认时间戳为 0.0
|
|
return last_times
|
|
|
|
|
|
def main(args: argparse.Namespace):
|
|
debug_mode = args.debug
|
|
logger.info(f"Starting with debug mode: {debug_mode}")
|
|
|
|
config = read_config(ENV_FILE)
|
|
pg_config_src = config['SOURCE']
|
|
conn_src = psycopg2.connect(
|
|
host=pg_config_src['DB_HOST'],
|
|
port=pg_config_src['DB_PORT'],
|
|
user=pg_config_src['DB_UNAME'],
|
|
password=pg_config_src['DB_PASSWD'],
|
|
dbname=pg_config_src['DB_DB_NAME']
|
|
)
|
|
pg_config_dst = config['DESTINATION']
|
|
conn_dst = psycopg2.connect(
|
|
host=pg_config_dst['DB_HOST'],
|
|
port=pg_config_dst['DB_PORT'],
|
|
user=pg_config_dst['DB_UNAME'],
|
|
password=pg_config_dst['DB_PASSWD'],
|
|
dbname=pg_config_dst['DB_DB_NAME']
|
|
)
|
|
|
|
production_lines = ['gzjc_data'] + ['monitor_data_'+str(i) for i in range(1, 21)]
|
|
last_times = get_last_times(conn_dst, production_lines)
|
|
if debug_mode:
|
|
last_times_print = {key: datetime.datetime.fromtimestamp(
|
|
value).astimezone().isoformat() for key, value in last_times.items()}
|
|
logger.info(f"[DEBUG] Last times: {last_times_print}")
|
|
|
|
while True:
|
|
for line in production_lines:
|
|
data = []
|
|
try:
|
|
data = fetch_data(conn_src, line, last_times[line], debug=debug_mode)
|
|
except (psycopg2.Error, Exception) as e:
|
|
logger.exception(f"Source database error: {e}")
|
|
if conn_src.closed:
|
|
conn_src = psycopg2.connect(
|
|
host=pg_config_src['DB_HOST'],
|
|
port=pg_config_src['DB_PORT'],
|
|
user=pg_config_src['DB_UNAME'],
|
|
password=pg_config_src['DB_PASSWD'],
|
|
dbname=pg_config_src['DB_DB_NAME']
|
|
)
|
|
continue
|
|
if debug_mode:
|
|
logger.info(f"[DEBUG] Fetched data for production line {line} {len(data)}")
|
|
|
|
if data:
|
|
try:
|
|
send_data(data, conn_dst, line, debug=debug_mode)
|
|
last_times[line] = float(
|
|
data[-1]['time'].timestamp()) if line == 'gzjc_data' else float(data[-1]['created_at'].timestamp())
|
|
logger.info(
|
|
f"[INFO] {line} updated to {datetime.datetime.fromtimestamp(last_times[line]).astimezone().isoformat()}")
|
|
except (psycopg2.Error, Exception) as e:
|
|
logger.exception(f"Destination database error: {e}")
|
|
if conn_dst.closed:
|
|
conn_dst = psycopg2.connect(
|
|
host=pg_config_dst['DB_HOST'],
|
|
port=pg_config_dst['DB_PORT'],
|
|
user=pg_config_dst['DB_UNAME'],
|
|
password=pg_config_dst['DB_PASSWD'],
|
|
dbname=pg_config_dst['DB_DB_NAME']
|
|
)
|
|
|
|
time.sleep(SEND_INTERVAL)
|
|
|
|
|
|
if __name__ == "__main__":
|
|
argparse = argparse.ArgumentParser(description="Sync PostgreSQL data.")
|
|
argparse.add_argument("--debug", action="store_true", help="Debug mode")
|
|
args = argparse.parse_args()
|
|
main(args)
|