sunholo 0.123.0__py3-none-any.whl → 0.123.2__py3-none-any.whl

sunholo/database/alloydb_client.py

@@ -531,128 +531,182 @@ class AlloyDBClient:
         
         return await self.engine._afetch(query=sql_statement)
 
-async def create_table_from_schema(self, table_name: str, schema_data: dict, users: list = None):
-    """
-    Creates or ensures a table exists based on the structure of the provided schema data.
-    
-    Args:
-        table_name (str): Name of the table to create
-        schema_data (dict): Data structure that matches the expected schema
-        users (list, optional): List of users to grant permissions to
+    async def check_connection(self):
+        """
+        Checks if the database connection is still valid.
         
-    Returns:
-        Result of SQL execution
-    """
-    # Generate column definitions from schema data
-    columns = []
-    for key, value in schema_data.items():
-        if isinstance(value, dict):
-            # For nested objects, store as JSONB
-            columns.append(f'"{key}" JSONB')
-        elif isinstance(value, list):
-            # For arrays, store as JSONB
-            columns.append(f'"{key}" JSONB')
-        elif isinstance(value, int):
-            columns.append(f'"{key}" INTEGER')
-        elif isinstance(value, float):
-            columns.append(f'"{key}" NUMERIC')
-        elif isinstance(value, bool):
-            columns.append(f'"{key}" BOOLEAN')
-        else:
-            # Default to TEXT for strings and other types
-            columns.append(f'"{key}" TEXT')
-    
-    # Add metadata columns
-    columns.extend([
-        '"source" TEXT',
-        '"extraction_date" TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP',
-        '"extraction_backend" TEXT',
-        '"extraction_model" TEXT'
-    ])
-    
-    # Create SQL statement for table creation
-    columns_sql = ", ".join(columns)
-    sql = f'''
-    CREATE TABLE IF NOT EXISTS "{table_name}" (
-        id SERIAL PRIMARY KEY,
-        {columns_sql}
-    )
-    '''
-    
-    # Execute SQL to create table
-    result = await self.execute_sql_async(sql)
-    log.info(f"Created or ensured table {table_name} exists")
-    
-    # Grant permissions if users are provided
-    if users:
-        for user in users:
-            await self.execute_sql_async(f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE "{table_name}" TO "{user}";')
-    
-    return result
+        Returns:
+            bool: True if connection is valid, False otherwise
+        """
+        try:
+            # Simple query to check connection
+            result = await self.execute_sql_async("SELECT 1")
+            return True
+        except Exception as e:
+            log.warning(f"Database connection check failed: {e}")
+            return False
 
-async def write_data_to_table(self, table_name: str, data: dict, metadata: dict = None):
-    """
-    Writes data to the specified table.
-    
-    Args:
-        table_name (str): Name of the table
-        data (dict): Data to write to the table
-        metadata (dict, optional): Additional metadata to include
+    async def ensure_connected(self):
+        """
+        Ensures the database connection is valid, attempting to reconnect if necessary.
         
-    Returns:
-        Result of SQL execution
-    """
-    # Create copies to avoid modifying the original data
-    insert_data = dict(data)
-    
-    # Add metadata if provided
-    if metadata:
-        insert_data["source"] = metadata.get("objectId", metadata.get("source", "unknown"))
-        insert_data["extraction_backend"] = metadata.get("extraction_backend", "unknown")
-        insert_data["extraction_model"] = metadata.get("extraction_model", "unknown")
-    
-    # Prepare column names and placeholders for values
-    columns = [f'"{key}"' for key in insert_data.keys()]
-    placeholders = []
-    values = []
-    
-    # Process values and create properly formatted placeholders
-    for key, value in insert_data.items():
-        values.append(json.dumps(value) if isinstance(value, (dict, list)) else value)
-        placeholders.append("%s")
-    
-    # Create SQL statement for insertion
-    columns_str = ", ".join(columns)
-    placeholders_str = ", ".join(placeholders)
-    
-    sql = f'''
-    INSERT INTO "{table_name}" ({columns_str})
-    VALUES ({placeholders_str})
-    RETURNING id
-    '''
-    
-    # Execute SQL to insert data
-    result = await self.execute_sql_async(sql, values)
-    log.info(f"Inserted data into table {table_name}")
-    
-    return result
+        Returns:
+            bool: True if connection is valid or reconnection successful, False otherwise
+        """
+        if await self.check_connection():
+            return True
+            
+        try:
+            # Attempt to reconnect - implementation depends on your database driver
+            if self.engine_type == "pg8000":
+                # Re-create the engine
+                self.engine = self._create_engine_from_pg8000(self.user, self.password, self.database)
+            elif self.engine_type == "langchain":
+                # Re-create the engine
+                self.engine = self._create_engine()
+                
+            log.info(f"Successfully reconnected to AlloyDB")
+            return True
+        except Exception as e:
+            log.error(f"Failed to reconnect to AlloyDB: {e}")
+            return False
 
