# FlightMonitor/data/storage.py import sqlite3 import os from datetime import datetime, timezone import time from typing import Optional, List, Dict # Relative imports # MODIFIED: Import config as app_config # WHY: Use the application-wide configuration module. # HOW: Added import. from . import config as app_config # MODIFIED: Import get_logger from the centralized logger utility. # WHY: Use the application's standard logger. # HOW: Added import. from ..utils.logger import get_logger # MODIFIED: Import CanonicalFlightState model. # WHY: Storage layer works with this standard format. # HOW: Added import. from .common_models import CanonicalFlightState # Import our canonical model # MODIFIED: Get module-level logger using the centralized get_logger. # WHY: Adhere to the application's logging standard. # HOW: Changed logging.getLogger(__name__) to get_logger(__name__). module_logger = get_logger(__name__) class DataStorage: """ Handles storage of flight data into daily SQLite database files. """ def __init__(self): self._current_db_path: Optional[str] = None self._connection: Optional[sqlite3.Connection] = None self._ensure_database_directory() module_logger.info("DataStorage initialized.") def _ensure_database_directory(self): # ... (implementation as before) ... try: abs_db_dir = os.path.abspath(app_config.DATABASE_DIRECTORY) if not os.path.exists(abs_db_dir): os.makedirs(abs_db_dir) module_logger.info(f"Database directory created: {abs_db_dir}") if not os.path.isdir(abs_db_dir): raise OSError(f"Path {abs_db_dir} exists but is not a directory.") if not os.access(abs_db_dir, os.W_OK): raise OSError(f"Directory {abs_db_dir} is not writable.") except OSError as e: module_logger.error( f"Database directory error for {app_config.DATABASE_DIRECTORY}: {e}", exc_info=True, ) raise def _get_db_path_for_date(self, target_date_utc: Optional[datetime] = None) -> str: # ... (implementation as before) ... if target_date_utc is None: target_date_utc = datetime.now(timezone.utc) db_filename = target_date_utc.strftime(app_config.DATABASE_FILENAME_FORMAT) db_path = os.path.join( os.path.abspath(app_config.DATABASE_DIRECTORY), db_filename ) return db_path def _get_db_connection( self, for_date_utc: Optional[datetime] = None ) -> Optional[sqlite3.Connection]: # ... (implementation as before, ensuring row_factory is set) ... target_db_path = self._get_db_path_for_date(for_date_utc) if self._connection and self._current_db_path == target_db_path: return self._connection if self._connection: try: self._connection.close() module_logger.info(f"Closed DB connection to: {self._current_db_path}") except sqlite3.Error as e: module_logger.error( f"Error closing previous DB connection to {self._current_db_path}: {e}" ) self._connection = None self._current_db_path = target_db_path try: self._connection = sqlite3.connect(target_db_path, timeout=10.0) self._connection.execute("PRAGMA foreign_keys = ON;") self._connection.row_factory = ( sqlite3.Row ) # Ensures select queries return dict-like rows module_logger.info(f"Opened DB connection to: {target_db_path}") self._init_tables_if_not_exist(self._connection) return self._connection except sqlite3.Error as e: module_logger.error( f"Failed to connect/initialize database {target_db_path}: {e}", exc_info=True, ) if self._connection: try: self._connection.close() except sqlite3.Error: pass self._connection = None self._current_db_path = None return None def _init_tables_if_not_exist(self, db_conn: sqlite3.Connection): # ... (implementation as before) ... # Schema includes: # flights (flight_id, icao24, callsign, origin_country, first_seen_day, last_seen_day) # positions (position_id, flight_id, detection_timestamp, latitude, longitude, # baro_altitude, geo_altitude, velocity, true_track, on_ground, recorded_at) cursor = None try: cursor = db_conn.cursor() cursor.execute( """ CREATE TABLE IF NOT EXISTS flights ( flight_id INTEGER PRIMARY KEY AUTOINCREMENT, icao24 TEXT NOT NULL UNIQUE, callsign TEXT, origin_country TEXT, first_seen_day REAL NOT NULL, last_seen_day REAL NOT NULL ) """ ) cursor.execute( """ CREATE TABLE IF NOT EXISTS positions ( position_id INTEGER PRIMARY KEY AUTOINCREMENT, flight_id INTEGER NOT NULL, detection_timestamp REAL NOT NULL, latitude REAL NOT NULL, longitude REAL NOT NULL, baro_altitude_m REAL, /* MODIFIED: Suffix _m for clarity */ geo_altitude_m REAL, /* MODIFIED: Suffix _m for clarity */ velocity_mps REAL, /* MODIFIED: Suffix _mps for clarity */ true_track_deg REAL, /* MODIFIED: Suffix _deg for clarity */ vertical_rate_mps REAL, /* ADDED: New field from CanonicalFlightState */ on_ground BOOLEAN NOT NULL, squawk TEXT, /* ADDED: New field */ spi BOOLEAN, /* ADDED: New field */ position_source TEXT, /* ADDED: New field (TEXT to be flexible) */ raw_data_provider TEXT, /* ADDED: New field */ recorded_at REAL NOT NULL, FOREIGN KEY (flight_id) REFERENCES flights(flight_id) ON DELETE CASCADE ) """ ) # Update indexes if new columns are frequently queried cursor.execute( "CREATE INDEX IF NOT EXISTS idx_positions_flight_id_timestamp ON positions (flight_id, detection_timestamp);" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_flights_icao24 ON flights (icao24);" ) cursor.execute( "CREATE INDEX IF NOT EXISTS idx_positions_detection_timestamp ON positions (detection_timestamp);" ) db_conn.commit() module_logger.debug( f"Database tables and indexes ensured in {self._current_db_path}." ) except sqlite3.Error as e: module_logger.error( f"Error initializing tables in {self._current_db_path}: {e}", exc_info=True, ) if db_conn: db_conn.rollback() finally: if cursor: cursor.close() def add_or_update_flight_daily( self, icao24: str, callsign: Optional[str], origin_country: Optional[str], detection_timestamp: float, ) -> Optional[int]: """ Adds a new flight to the daily DB or updates its last_seen_day. Operates on the DB file for the UTC date of detection_timestamp. Args: icao24 (str): ICAO24 address. callsign (Optional[str]): Callsign. origin_country (Optional[str]): Origin country. detection_timestamp (float): Unix timestamp (UTC) of flight data detection. Returns: Optional[int]: The flight_id, or None on error. """ if not icao24: # Should be caught by CanonicalFlightState constructor too module_logger.warning("icao24 is required to add or update a flight.") return None event_date_utc = datetime.fromtimestamp(detection_timestamp, timezone.utc) conn = self._get_db_connection(for_date_utc=event_date_utc) if not conn: module_logger.error( "Failed to get DB connection for add_or_update_flight_daily." ) return None cursor = None flight_id: Optional[int] = None try: cursor = conn.cursor() cursor.execute( "SELECT flight_id, callsign FROM flights WHERE icao24 = ?", (icao24,) ) row = cursor.fetchone() if row: flight_id = row["flight_id"] existing_callsign = row["callsign"] # Update callsign only if the new one is provided and different, # or if the existing one was None/empty and the new one is provided. final_callsign = existing_callsign if callsign and callsign != existing_callsign: final_callsign = callsign elif ( not existing_callsign and callsign ): # Existing was empty, new one is not final_callsign = callsign cursor.execute( """ UPDATE flights SET last_seen_day = ?, callsign = ? WHERE flight_id = ? """, (detection_timestamp, final_callsign, flight_id), ) module_logger.debug( f"Updated flight icao24={icao24}, flight_id={flight_id}, last_seen_day={detection_timestamp:.0f} in {self._current_db_path}" ) else: cursor.execute( """ INSERT INTO flights (icao24, callsign, origin_country, first_seen_day, last_seen_day) VALUES (?, ?, ?, ?, ?) """, ( icao24, callsign, origin_country, detection_timestamp, detection_timestamp, ), ) flight_id = cursor.lastrowid module_logger.info( f"Added new flight icao24={icao24}, flight_id={flight_id} to {self._current_db_path}" ) conn.commit() except sqlite3.Error as e: module_logger.error( f"DB error in add_or_update_flight_daily for icao24={icao24}: {e}", exc_info=True, ) if conn: conn.rollback() flight_id = None finally: if cursor: cursor.close() return flight_id def add_position_daily( self, flight_id: int, flight_state_obj: CanonicalFlightState ) -> Optional[int]: # MODIFIED: Accepts CanonicalFlightState """ Adds a new position record using data from a CanonicalFlightState object. Operates on the DB file for the UTC date of flight_state_obj.timestamp. Args: flight_id (int): The flight_id from the 'flights' table. flight_state_obj (CanonicalFlightState): The canonical flight state object. Returns: Optional[int]: The position_id of the inserted record, or None on error. """ if not flight_id: module_logger.warning("flight_id is required to add a position.") return None if not isinstance(flight_state_obj, CanonicalFlightState): module_logger.error( f"Invalid type for flight_state_obj: expected CanonicalFlightState, got {type(flight_state_obj)}" ) return None # Use the timestamp from the flight state object to determine the correct daily DB event_date_utc = datetime.fromtimestamp( flight_state_obj.timestamp, timezone.utc ) conn = self._get_db_connection(for_date_utc=event_date_utc) if not conn: module_logger.error( f"Failed to get DB connection for add_position_daily (flight_id={flight_id})." ) return None cursor = None position_id: Optional[int] = None recorded_at_ts = time.time() # Timestamp of actual DB record insertion # Validate essential position data if flight_state_obj.latitude is None or flight_state_obj.longitude is None: module_logger.warning( f"Skipping position for flight_id={flight_id} (ICAO: {flight_state_obj.icao24}) " f"due to missing latitude/longitude. Timestamp: {flight_state_obj.timestamp:.0f}" ) return None try: cursor = conn.cursor() # MODIFIED: SQL query to include new fields from CanonicalFlightState sql = """ INSERT INTO positions ( flight_id, detection_timestamp, latitude, longitude, baro_altitude_m, geo_altitude_m, velocity_mps, true_track_deg, vertical_rate_mps, on_ground, squawk, spi, position_source, raw_data_provider, recorded_at ) VALUES (?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?, ?) """ params = ( flight_id, flight_state_obj.timestamp, # This is the primary timestamp from the canonical object flight_state_obj.latitude, flight_state_obj.longitude, flight_state_obj.baro_altitude_m, flight_state_obj.geo_altitude_m, flight_state_obj.velocity_mps, flight_state_obj.true_track_deg, flight_state_obj.vertical_rate_mps, # New field flight_state_obj.on_ground, flight_state_obj.squawk, # New field flight_state_obj.spi, # New field ( str(flight_state_obj.position_source) if flight_state_obj.position_source is not None else None ), # Store as TEXT flight_state_obj.raw_data_provider, # New field recorded_at_ts, ) cursor.execute(sql, params) position_id = cursor.lastrowid conn.commit() module_logger.debug( f"Added position for flight_id={flight_id} (ICAO: {flight_state_obj.icao24}) " f"at {flight_state_obj.timestamp:.0f} (pos_id={position_id}) in {self._current_db_path}" ) except sqlite3.Error as e: module_logger.error( f"DB error in add_position_daily for flight_id={flight_id} (ICAO: {flight_state_obj.icao24}): {e}", exc_info=True, ) if conn: conn.rollback() position_id = None finally: if cursor: cursor.close() return position_id def close_connection(self): # ... (implementation as before) ... if self._connection: try: self._connection.close() module_logger.info( f"DB connection to {self._current_db_path} closed by request." ) except sqlite3.Error as e: module_logger.error( f"Error explicitly closing DB connection to {self._current_db_path}: {e}" ) finally: self._connection = None self._current_db_path = None # Standalone test block if __name__ == "__main__": # ... (Test setup as before, using MockConfig and CanonicalFlightState) ... import logging # Ensure logging is imported for basicConfig from datetime import timedelta # Ensure timedelta is imported logging.basicConfig( level=logging.DEBUG, format="%(asctime)s [%(levelname)-8s] %(name)-20s : %(message)s", datefmt="%Y-%m-%d %H:%M:%S", ) module_logger.info( "--- Running DataStorage Standalone Test (with CanonicalFlightState) ---" ) class MockConfig: DATABASE_DIRECTORY = os.path.join( os.getcwd(), "test_flight_data_storage_canonical" ) DATABASE_FILENAME_FORMAT = "test_flights_canonical_%Y-%m-%d.db" # MODIFIED: Added LOG_LEVEL to MockConfig to avoid error if it's accessed by get_logger in storage.py (though it shouldn't be after the logger refactor) # WHY: Defensive programming, ensures MockConfig is more complete if modules unexpectedly access it. # HOW: Added LOG_LEVEL. LOG_LEVEL = "DEBUG" original_app_config = app_config # MODIFIED: Get a reference to the module-level app_config before overwriting. # WHY: Need to restore it at the end of the test. # HOW: Used globals() to get the module reference. app_config_module_ref = globals()[ "app_config" ] # Get a reference to the module-level app_config globals()[ "app_config" ] = MockConfig() # Temporarily override module-level app_config for test storage = None try: storage = DataStorage() current_ts_utc = time.time() # Flight 1 state1_obj = CanonicalFlightState( icao24="CANON01", callsign="CANONFLT1", origin_country="Canonistan", timestamp=current_ts_utc, last_contact_timestamp=current_ts_utc, latitude=51.0, longitude=11.0, on_ground=False, baro_altitude_m=9000.0, velocity_mps=200.0, true_track_deg=90.0, vertical_rate_mps=10.0, squawk="1234", spi=False, position_source=0, raw_data_provider="TestAdapter", ) flight_id1 = storage.add_or_update_flight_daily( state1_obj.icao24, state1_obj.callsign, state1_obj.origin_country, state1_obj.timestamp, ) if flight_id1: module_logger.info( f"Flight 1 (icao={state1_obj.icao24}) added/updated, flight_id: {flight_id1}" ) pos_id1 = storage.add_position_daily(flight_id1, state1_obj) if pos_id1: module_logger.info( f"Position 1 for flight_id {flight_id1} added, pos_id: {pos_id1}" ) # Simulate update for Flight 1 time.sleep(0.01) # Ensure different timestamp updated_ts_utc = time.time() state1_update_obj = CanonicalFlightState( icao24="CANON01", callsign="CANONFLT1A", origin_country="Canonistan", # Callsign changed timestamp=updated_ts_utc, last_contact_timestamp=updated_ts_utc, # New timestamp latitude=51.5, longitude=11.5, on_ground=False, baro_altitude_m=9050.0, velocity_mps=205.0, true_track_deg=95.0, vertical_rate_mps=0.0, # Leveling off squawk="1234", spi=False, position_source=0, raw_data_provider="TestAdapter", ) flight_id1_upd = storage.add_or_update_flight_daily( state1_update_obj.icao24, state1_update_obj.callsign, state1_update_obj.origin_country, state1_update_obj.timestamp, ) assert ( flight_id1_upd == flight_id1 ), "Flight ID should remain the same on update" module_logger.info( f"Flight 1 (icao={state1_update_obj.icao24}) updated, flight_id: {flight_id1_upd}" ) if flight_id1_upd: pos_id1_upd = storage.add_position_daily(flight_id1_upd, state1_update_obj) if pos_id1_upd: module_logger.info( f"Position 2 for flight_id {flight_id1_upd} added, pos_id: {pos_id1_upd}" ) # Flight 2 - for "tomorrow" tomorrow_dt = datetime.now(timezone.utc) + timedelta(days=1) tomorrow_ts_utc = tomorrow_dt.timestamp() state2_obj = CanonicalFlightState( icao24="CANON02", callsign="CANONFLT2", origin_country="OtherCanonPlace", timestamp=tomorrow_ts_utc, last_contact_timestamp=tomorrow_ts_utc, latitude=34.0, longitude=-118.0, on_ground=False, baro_altitude_m=12000.0, velocity_mps=250.0, true_track_deg=270.0, vertical_rate_mps=-5.0, # Descending squawk="7700", spi=True, position_source=1, raw_data_provider="TestAdapter", # Emergency squawk, different source ) flight_id2 = storage.add_or_update_flight_daily( state2_obj.icao24, state2_obj.callsign, state2_obj.origin_country, state2_obj.timestamp, ) if flight_id2: module_logger.info( f"Flight 2 (icao={state2_obj.icao24}) for tomorrow added, flight_id: {flight_id2}" ) pos_id2 = storage.add_position_daily(flight_id2, state2_obj) if pos_id2: module_logger.info( f"Position for flight_id {flight_id2} (tomorrow) added, pos_id: {pos_id2}" ) except Exception as e: module_logger.critical( f"Error during DataStorage standalone test (Canonical): {e}", exc_info=True ) finally: if storage: storage.close_connection() # MODIFIED: Restore the original app_config after the test. # WHY: Avoid affecting other parts of the application if this script is run as part of a larger test suite. # HOW: Assigned the saved reference back to the module's app_config using globals(). globals()["app_config"] = original_app_config # Restore original app_config test_dir_to_check = os.path.join( os.getcwd(), "test_flight_data_storage_canonical" ) if os.path.exists(test_dir_to_check): module_logger.info( f"Test directory {test_dir_to_check} was used. Please check/clean manually." ) # import shutil # shutil.rmtree(test_dir_to_check) # module_logger.info(f"Test directory {test_dir_to_check} cleaned up.") module_logger.info("--- DataStorage Standalone Test (Canonical) Finished ---")