mlukac/xrf-explorer-dev
0
1"""2Database models and operations for XRF Explorer well database.3 4This module provides:5- Well and XRF file data models6- Database connection management 7- Query operations for the UI layer8"""9 10import sqlite311from dataclasses import dataclass12from pathlib import Path13from typing import List, Optional, Tuple14 15from xrf_explorer.core.types import DatasetType16 17 18@dataclass19class Well:20 """Data model for a well record."""21 id: int22 name: str23 api_number: Optional[str] = None24 latitude: Optional[float] = None25 longitude: Optional[float] = None26 county: Optional[str] = None27 state: Optional[str] = None28 29 30@dataclass 31class WellFile:32 """Data model for a well file record (XRF, mass spec, drilling, gamma ray, etc.)."""33 id: int34 well_id: int35 file_path: str # Relative to data/processed/36 dataset_type: str # DatasetType enum value (xrf, mass_spec, drilling, gamma_ray)37 file_type: str38 service_company: Optional[str] = None39 start_depth: Optional[float] = None40 stop_depth: Optional[float] = None41 42 43@dataclass44class XRFFile(WellFile):45 """Backward compatibility alias for XRF files."""46 pass47 48 49class WellDatabase:50 """Database operations for well and XRF file management."""51 52 def __init__(self, db_path: str = "data/wells.db"):53 """54 Initialize database connection.55 56 Args:57 db_path: Path to SQLite database file58 """59 self.db_path = db_path60 self._check_database_exists()61 62 def _check_database_exists(self) -> None:63 """Check if database file exists and raise helpful error if not."""64 if not Path(self.db_path).exists():65 raise FileNotFoundError(66 f"Database not found at {self.db_path}. "67 f"Run 'python database/seed_data.py' to create it."68 )69 70 def _get_connection(self) -> sqlite3.Connection:71 """Get database connection with row factory for easier access."""72 conn = sqlite3.connect(self.db_path)73 conn.row_factory = sqlite3.Row # Enables dict-like access to rows74 return conn75 76 def get_wells(self) -> List[Well]:77 """Get all wells from database."""78 conn = self._get_connection()79 cursor = conn.cursor()80 81 cursor.execute("""82 SELECT id, name, api_number, latitude, longitude, county, state83 FROM wells84 ORDER BY name85 """)86 87 wells = []88 for row in cursor.fetchall():89 wells.append(Well(90 id=row["id"],91 name=row["name"], 92 api_number=row["api_number"],93 latitude=row["latitude"],94 longitude=row["longitude"],95 county=row["county"],96 state=row["state"]97 ))98 99 conn.close()100 return wells101 102 def get_well_by_id(self, well_id: int) -> Optional[Well]:103 """Get a specific well by ID."""104 conn = self._get_connection()105 cursor = conn.cursor()106 107 cursor.execute("""108 SELECT id, name, api_number, latitude, longitude, county, state109 FROM wells110 WHERE id = ?111 """, (well_id,))112 113 row = cursor.fetchone()114 conn.close()115 116 if row:117 return Well(118 id=row["id"],119 name=row["name"],120 api_number=row["api_number"], 121 latitude=row["latitude"],122 longitude=row["longitude"],123 county=row["county"],124 state=row["state"]125 )126 return None127 128 def get_well_files_for_well(self, well_id: int, dataset_type: Optional[DatasetType] = None) -> List[WellFile]:129 """Get all files for a specific well, optionally filtered by dataset type."""130 conn = self._get_connection()131 cursor = conn.cursor()132 133 if dataset_type:134 cursor.execute("""135 SELECT id, well_id, file_path, dataset_type, file_type, service_company, start_depth, stop_depth136 FROM well_files 137 WHERE well_id = ? AND dataset_type = ?138 ORDER BY start_depth139 """, (well_id, dataset_type.value))140 else:141 cursor.execute("""142 SELECT id, well_id, file_path, dataset_type, file_type, service_company, start_depth, stop_depth143 FROM well_files 144 WHERE well_id = ?145 ORDER BY dataset_type, start_depth146 """, (well_id,))147 148 well_files = []149 for row in cursor.fetchall():150 well_files.append(WellFile(151 id=row["id"],152 well_id=row["well_id"],153 file_path=row["file_path"],154 dataset_type=row["dataset_type"],155 file_type=row["file_type"],156 service_company=row["service_company"],157 start_depth=row["start_depth"],158 stop_depth=row["stop_depth"]159 ))160 161 conn.close()162 return well_files163 164 def get_xrf_files_for_well(self, well_id: int) -> List[XRFFile]:165 """Get all XRF files for a specific well (backward compatibility)."""166 well_files = self.get_well_files_for_well(well_id, DatasetType.XRF)167 return [XRFFile(**file.__dict__) for file in well_files]168 169 def get_well_file_path(self, well_id: int, dataset_type: DatasetType) -> Optional[str]:170 """171 Get the file path for the first file of a specific dataset type for a well.172 173 Args:174 well_id: ID of the well175 dataset_type: Type of dataset (XRF, MASS_SPEC, etc.)176 177 Returns:178 Absolute path to file, or None if no file found179 """180 well_files = self.get_well_files_for_well(well_id, dataset_type)181 182 if not well_files:183 return None184 185 # Get first file of this type (assume one file per dataset type per well)186 well_file = well_files[0]187 188 # Convert relative path to absolute189 relative_path = Path("data/processed") / well_file.file_path190 absolute_path = Path.cwd() / relative_path191 192 if not absolute_path.exists():193 raise FileNotFoundError(f"{dataset_type.value} file not found: {absolute_path}")194 195 return str(absolute_path)196 197 def get_xrf_file_path(self, well_id: int) -> Optional[str]:198 """199 Get the file path for the first XRF file of a well (backward compatibility).200 201 Args:202 well_id: ID of the well203 204 Returns:205 Absolute path to XRF file, or None if no file found206 """207 return self.get_well_file_path(well_id, DatasetType.XRF)208 209 def get_available_dataset_types(self, well_id: int) -> List[DatasetType]:210 """211 Get list of available dataset types for a well.212 213 Args:214 well_id: ID of the well215 216 Returns:217 List of DatasetType enums for files available for this well218 """219 conn = self._get_connection()220 cursor = conn.cursor()221 222 cursor.execute("""223 SELECT DISTINCT dataset_type224 FROM well_files 225 WHERE well_id = ?226 ORDER BY dataset_type227 """, (well_id,))228 229 dataset_types = []230 for row in cursor.fetchall():231 try:232 dataset_types.append(DatasetType(row["dataset_type"]))233 except ValueError:234 # Skip unknown dataset types235 continue236 237 conn.close()238 return dataset_types239 240 def get_wells_with_xrf_files(self) -> List[Tuple[Well, List[XRFFile]]]:241 """Get all wells along with their XRF files."""242 wells = self.get_wells()243 results = []244 245 for well in wells:246 xrf_files = self.get_xrf_files_for_well(well.id)247 results.append((well, xrf_files))248 249 return results250 251 def get_well_choices_for_ui(self) -> dict[str, int]:252 """253 Get well choices formatted for UI dropdown.254 255 Returns:256 Dictionary mapping display_name to well_id257 """258 wells = self.get_wells()259 choices = {}260 261 for well in wells:262 # Create descriptive display name263 display_name = well.name264 if well.api_number:265 display_name += f" (API: {well.api_number})"266 267 choices[display_name] = well.id268 269 return choices270 271 def get_well_locations(self) -> List[dict]:272 """273 Get well location data for mapping.274 275 Returns:276 List of dictionaries with well location and metadata for mapping277 """278 wells = self.get_wells()279 locations = []280 281 for well in wells:282 # Only include wells with valid coordinates283 if well.latitude is not None and well.longitude is not None:284 locations.append({285 'id': well.id,286 'name': well.name,287 'api_number': well.api_number or 'N/A',288 'latitude': well.latitude,289 'longitude': well.longitude,290 'county': well.county or 'Unknown',291 'state': well.state or 'Unknown',292 # Short name for map labels293 'short_name': well.name.split()[0] if well.name else f"Well {well.id}"294 })295 296 return locations297 298 299# Convenience function for easy imports300def get_database() -> WellDatabase:301 """Get a WellDatabase instance with default settings."""302 return WellDatabase()