mlukac/xrf-explorer-dev
0
1#!/usr/bin/env python32"""3Database seeding script for XRF Explorer well database.4 5This script:61. Creates the data/ directory structure72. Copies test files to data/processed/83. Creates SQLite database with schema94. Populates database with well and file data10 11Run with: python xrf_explorer/database/seed_data.py12"""13 14import os15import shutil16import sqlite317from pathlib import Path18 19 20def create_directory_structure():21 """Create the data directory structure."""22 print("📁 Creating directory structure...")23 24 directories = [25 "data",26 "data/processed", 27 "data/processed/well_001_ella",28 "data/processed/well_002_lamp",29 "data/processed/well_003_broken"30 ]31 32 for directory in directories:33 Path(directory).mkdir(parents=True, exist_ok=True)34 print(f" ✓ Created {directory}/")35 36 37def copy_test_files():38 """Copy test files from tests/data/ to data/processed/."""39 print("\n📄 Copying test files...")40 41 file_mappings = [42 # ELLA well files43 {44 "source": "tests/data/ella_xrf.xlsx",45 "dest": "data/processed/well_001_ella/xrf_data.xlsx",46 "description": "ELLA well XRF data (Bruker format)"47 },48 {49 "source": "tests/data/ella_mass_spec.las",50 "dest": "data/processed/well_001_ella/mass_spec_data.las",51 "description": "ELLA well Mass Spec data"52 },53 {54 "source": "tests/data/ella_drilling.las",55 "dest": "data/processed/well_001_ella/drilling_data.las",56 "description": "ELLA well Drilling data"57 },58 {59 "source": "tests/data/ella_gamma.las",60 "dest": "data/processed/well_001_ella/gamma_ray_data.las",61 "description": "ELLA well Gamma Ray data"62 },63 # Lamp well files (using rig1 data)64 {65 "source": "tests/data/rig1_xrf.xlsx",66 "dest": "data/processed/well_002_lamp/xrf_data.xlsx",67 "description": "Lamp well XRF data (using rig1 data)"68 },69 {70 "source": "tests/data/rig1_mass_spec.las",71 "dest": "data/processed/well_002_lamp/mass_spec_data.las",72 "description": "Lamp well Mass Spec data (using rig1 data)"73 },74 {75 "source": "tests/data/rig1_gamma.las",76 "dest": "data/processed/well_002_lamp/gamma_ray_data.las",77 "description": "Lamp well Gamma Ray data (using rig1 data)"78 },79 # Broken well (intentionally missing file for error testing)80 {81 "source": "tests/data/lamp_xrf.csv", 82 "dest": "data/processed/well_003_broken/xrf_data.csv",83 "description": "Broken well XRF data (file doesn't exist - for error testing)"84 }85 ]86 87 for mapping in file_mappings:88 source = Path(mapping["source"])89 dest = Path(mapping["dest"])90 91 if source.exists():92 shutil.copy2(source, dest)93 print(f" ✓ Copied {mapping['description']}")94 print(f" {source} → {dest}")95 else:96 print(f" ⚠️ Source file not found: {source}")97 98 99def create_database_schema():100 """Create SQLite database and tables."""101 print("\n🗄️ Creating database schema...")102 103 db_path = "data/wells.db"104 105 # Remove existing database to start fresh106 if Path(db_path).exists():107 Path(db_path).unlink()108 print(" ✓ Removed existing database")109 110 conn = sqlite3.connect(db_path)111 cursor = conn.cursor()112 113 # Create wells table114 cursor.execute("""115 CREATE TABLE wells (116 id INTEGER PRIMARY KEY,117 name TEXT NOT NULL,118 api_number TEXT,119 latitude REAL,120 longitude REAL,121 county TEXT,122 state TEXT,123 created_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP124 )125 """)126 print(" ✓ Created wells table")127 128 # Create well_files table (renamed from xrf_files to be more general)129 cursor.execute("""130 CREATE TABLE well_files (131 id INTEGER PRIMARY KEY,132 well_id INTEGER,133 file_path TEXT NOT NULL,134 dataset_type TEXT NOT NULL,135 file_type TEXT,136 service_company TEXT,137 start_depth REAL,138 stop_depth REAL,139 created_date TIMESTAMP DEFAULT CURRENT_TIMESTAMP,140 FOREIGN KEY (well_id) REFERENCES wells(id)141 )142 """)143 print(" ✓ Created well_files table")144 145 conn.commit()146 return conn147 148 149def populate_database(conn):150 """Populate database with well and file data."""151 print("\n📝 Populating database...")152 153 cursor = conn.cursor()154 155 # Insert wells data156 wells_data = [157 {158 "name": "ELLA UT 14-7 4-2-7-6-7H",159 "api_number": "43013546560000", 160 "latitude": 41.234,161 "longitude": -109.876,162 "county": "Uintah",163 "state": "Utah"164 },165 {166 "name": "Lamp 75-3229-21L",167 "api_number": "43013545720000",168 "latitude": 41.567, 169 "longitude": -109.234,170 "county": "Uintah",171 "state": "Utah"172 },173 {174 "name": "Broken Test Well",175 "api_number": "99999999999999",176 "latitude": 40.000, 177 "longitude": -110.000,178 "county": "Test",179 "state": "Utah"180 }181 ]182 183 for well in wells_data:184 cursor.execute("""185 INSERT INTO wells (name, api_number, latitude, longitude, county, state)186 VALUES (?, ?, ?, ?, ?, ?)187 """, (well["name"], well["api_number"], well["latitude"], 188 well["longitude"], well["county"], well["state"]))189 190 print(f" ✓ Added well: {well['name']}")191 192 # Insert well files data (multiple file types per well)193 well_files_data = [194 # ELLA well files195 {196 "well_id": 1, # ELLA well197 "file_path": "well_001_ella/xrf_data.xlsx",198 "dataset_type": "xrf",199 "file_type": "BRUKER_XLSX",200 "service_company": "Bruker",201 "start_depth": 10300.0,202 "stop_depth": 14500.0203 },204 {205 "well_id": 1, # ELLA well - Mass Spec206 "file_path": "well_001_ella/mass_spec_data.las",207 "dataset_type": "mass_spec",208 "file_type": "LAS",209 "service_company": "Various",210 "start_depth": 10300.0,211 "stop_depth": 14500.0212 },213 {214 "well_id": 1, # ELLA well - Drilling215 "file_path": "well_001_ella/drilling_data.las",216 "dataset_type": "drilling",217 "file_type": "LAS",218 "service_company": "Drilling Co",219 "start_depth": 10300.0,220 "stop_depth": 14500.0221 },222 {223 "well_id": 1, # ELLA well - Gamma Ray224 "file_path": "well_001_ella/gamma_ray_data.las",225 "dataset_type": "gamma_ray",226 "file_type": "LAS",227 "service_company": "Logging Co",228 "start_depth": 10300.0,229 "stop_depth": 14500.0230 },231 # Lamp well files (using rig1 data)232 {233 "well_id": 2, # Lamp well - XRF234 "file_path": "well_002_lamp/xrf_data.xlsx",235 "dataset_type": "xrf",236 "file_type": "BRUKER_XLSX",237 "service_company": "Bruker",238 "start_depth": 8000.0,239 "stop_depth": 12000.0240 },241 {242 "well_id": 2, # Lamp well - Mass Spec243 "file_path": "well_002_lamp/mass_spec_data.las",244 "dataset_type": "mass_spec",245 "file_type": "LAS",246 "service_company": "Various",247 "start_depth": 8000.0,248 "stop_depth": 12000.0249 },250 {251 "well_id": 2, # Lamp well - Gamma Ray252 "file_path": "well_002_lamp/gamma_ray_data.las",253 "dataset_type": "gamma_ray",254 "file_type": "LAS",255 "service_company": "Logging Co",256 "start_depth": 8000.0,257 "stop_depth": 12000.0258 },259 # Broken well (intentionally missing file for error testing)260 {261 "well_id": 3, # Broken well262 "file_path": "well_003_broken/xrf_data.csv",263 "dataset_type": "xrf",264 "file_type": "CSV",265 "service_company": "Unknown",266 "start_depth": 5000.0,267 "stop_depth": 8000.0268 }269 ]270 271 for well_file in well_files_data:272 cursor.execute("""273 INSERT INTO well_files (well_id, file_path, dataset_type, file_type, service_company, start_depth, stop_depth)274 VALUES (?, ?, ?, ?, ?, ?, ?)275 """, (well_file["well_id"], well_file["file_path"], well_file["dataset_type"], well_file["file_type"],276 well_file["service_company"], well_file["start_depth"], well_file["stop_depth"]))277 278 print(f" ✓ Added {well_file['dataset_type']} file: {well_file['file_path']}")279 280 conn.commit()281 282 283def verify_database():284 """Verify database was created correctly."""285 print("\n✅ Verifying database...")286 287 conn = sqlite3.connect("data/wells.db")288 cursor = conn.cursor()289 290 # Check wells291 cursor.execute("SELECT COUNT(*) FROM wells")292 well_count = cursor.fetchone()[0]293 print(f" ✓ Wells in database: {well_count}")294 295 # Check well files296 cursor.execute("SELECT COUNT(*) FROM well_files") 297 file_count = cursor.fetchone()[0]298 print(f" ✓ Well files in database: {file_count}")299 300 # Show well details301 print("\n📋 Well Summary:")302 cursor.execute("""303 SELECT w.name, wf.file_path, wf.dataset_type, wf.file_type 304 FROM wells w305 LEFT JOIN well_files wf ON w.id = wf.well_id306 ORDER BY w.name, wf.dataset_type307 """)308 309 current_well = None310 for row in cursor.fetchall():311 well_name, file_path, dataset_type, file_type = row312 if well_name != current_well:313 print(f" • {well_name}")314 current_well = well_name315 if file_path:316 print(f" 📄 {dataset_type}: {file_path} ({file_type})")317 else:318 print(f" 📄 No files")319 320 conn.close()321 322 323def main():324 """Main seeding function."""325 print("🌱 XRF Explorer Database Seeding")326 print("=" * 40)327 328 try:329 create_directory_structure()330 copy_test_files()331 conn = create_database_schema()332 populate_database(conn)333 conn.close()334 verify_database()335 336 print("\n🎉 Database seeding completed successfully!")337 print(" Database location: data/wells.db")338 print(" Data files location: data/processed/")339 340 except Exception as e:341 print(f"\n❌ Error during seeding: {e}")342 raise343 344 345if __name__ == "__main__":346 main()