Done brand and body type detection with postgresql integration

siwatsystem/database.py

@@ -2,6 +2,7 @@ import psycopg2
 import psycopg2.extras
 from typing import Optional, Dict, Any
 import logging
+import uuid
 
 logger = logging.getLogger(__name__)
 
@@ -90,8 +91,11 @@ class DatabaseManager:
                     set_clauses.append(f"{field} = %s")
                     values.append(value)
             
+            # Add schema prefix if table doesn't already have it
+            full_table_name = table if '.' in table else f"gas_station_1.{table}"
+            
             query = f"""
-            INSERT INTO {table} ({key_field}, {', '.join(fields.keys())})
+            INSERT INTO {full_table_name} ({key_field}, {', '.join(fields.keys())})
             VALUES (%s, {', '.join(['%s'] * len(fields))})
             ON CONFLICT ({key_field})
             DO UPDATE SET {', '.join(set_clauses)}
@@ -109,4 +113,80 @@ class DatabaseManager:
             logger.error(f"Failed to execute update on {table}: {e}")
             if self.connection:
                 self.connection.rollback()
-            return False
+            return False
+    
+    def create_car_frontal_info_table(self) -> bool:
+        """Create the car_frontal_info table in gas_station_1 schema if it doesn't exist."""
+        if not self.is_connected():
+            if not self.connect():
+                return False
+        
+        try:
+            cur = self.connection.cursor()
+            
+            # Create schema if it doesn't exist
+            cur.execute("CREATE SCHEMA IF NOT EXISTS gas_station_1")
+            
+            # Create table if it doesn't exist
+            create_table_query = """
+            CREATE TABLE IF NOT EXISTS gas_station_1.car_frontal_info (
+                camera_id VARCHAR(255),
+                captured_timestamp VARCHAR(255),
+                session_id VARCHAR(255) PRIMARY KEY,
+                license_character VARCHAR(255) DEFAULT NULL,
+                license_type VARCHAR(255) DEFAULT 'No model available',
+                car_brand VARCHAR(255) DEFAULT NULL,
+                car_model VARCHAR(255) DEFAULT NULL,
+                car_body_type VARCHAR(255) DEFAULT NULL,
+                created_at TIMESTAMP DEFAULT NOW(),
+                updated_at TIMESTAMP DEFAULT NOW()
+            )
+            """
+            
+            cur.execute(create_table_query)
+            self.connection.commit()
+            cur.close()
+            logger.info("Successfully created/verified car_frontal_info table in gas_station_1 schema")
+            return True
+            
+        except Exception as e:
+            logger.error(f"Failed to create car_frontal_info table: {e}")
+            if self.connection:
+                self.connection.rollback()
+            return False
+    
+    def insert_initial_detection(self, camera_id: str, captured_timestamp: str, session_id: str = None) -> str:
+        """Insert initial detection record and return the session_id."""
+        if not self.is_connected():
+            if not self.connect():
+                return None
+        
+        # Generate session_id if not provided
+        if not session_id:
+            session_id = str(uuid.uuid4())
+        
+        try:
+            # Ensure table exists
+            if not self.create_car_frontal_info_table():
+                logger.error("Failed to create/verify table before insertion")
+                return None
+            
+            cur = self.connection.cursor()
+            insert_query = """
+            INSERT INTO gas_station_1.car_frontal_info 
+            (camera_id, captured_timestamp, session_id, license_character, license_type, car_brand, car_model, car_body_type)
+            VALUES (%s, %s, %s, NULL, 'No model available', NULL, NULL, NULL)
+            ON CONFLICT (session_id) DO NOTHING
+            """
+            
+            cur.execute(insert_query, (camera_id, captured_timestamp, session_id))
+            self.connection.commit()
+            cur.close()
+            logger.info(f"Inserted initial detection record with session_id: {session_id}")
+            return session_id
+            
+        except Exception as e:
+            logger.error(f"Failed to insert initial detection record: {e}")
+            if self.connection:
+                self.connection.rollback()
+            return None