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

Auto-update: Tue Aug 6 22:19:38 PDT 2024

Browse source
This commit is contained in:
sanj 2024-08-06 22:19:38 -07:00
parent d8514b57a6
commit e527b0e391
1 changed files with 27 additions and 24 deletions
Show stats Download patch file Download diff file Expand all files Collapse all files

49
sijapi/classes.py
View file

@ -832,7 +832,6 @@ class APIConfig(BaseModel):
return [dict(r) for r in latest_result] # Convert Record objects to dictionaries
async def execute_write_query(self, query: str, *args, table_name: str):
local_ts_id = os.environ.get('TS_ID')
online_hosts = await self.get_online_hosts()
@ -850,11 +849,12 @@ class APIConfig(BaseModel):
if table_name in self.SPECIAL_TABLES:
result = await self._execute_special_table_write(conn, query, *args, table_name=table_name)
else:
# Modify the query to be an UPSERT operation
if query.strip().upper().startswith("INSERT"):
table_columns = await self.get_table_columns(conn, table_name)
primary_key = await self.get_primary_key(conn, table_name)
if primary_key:
if primary_key and len(table_columns) == len(args):
set_clause = ", ".join([f"{col} = EXCLUDED.{col}" for col in table_columns if col != primary_key])
query = f"""
INSERT INTO {table_name} ({', '.join(table_columns)})
@ -862,11 +862,12 @@ class APIConfig(BaseModel):
ON CONFLICT ({primary_key}) DO UPDATE SET
{set_clause}
"""
else:
err(f"Column count mismatch for table {table_name}. Columns: {len(table_columns)}, Values: {len(args)}")
return []
result = await conn.fetch(query, *args)
asyncio.create_task(self._sync_changes(table_name, await self.get_primary_key(conn, table_name)))
return result
except Exception as e:
err(f"Error executing write query on {pool_entry['ts_id']}: {str(e)}")
@ -878,6 +879,27 @@ class APIConfig(BaseModel):
return []
async def get_table_columns(self, conn, table_name: str) -> List[str]:
query = """
SELECT column_name
FROM information_schema.columns
WHERE table_schema = 'public' AND table_name = $1
ORDER BY ordinal_position
"""
columns = await conn.fetch(query, table_name)
return [col['column_name'] for col in columns]
async def get_primary_key(self, conn, table_name: str) -> str:
query = """
SELECT a.attname
FROM pg_index i
JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
WHERE i.indrelid = $1::regclass AND i.indisprimary
"""
result = await conn.fetchval(query, table_name)
return result
async def _sync_write_to_other_dbs(self, query: str, args: tuple, table_name: str):
local_ts_id = os.environ.get('TS_ID')
@ -1165,25 +1187,6 @@ class APIConfig(BaseModel):
return 0
async def get_table_columns(self, conn, table_name: str) -> List[str]:
query = """
SELECT column_name
FROM information_schema.columns
WHERE table_schema = 'public' AND table_name = $1
"""
columns = await conn.fetch(query, table_name)
return [col['column_name'] for col in columns]
async def get_primary_key(self, conn, table_name: str) -> str:
query = """
SELECT a.attname
FROM pg_index i
JOIN pg_attribute a ON a.attrelid = i.indrelid AND a.attnum = ANY(i.indkey)
WHERE i.indrelid = $1::regclass AND i.indisprimary
"""
result = await conn.fetchval(query, table_name)
return result
async def add_primary_keys_to_local_tables(self):