diff --git a/sijapi/classes.py b/sijapi/classes.py index 78f335f..ce560b6 100644 --- a/sijapi/classes.py +++ b/sijapi/classes.py @@ -14,7 +14,7 @@ import reverse_geocoder as rg from pathlib import Path from typing import Any, Dict, List, Optional, Tuple, Union, TypeVar from dotenv import load_dotenv -from pydantic import BaseModel, Field, create_model, PrivateAttr +from pydantic import BaseModel, Field, create_model from concurrent.futures import ThreadPoolExecutor from contextlib import asynccontextmanager from datetime import datetime, timedelta, timezone @@ -165,32 +165,6 @@ class Configuration(BaseModel): -class DatabasePool: - def __init__(self): - self.pools = {} - - async def get_connection(self, pool_entry): - pool_key = f"{pool_entry['ts_ip']}:{pool_entry['db_port']}" - if pool_key not in self.pools: - self.pools[pool_key] = await asyncpg.create_pool( - host=pool_entry['ts_ip'], - port=pool_entry['db_port'], - user=pool_entry['db_user'], - password=pool_entry['db_pass'], - database=pool_entry['db_name'], - min_size=1, - max_size=10 - ) - return await self.pools[pool_key].acquire() - - async def release_connection(self, pool_entry, connection): - pool_key = f"{pool_entry['ts_ip']}:{pool_entry['db_port']}" - await self.pools[pool_key].release(connection) - - async def close_all(self): - for pool in self.pools.values(): - await pool.close() - class APIConfig(BaseModel): HOST: str PORT: int @@ -204,18 +178,6 @@ class APIConfig(BaseModel): TZ: str KEYS: List[str] GARBAGE: Dict[str, Any] - _db_pool: DatabasePool = PrivateAttr(default_factory=DatabasePool) - - class Config: - arbitrary_types_allowed = True - - @property - def db_pool(self): - return self._db_pool - - @property - def db_pool(self): - return self._db_pool @classmethod def load(cls, config_path: Union[str, Path], secrets_path: Union[str, Path]): @@ -258,7 +220,6 @@ class APIConfig(BaseModel): config_data['EXTENSIONS'] = cls._create_dynamic_config(config_data.get('EXTENSIONS', {}), 'DynamicExtensionsConfig') return cls(**config_data) - @classmethod def _create_dynamic_config(cls, data: Dict[str, Any], model_name: str): fields = {} @@ -337,98 +298,78 @@ class APIConfig(BaseModel): if pool_entry is None: pool_entry = self.local_db - for attempt in range(3): # Retry up to 3 times + info(f"Attempting to connect to database: {pool_entry}") + try: + conn = await asyncpg.connect( + host=pool_entry['ts_ip'], + port=pool_entry['db_port'], + user=pool_entry['db_user'], + password=pool_entry['db_pass'], + database=pool_entry['db_name'] + ) try: - conn = await asyncpg.connect( - host=pool_entry['ts_ip'], - port=pool_entry['db_port'], - user=pool_entry['db_user'], - password=pool_entry['db_pass'], - database=pool_entry['db_name'] - ) - try: - yield conn - finally: - await conn.close() - return - except asyncpg.exceptions.CannotConnectNowError: - if attempt < 2: # Don't sleep on the last attempt - await asyncio.sleep(1) # Wait before retrying - except Exception as e: - err(f"Failed to connect to database: {pool_entry['ts_ip']}:{pool_entry['db_port']}") - err(f"Error: {str(e)}") - if attempt == 2: # Raise the exception on the last attempt - raise + yield conn + finally: + await conn.close() + except Exception as e: + warn(f"Failed to connect to database: {pool_entry['ts_ip']}:{pool_entry['db_port']}") + err(f"Error: {str(e)}") + raise - raise Exception(f"Failed to connect to database after 3 attempts: {pool_entry['ts_ip']}:{pool_entry['db_port']}") async def initialize_sync(self): for pool_entry in self.POOL: - for attempt in range(3): # Retry up to 3 times - try: - async with self.get_connection(pool_entry) as conn: - tables = await conn.fetch(""" - SELECT tablename FROM pg_tables - WHERE schemaname = 'public' + try: + async with self.get_connection(pool_entry) as conn: + tables = await conn.fetch(""" + SELECT tablename FROM pg_tables + WHERE schemaname = 'public' + """) + + for table in tables: + table_name = table['tablename'] + # Add version and server_id columns if they don't exist + await conn.execute(f""" + DO $$ + BEGIN + BEGIN + ALTER TABLE "{table_name}" + ADD COLUMN IF NOT EXISTS version INTEGER DEFAULT 1, + ADD COLUMN IF NOT EXISTS server_id TEXT DEFAULT '{os.environ.get('TS_ID')}'; + EXCEPTION + WHEN duplicate_column THEN + -- Do nothing, column already exists + END; + END $$; """) - - for table in tables: - table_name = table['tablename'] - await self.ensure_sync_columns(conn, table_name) - await self.create_sync_trigger(conn, table_name) - info(f"Sync initialization complete for {pool_entry['ts_ip']}. All tables now have version and server_id columns with appropriate triggers.") - break # If successful, break the retry loop - except asyncpg.exceptions.ConnectionFailureError: - err(f"Failed to connect to database during initialization: {pool_entry['ts_ip']}:{pool_entry['db_port']}") - if attempt < 2: # Don't sleep on the last attempt - await asyncio.sleep(1) # Wait before retrying - except Exception as e: - err(f"Error initializing sync for {pool_entry['ts_ip']}: {str(e)}") - err(f"Traceback: {traceback.format_exc()}") - break # Don't retry for unexpected errors + # Create or replace the trigger function + await conn.execute(f""" + CREATE OR REPLACE FUNCTION update_version_and_server_id() + RETURNS TRIGGER AS $$ + BEGIN + NEW.version = COALESCE(OLD.version, 0) + 1; + NEW.server_id = '{os.environ.get('TS_ID')}'; + RETURN NEW; + END; + $$ LANGUAGE plpgsql; + """) - async def ensure_sync_columns(self, conn, table_name): - try: - await conn.execute(f""" - DO $$ - BEGIN - BEGIN - ALTER TABLE "{table_name}" - ADD COLUMN IF NOT EXISTS version INTEGER DEFAULT 1, - ADD COLUMN IF NOT EXISTS server_id TEXT DEFAULT '{os.environ.get('TS_ID')}'; - EXCEPTION - WHEN duplicate_column THEN - NULL; -- Silently handle duplicate column - END; - END $$; - """) - info(f"Ensured sync columns for table {table_name}") - except Exception as e: - err(f"Error ensuring sync columns for table {table_name}: {str(e)}") - err(f"Traceback: {traceback.format_exc()}") + # Create the trigger if it doesn't exist + await conn.execute(f""" + DO $$ + BEGIN + IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'update_version_and_server_id_trigger' AND tgrelid = '{table_name}'::regclass) THEN + CREATE TRIGGER update_version_and_server_id_trigger + BEFORE INSERT OR UPDATE ON "{table_name}" + FOR EACH ROW EXECUTE FUNCTION update_version_and_server_id(); + END IF; + END $$; + """) - - async def create_sync_trigger(self, conn, table_name): - await conn.execute(f""" - CREATE OR REPLACE FUNCTION update_version_and_server_id() - RETURNS TRIGGER AS $$ - BEGIN - NEW.version = COALESCE(OLD.version, 0) + 1; - NEW.server_id = '{os.environ.get('TS_ID')}'; - RETURN NEW; - END; - $$ LANGUAGE plpgsql; - - DO $$ - BEGIN - IF NOT EXISTS (SELECT 1 FROM pg_trigger WHERE tgname = 'update_version_and_server_id_trigger' AND tgrelid = '{table_name}'::regclass) THEN - CREATE TRIGGER update_version_and_server_id_trigger - BEFORE INSERT OR UPDATE ON "{table_name}" - FOR EACH ROW EXECUTE FUNCTION update_version_and_server_id(); - END IF; - END $$; - """) + info(f"Sync initialization complete for {pool_entry['ts_ip']}. All tables now have version and server_id columns with appropriate triggers.") + except Exception as e: + err(f"Error initializing sync for {pool_entry['ts_ip']}: {str(e)}") async def get_most_recent_source(self): most_recent_source = None @@ -438,52 +379,25 @@ class APIConfig(BaseModel): if pool_entry['ts_id'] == os.environ.get('TS_ID'): continue - for _ in range(3): # Retry up to 3 times - try: - async with self.get_connection(pool_entry) as conn: - # Check if the version column exists in any table - version_exists = await conn.fetchval(""" - SELECT EXISTS ( - SELECT 1 - FROM information_schema.columns - WHERE table_schema = 'public' - AND column_name = 'version' - ) - """) - - if not version_exists: - info(f"Version column does not exist in any table for {pool_entry['ts_id']}") - break # Move to the next pool entry - - version = await conn.fetchval(""" - SELECT COALESCE(MAX(version), -1) - FROM ( - SELECT MAX(version) as version - FROM information_schema.columns - WHERE table_schema = 'public' - AND column_name = 'version' - AND is_updatable = 'YES' - ) as subquery - """) - info(f"Max version for {pool_entry['ts_id']}: {version}") - if version > max_version: - max_version = version - most_recent_source = pool_entry - break # If successful, break the retry loop - except asyncpg.exceptions.PostgresError as e: - err(f"Error checking version for {pool_entry['ts_id']}: {str(e)}") - await asyncio.sleep(1) # Wait before retrying - except Exception as e: - err(f"Unexpected error for {pool_entry['ts_id']}: {str(e)}") - break # Don't retry for unexpected errors - - if most_recent_source: - info(f"Most recent source: {most_recent_source['ts_id']} with version {max_version}") - else: - info("No valid source found with version information") + try: + async with self.get_connection(pool_entry) as conn: + version = await conn.fetchval(""" + SELECT COALESCE(MAX(version), -1) FROM ( + SELECT MAX(version) as version FROM information_schema.columns + WHERE table_schema = 'public' AND column_name = 'version' + ) as subquery + """) + if version > max_version: + max_version = version + most_recent_source = pool_entry + except Exception as e: + err(f"Error checking version for {pool_entry['ts_id']}: {str(e)}") return most_recent_source + + + async def pull_changes(self, source_pool_entry, batch_size=10000): if source_pool_entry['ts_id'] == os.environ.get('TS_ID'): info("Skipping self-sync") @@ -500,10 +414,6 @@ class APIConfig(BaseModel): try: async with self.get_connection(source_pool_entry) as source_conn: async with self.get_connection(self.local_db) as dest_conn: - # Sync schema first - schema_changes = await self.detect_schema_changes(source_conn, dest_conn) - await self.apply_schema_changes(dest_conn, schema_changes) - tables = await source_conn.fetch(""" SELECT tablename FROM pg_tables WHERE schemaname = 'public' @@ -557,22 +467,12 @@ class APIConfig(BaseModel): columns = changes[0].keys() await conn.copy_records_to_table(temp_table_name, records=[tuple(change[col] for col in columns) for change in changes]) - # Perform upsert with spatial awareness + # Perform upsert result = await conn.execute(f""" INSERT INTO "{table_name}" - SELECT tc.* - FROM {temp_table_name} tc - LEFT JOIN "{table_name}" t ON t.id = tc.id - WHERE t.id IS NULL + SELECT * FROM {temp_table_name} ON CONFLICT (id) DO UPDATE SET {', '.join(f"{col} = EXCLUDED.{col}" for col in columns if col != 'id')} - WHERE ( - CASE - WHEN "{table_name}".geometry IS NOT NULL AND EXCLUDED.geometry IS NOT NULL - THEN NOT ST_Equals("{table_name}".geometry, EXCLUDED.geometry) - ELSE FALSE - END - ) OR {' OR '.join(f"COALESCE({col} <> EXCLUDED.{col}, TRUE)" for col in columns if col not in ['id', 'geometry'])} """) # Parse the result to get the number of affected rows @@ -583,6 +483,7 @@ class APIConfig(BaseModel): # Ensure temporary table is dropped await conn.execute(f"DROP TABLE IF EXISTS {temp_table_name}") + async def push_changes_to_all(self): for pool_entry in self.POOL: if pool_entry['ts_id'] != os.environ.get('TS_ID'): @@ -591,14 +492,11 @@ class APIConfig(BaseModel): except Exception as e: err(f"Error pushing changes to {pool_entry['ts_id']}: {str(e)}") - async def push_changes_to_one(self, pool_entry, batch_size=10000): + + async def push_changes_to_one(self, pool_entry): try: async with self.get_connection() as local_conn: async with self.get_connection(pool_entry) as remote_conn: - # Sync schema first - schema_changes = await self.detect_schema_changes(local_conn, remote_conn) - await self.apply_schema_changes(remote_conn, schema_changes) - tables = await local_conn.fetch(""" SELECT tablename FROM pg_tables WHERE schemaname = 'public' @@ -608,29 +506,49 @@ class APIConfig(BaseModel): table_name = table['tablename'] last_synced_version = await self.get_last_synced_version(remote_conn, table_name, os.environ.get('TS_ID')) - while True: - changes = await local_conn.fetch(f""" - SELECT * FROM "{table_name}" - WHERE version > $1 AND server_id = $2 - ORDER BY version ASC - LIMIT $3 - """, last_synced_version, os.environ.get('TS_ID'), batch_size) + changes = await local_conn.fetch(f""" + SELECT * FROM "{table_name}" + WHERE version > $1 AND server_id = $2 + ORDER BY version ASC + """, last_synced_version, os.environ.get('TS_ID')) + + for change in changes: + columns = list(change.keys()) + values = [change[col] for col in columns] + placeholders = [f'${i+1}' for i in range(len(columns))] - if not changes: - break - - changes_count = await self.apply_batch_changes(remote_conn, table_name, changes) + insert_query = f""" + INSERT INTO "{table_name}" ({', '.join(columns)}) + VALUES ({', '.join(placeholders)}) + ON CONFLICT (id) DO UPDATE SET + {', '.join(f"{col} = EXCLUDED.{col}" for col in columns if col != 'id')} + """ - last_synced_version = changes[-1]['version'] - await self.update_last_synced_version(remote_conn, table_name, os.environ.get('TS_ID'), last_synced_version) - - info(f"Pushed batch for {table_name}: {changes_count} changes to {pool_entry['ts_id']}") + await remote_conn.execute(insert_query, *values) + + if changes: + await self.update_last_synced_version(remote_conn, table_name, os.environ.get('TS_ID'), changes[-1]['version']) info(f"Successfully pushed changes to {pool_entry['ts_id']}") except Exception as e: err(f"Error pushing changes to {pool_entry['ts_id']}: {str(e)}") err(f"Traceback: {traceback.format_exc()}") + async def ensure_sync_columns(self, conn, table_name): + await conn.execute(f""" + DO $$ + BEGIN + BEGIN + ALTER TABLE "{table_name}" + ADD COLUMN IF NOT EXISTS version INTEGER DEFAULT 1, + ADD COLUMN IF NOT EXISTS server_id TEXT DEFAULT '{os.environ.get('TS_ID')}'; + EXCEPTION + WHEN duplicate_column THEN + -- Do nothing, column already exists + END; + END $$; + """) + async def get_last_synced_version(self, conn, table_name, server_id): return await conn.fetchval(f""" SELECT COALESCE(MAX(version), 0) @@ -666,85 +584,7 @@ class APIConfig(BaseModel): END $$; """) - async def detect_schema_changes(self, source_conn, dest_conn): - schema_changes = { - 'new_tables': [], - 'new_columns': {} - } - - # Detect new tables - source_tables = await source_conn.fetch("SELECT tablename FROM pg_tables WHERE schemaname = 'public'") - dest_tables = await dest_conn.fetch("SELECT tablename FROM pg_tables WHERE schemaname = 'public'") - - source_table_names = set(table['tablename'] for table in source_tables) - dest_table_names = set(table['tablename'] for table in dest_tables) - - new_tables = source_table_names - dest_table_names - schema_changes['new_tables'] = list(new_tables) - - # Detect new columns - for table_name in source_table_names: - if table_name in dest_table_names: - source_columns = await source_conn.fetch(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}'") - dest_columns = await dest_conn.fetch(f"SELECT column_name, data_type FROM information_schema.columns WHERE table_name = '{table_name}'") - - source_column_names = set(column['column_name'] for column in source_columns) - dest_column_names = set(column['column_name'] for column in dest_columns) - - new_columns = source_column_names - dest_column_names - if new_columns: - schema_changes['new_columns'][table_name] = [ - {'name': column['column_name'], 'type': column['data_type']} - for column in source_columns if column['column_name'] in new_columns - ] - - return schema_changes - async def apply_schema_changes(self, conn, schema_changes): - for table_name in schema_changes['new_tables']: - create_table_sql = await self.get_table_creation_sql(conn, table_name) - await conn.execute(create_table_sql) - info(f"Created new table: {table_name}") - - for table_name, columns in schema_changes['new_columns'].items(): - for column in columns: - await conn.execute(f""" - ALTER TABLE "{table_name}" - ADD COLUMN IF NOT EXISTS {column['name']} {column['type']} - """) - info(f"Added new column {column['name']} to table {table_name}") - - async def get_table_creation_sql(self, conn, table_name): - create_table_sql = await conn.fetchval(f""" - SELECT pg_get_tabledef('{table_name}'::regclass::oid) - """) - return create_table_sql - - async def table_exists(self, conn, table_name): - exists = await conn.fetchval(f""" - SELECT EXISTS ( - SELECT FROM information_schema.tables - WHERE table_schema = 'public' - AND table_name = $1 - ) - """, table_name) - return exists - - async def column_exists(self, conn, table_name, column_name): - exists = await conn.fetchval(f""" - SELECT EXISTS ( - SELECT FROM information_schema.columns - WHERE table_schema = 'public' - AND table_name = $1 - AND column_name = $2 - ) - """, table_name, column_name) - return exists - - - async def close_db_pools(self): - if self._db_pool: - await self._db_pool.close_all() class Location(BaseModel):