From 99e2729fde4aa15f2a9805964207e841dc60485b Mon Sep 17 00:00:00 2001
From: sanj <67624670+iodrift@users.noreply.github.com>
Date: Tue, 23 Jul 2024 21:54:54 -0700
Subject: [PATCH] Auto-update: Tue Jul 23 21:54:54 PDT 2024

---
 sijapi/classes.py | 31 ++++++++++++++++++-------------
 1 file changed, 18 insertions(+), 13 deletions(-)

diff --git a/sijapi/classes.py b/sijapi/classes.py
index cc8877c..332f579 100644
--- a/sijapi/classes.py
+++ b/sijapi/classes.py
@@ -277,7 +277,8 @@ class APIConfig(BaseModel):
 
     @property
     def local_db(self):
-        return self.POOL[0]
+        ts_id = os.environ.get('TS_ID')
+        return next((db for db in self.POOL if db['ts_id'] == ts_id), None)
 
     @asynccontextmanager
     async def get_connection(self, pool_entry: Dict[str, Any] = None):
@@ -319,9 +320,24 @@ class APIConfig(BaseModel):
             for conn in connections:
                 await conn.__aexit__(None, None, None)
 
+    async def get_default_source(self):
+        local = self.local_db
+        for db in self.POOL:
+            if db != local:
+                try:
+                    async with self.get_connection(db):
+                        return db
+                except:
+                    continue
+        return None
+
     async def pull_changes(self, source_pool_entry: Dict[str, Any] = None):
         if source_pool_entry is None:
-            source_pool_entry = self.POOL[1]  # Default to the second database in the pool
+            source_pool_entry = await self.get_default_source()
+        
+        if source_pool_entry is None:
+            logger.error("No available source for pulling changes")
+            return
         
         logger = Logger("DatabaseReplication")
         async with self.get_connection(source_pool_entry) as source_conn:
@@ -440,17 +456,6 @@ class APIConfig(BaseModel):
                     await conn.execute(f"ALTER TABLE {target_con['table_name']} DROP CONSTRAINT {con_name}")
 
 
-    async def apply_schema(self, pool_entry: Dict[str, Any], schema):
-        async with self.get_connection(pool_entry) as conn:
-            # This is a simplified version. You'd need to handle creating/altering tables,
-            # adding/removing columns, changing data types, etc.
-            for table in schema:
-                await conn.execute(f"""
-                    CREATE TABLE IF NOT EXISTS {table['table_name']} (
-                        {table['column_name']} {table['data_type']}
-                    )
-                """)
-
 
 
 class Location(BaseModel):