106 lines
3.1 KiB
Python
106 lines
3.1 KiB
Python
import psycopg2
|
|
from psycopg2 import sql
|
|
from config import DB_HOST, DB_PORT, DB_NAME, DB_USER, DB_PASSWORD
|
|
|
|
def get_db_connection():
|
|
"""Connect to the PostgreSQL database."""
|
|
return psycopg2.connect(
|
|
host=DB_HOST,
|
|
port=DB_PORT,
|
|
dbname=DB_NAME,
|
|
user=DB_USER,
|
|
password=DB_PASSWORD
|
|
)
|
|
|
|
def create_tables():
|
|
"""Create database tables if they don't exist."""
|
|
commands = [
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS trains (
|
|
train_id VARCHAR PRIMARY KEY,
|
|
train_type VARCHAR,
|
|
model VARCHAR,
|
|
capacity INT,
|
|
age_years FLOAT,
|
|
maintenance_history JSONB,
|
|
technical_specs JSONB
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS schedules (
|
|
schedule_id UUID PRIMARY KEY,
|
|
train_id VARCHAR,
|
|
route_id VARCHAR,
|
|
departure_station VARCHAR,
|
|
arrival_station VARCHAR,
|
|
scheduled_departure TIMESTAMP,
|
|
scheduled_arrival TIMESTAMP,
|
|
platform_departure VARCHAR,
|
|
platform_arrival VARCHAR,
|
|
distance_km FLOAT,
|
|
scheduled_stops INT,
|
|
service_type VARCHAR,
|
|
FOREIGN KEY (train_id) REFERENCES trains(train_id)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS actual_journeys (
|
|
journey_id UUID PRIMARY KEY,
|
|
schedule_id UUID,
|
|
actual_departure TIMESTAMP,
|
|
actual_arrival TIMESTAMP,
|
|
actual_platform_departure VARCHAR,
|
|
actual_platform_arrival VARCHAR,
|
|
delay_departure_minutes INT,
|
|
delay_arrival_minutes INT,
|
|
cancellation_flag BOOLEAN,
|
|
cancellation_reason VARCHAR,
|
|
passenger_count INT,
|
|
load_factor FLOAT,
|
|
FOREIGN KEY (schedule_id) REFERENCES schedules(schedule_id)
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS stations (
|
|
station_id VARCHAR PRIMARY KEY,
|
|
station_name VARCHAR,
|
|
latitude FLOAT,
|
|
longitude FLOAT,
|
|
elevation FLOAT,
|
|
number_of_platforms INT,
|
|
daily_passenger_volume INT,
|
|
station_category INT,
|
|
facilities JSONB,
|
|
connection_types JSONB
|
|
)
|
|
""",
|
|
"""
|
|
CREATE TABLE IF NOT EXISTS weather_history (
|
|
weather_id UUID PRIMARY KEY,
|
|
station_id VARCHAR,
|
|
timestamp TIMESTAMP,
|
|
temperature FLOAT,
|
|
precipitation FLOAT,
|
|
wind_speed FLOAT,
|
|
wind_direction INT,
|
|
humidity FLOAT,
|
|
pressure FLOAT,
|
|
visibility FLOAT,
|
|
weather_condition VARCHAR,
|
|
FOREIGN KEY (station_id) REFERENCES stations(station_id)
|
|
)
|
|
"""
|
|
]
|
|
conn = None
|
|
try:
|
|
conn = get_db_connection()
|
|
cur = conn.cursor()
|
|
for command in commands:
|
|
cur.execute(command)
|
|
cur.close()
|
|
conn.commit()
|
|
except Exception as e:
|
|
print(f"Error creating tables: {e}")
|
|
finally:
|
|
if conn is not None:
|
|
conn.close() |