-async def execute_sql_async(self, sql_statement, values=None):
-    """
-    Executes a given SQL statement asynchronously with optional parameter values.
-    
-    Args:
-        sql_statement (str): The SQL statement to execute
-        values (list, optional): Values for parameterized query
+    async def close(self):
+        """
+        Properly close the database connection.
+        """
+        try:
+            if self.engine_type == "pg8000":
+                # Close engine or connector
+                if hasattr(self, 'connector'):
+                    await self.connector.close()
+            # For langchain engine, additional cleanup might be needed
+            log.info("Closed AlloyDB connection")
+        except Exception as e:
+            log.warning(f"Error closing AlloyDB connection: {e}")
+
+    async def create_table_from_schema(self, table_name: str, schema_data: dict, users: list = None):
+        """
+        Creates or ensures a table exists based on the structure of the provided schema data.
         
-    Returns:
-        Result of SQL execution
-    """
-    log.info(f"Executing async SQL statement: {sql_statement}")
-    if self.engine_type == "pg8000":
-        result = await self._execute_sql_async_pg8000(sql_statement, values)
-    elif self.engine_type == "langchain":
-        result = await self._execute_sql_async_langchain(sql_statement, values)
-    
-    return result
+        Args:
+            table_name (str): Name of the table to create
+            schema_data (dict): Data structure that matches the expected schema
+            users (list, optional): List of users to grant permissions to
+            
+        Returns:
+            Result of SQL execution
+        """
+        # Generate column definitions from schema data
+        columns = []
+        for key, value in schema_data.items():
+            if isinstance(value, dict):
+                # For nested objects, store as JSONB
+                columns.append(f'"{key}" JSONB')
+            elif isinstance(value, list):
+                # For arrays, store as JSONB
+                columns.append(f'"{key}" JSONB')
+            elif isinstance(value, int):
+                columns.append(f'"{key}" INTEGER')
+            elif isinstance(value, float):
+                columns.append(f'"{key}" NUMERIC')
+            elif isinstance(value, bool):
+                columns.append(f'"{key}" BOOLEAN')
+            else:
+                # Default to TEXT for strings and other types
+                columns.append(f'"{key}" TEXT')
+        
+        # Add metadata columns
+        columns.extend([
+            '"source" TEXT',
+            '"extraction_date" TIMESTAMPTZ DEFAULT CURRENT_TIMESTAMP',
+            '"extraction_backend" TEXT',
+            '"extraction_model" TEXT'
+        ])
+        
+        # Create SQL statement for table creation
+        columns_sql = ", ".join(columns)
+        sql = f'''
+        CREATE TABLE IF NOT EXISTS "{table_name}" (
+            id SERIAL PRIMARY KEY,
+            {columns_sql}
+        )
+        '''
+        
+        # Execute SQL to create table
+        result = await self.execute_sql_async(sql)
+        log.info(f"Created or ensured table {table_name} exists")
+        
+        # Grant permissions if users are provided
+        if users:
+            for user in users:
+                await self.execute_sql_async(f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLE "{table_name}" TO "{user}";')
+        
+        return result
+
+    async def write_data_to_table(self, table_name: str, data: dict, metadata: dict = None):
+        """
+        Writes data to the specified table.
+        
+        Args:
+            table_name (str): Name of the table
+            data (dict): Data to write to the table
+            metadata (dict, optional): Additional metadata to include
+            
+        Returns:
+            Result of SQL execution
+        """
+        # Create copies to avoid modifying the original data
+        insert_data = dict(data)
+        
+        # Add metadata if provided
+        if metadata:
+            insert_data["source"] = metadata.get("objectId", metadata.get("source", "unknown"))
+            insert_data["extraction_backend"] = metadata.get("extraction_backend", "unknown")
+            insert_data["extraction_model"] = metadata.get("extraction_model", "unknown")
+        
+        # Prepare column names and placeholders for values
+        columns = [f'"{key}"' for key in insert_data.keys()]
+        placeholders = []
+        values = []
+        
+        # Process values and create properly formatted placeholders
+        for key, value in insert_data.items():
+            values.append(json.dumps(value) if isinstance(value, (dict, list)) else value)
+            placeholders.append("%s")
+        
+        # Create SQL statement for insertion
+        columns_str = ", ".join(columns)
+        placeholders_str = ", ".join(placeholders)
+        
+        sql = f'''
+        INSERT INTO "{table_name}" ({columns_str})
+        VALUES ({placeholders_str})
+        RETURNING id
+        '''
+        
+        # Execute SQL to insert data
+        result = await self.execute_sql_async(sql, values)
+        log.info(f"Inserted data into table {table_name}")
+        
+        return result
+
+    async def execute_sql_async(self, sql_statement, values=None):
+        """
+        Executes a given SQL statement asynchronously with optional parameter values.
+        
+        Args:
+            sql_statement (str): The SQL statement to execute
+            values (list, optional): Values for parameterized query
+            
+        Returns:
+            Result of SQL execution
+        """
+        log.info(f"Executing async SQL statement: {sql_statement}")
+        if self.engine_type == "pg8000":
+            result = await self._execute_sql_async_pg8000(sql_statement, values)
+        elif self.engine_type == "langchain":
+            result = await self._execute_sql_async_langchain(sql_statement, values)
+        
+        return result
 

{sunholo-0.123.0.dist-info → sunholo-0.123.2.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sunholo
-Version: 0.123.0
+Version: 0.123.2
 Summary: Large Language Model DevOps - a package to help deploy LLMs to the Cloud.
 Author-email: Holosun ApS <multivac@sunholo.com>
 License: Apache License, Version 2.0

{sunholo-0.123.0.dist-info → sunholo-0.123.2.dist-info}/RECORD

@@ -60,7 +60,7 @@ sunholo/components/retriever.py,sha256=Wmchv3huAM4w7DIS-a5Lp9Hi7M8pE6vZdxgseiT9S
 sunholo/components/vectorstore.py,sha256=k7GS1Y5c6ZGXSDAJvyCes6dTjhDAi0fjGbVLqpyfzBc,5918
 sunholo/database/__init__.py,sha256=bpB5Nk21kwqYj-qdVnvNgXjLsbflnH4g-San7OHMqR4,283
 sunholo/database/alloydb.py,sha256=x1zUMB-EVWbE2Zvp4nAs2Z-tB_kOZmS45H2lwVHdYnk,11678
-sunholo/database/alloydb_client.py,sha256=uUJN-FYPmGAb9rzGVdg7ds38YhwKBQCXPqUtdOzBQCg,25123
+sunholo/database/alloydb_client.py,sha256=uPRcdfrE1YXCUoDbaKopjlKXXP5deceEn08bJKUeWI4,27596
 sunholo/database/database.py,sha256=VqhZdkXUNdvWn8sUcUV3YNby1JDVf7IykPVXWBtxo9U,7361
 sunholo/database/lancedb.py,sha256=DyfZntiFKBlVPaFooNN1Z6Pl-LAs4nxWKKuq8GBqN58,715
 sunholo/database/static_dbs.py,sha256=8cvcMwUK6c32AS2e_WguKXWMkFf5iN3g9WHzsh0C07Q,442
@@ -168,9 +168,9 @@ sunholo/vertex/init.py,sha256=1OQwcPBKZYBTDPdyU7IM4X4OmiXLdsNV30C-fee2scQ,2875
 sunholo/vertex/memory_tools.py,sha256=tBZxqVZ4InTmdBvLlOYwoSEWu4-kGquc-gxDwZCC4FA,7667
 sunholo/vertex/safety.py,sha256=S9PgQT1O_BQAkcqauWncRJaydiP8Q_Jzmu9gxYfy1VA,2482
 sunholo/vertex/type_dict_to_json.py,sha256=uTzL4o9tJRao4u-gJOFcACgWGkBOtqACmb6ihvCErL8,4694
-sunholo-0.123.0.dist-info/licenses/LICENSE.txt,sha256=SdE3QjnD3GEmqqg9EX3TM9f7WmtOzqS1KJve8rhbYmU,11345
-sunholo-0.123.0.dist-info/METADATA,sha256=ue-wXoQBj7xkEm3nzOg0to1O3counnyi7K3RJfzcv0M,10001
-sunholo-0.123.0.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
-sunholo-0.123.0.dist-info/entry_points.txt,sha256=bZuN5AIHingMPt4Ro1b_T-FnQvZ3teBes-3OyO0asl4,49
-sunholo-0.123.0.dist-info/top_level.txt,sha256=wt5tadn5--5JrZsjJz2LceoUvcrIvxjHJe-RxuudxAk,8
-sunholo-0.123.0.dist-info/RECORD,,
+sunholo-0.123.2.dist-info/licenses/LICENSE.txt,sha256=SdE3QjnD3GEmqqg9EX3TM9f7WmtOzqS1KJve8rhbYmU,11345
+sunholo-0.123.2.dist-info/METADATA,sha256=0w8jC9b0-_LEB8VnRcKIIbSlHCIJRFPyiuqwfTnbsm8,10001
+sunholo-0.123.2.dist-info/WHEEL,sha256=CmyFI0kx5cdEMTLiONQRbGQwjIoR1aIYB7eCAQ4KPJ0,91
+sunholo-0.123.2.dist-info/entry_points.txt,sha256=bZuN5AIHingMPt4Ro1b_T-FnQvZ3teBes-3OyO0asl4,49
+sunholo-0.123.2.dist-info/top_level.txt,sha256=wt5tadn5--5JrZsjJz2LceoUvcrIvxjHJe-RxuudxAk,8
+sunholo-0.123.2.dist-info/RECORD,,