sij/sijapi
1
1
Fork 0
Code Issues 3 Pull requests Projects Releases Packages Wiki Activity Actions

Auto-update: Thu Jul 25 00:16:07 PDT 2024

Browse source
This commit is contained in:
sanj 2024-07-25 00:16:07 -07:00
parent 03ccf529c9
commit 8775c05927
1 changed files with 60 additions and 45 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

105
sijapi/classes.py
View file

@ -5,6 +5,7 @@ import yaml
import math import math
import os import os
import re import re
import traceback
import aiofiles import aiofiles
import aiohttp import aiohttp
import asyncpg import asyncpg
@ -339,53 +340,67 @@ class APIConfig(BaseModel):
async def pull_changes(self, source_pool_entry: Dict[str, Any] = None): async def pull_changes(self, source_pool_entry: Dict[str, Any] = None):
if source_pool_entry is None: try:
source_pool_entry = await self.get_default_source() if source_pool_entry is None:
source_pool_entry = await self.get_default_source()
if source_pool_entry is None:
err("No available source for pulling changes") if source_pool_entry is None:
return err("No available source for pulling changes")
return
async with self.get_connection(source_pool_entry) as source_conn:
async with self.get_connection() as dest_conn: async with self.get_connection(source_pool_entry) as source_conn:
tables = await source_conn.fetch( async with self.get_connection() as dest_conn:
"SELECT tablename FROM pg_tables WHERE schemaname = 'public'" tables = await source_conn.fetch(
) "SELECT tablename FROM pg_tables WHERE schemaname = 'public'"
for table in tables: )
table_name = table['tablename'] for table in tables:
info(f"Processing table: {table_name}") table_name = table['tablename']
info(f"Processing table: {table_name}")
# Get primary key column(s)
pk_columns = await source_conn.fetch(""" # Get primary key column(s)
SELECT a.attname pk_columns = await source_conn.fetch("""
FROM pg_index i SELECT a.attname
JOIN pg_attribute a ON a.attrelid = i.indrelid FROM pg_index i
AND a.attnum = ANY(i.indkey) JOIN pg_attribute a ON a.attrelid = i.indrelid
WHERE i.indrelid = $1::regclass AND a.attnum = ANY(i.indkey)
AND i.indisprimary; WHERE i.indrelid = $1::regclass
""", table_name) AND i.indisprimary;
""", table_name)
pk_cols = [col['attname'] for col in pk_columns]
if not pk_cols: pk_cols = [col['attname'] for col in pk_columns]
warn(f"No primary key found for table {table_name}. Skipping.") info(f"Primary key columns for {table_name}: {pk_cols}")
continue if not pk_cols:
warn(f"No primary key found for table {table_name}. Skipping.")
continue
# Fetch all rows from the source table # Fetch all rows from the source table
rows = await source_conn.fetch(f"SELECT * FROM {table_name}") rows = await source_conn.fetch(f"SELECT * FROM {table_name}")
if rows: info(f"Fetched {len(rows)} rows from {table_name}")
columns = rows[0].keys() if rows:
# Upsert records to the destination table columns = list(rows[0].keys())
for row in rows: info(f"Columns for {table_name}: {columns}")
await dest_conn.execute(f""" # Upsert records to the destination table
INSERT INTO {table_name} ({', '.join(columns)}) for row in rows:
VALUES ({', '.join(f'${i+1}' for i in range(len(columns)))}) try:
ON CONFLICT ({', '.join(pk_cols)}) DO UPDATE SET query = f"""
{', '.join(f"{col} = EXCLUDED.{col}" for col in columns if col not in pk_cols)} INSERT INTO {table_name} ({', '.join(columns)})
""", *[row[col] for col in columns]) VALUES ({', '.join(f'${i+1}' for i in range(len(columns)))})
ON CONFLICT ({', '.join(pk_cols)}) DO UPDATE SET
info(f"Completed processing table: {table_name}") {', '.join(f"{col} = EXCLUDED.{col}" for col in columns if col not in pk_cols)}
"""
info(f"Executing query: {query}")
info(f"With values: {[row[col] for col in columns]}")
await dest_conn.execute(query, *[row[col] for col in columns])
except Exception as e:
err(f"Error processing row in {table_name}: {str(e)}")
err(f"Problematic row: {row}")
info(f"Completed processing table: {table_name}")
info(f"Successfully pulled changes from {source_pool_entry['ts_ip']}") info(f"Successfully pulled changes from {source_pool_entry['ts_ip']}")
except Exception as e:
err(f"Unexpected error in pull_changes: {str(e)}")
err(f"Traceback: {traceback.format_exc()}")