51 lines
1.4 KiB
Python
51 lines
1.4 KiB
Python
#!/usr/bin/env python3
|
|
"""
|
|
Database initialization script to create PostgreSQL tables
|
|
"""
|
|
import psycopg2
|
|
from config import get_database_url
|
|
|
|
def create_database_tables():
|
|
"""Create the required tables in PostgreSQL"""
|
|
conn = psycopg2.connect(get_database_url())
|
|
cur = conn.cursor()
|
|
|
|
# Create allowed_keys table
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS allowed_keys (
|
|
key VARCHAR(64) PRIMARY KEY,
|
|
label TEXT
|
|
)
|
|
""")
|
|
|
|
# Create access_windows table
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS access_windows (
|
|
id SERIAL PRIMARY KEY,
|
|
key VARCHAR(64) REFERENCES allowed_keys(key) ON DELETE CASCADE,
|
|
start_datetime TIMESTAMP,
|
|
end_datetime TIMESTAMP
|
|
)
|
|
""")
|
|
|
|
# Create token_usage table
|
|
cur.execute("""
|
|
CREATE TABLE IF NOT EXISTS token_usage (
|
|
id SERIAL PRIMARY KEY,
|
|
key VARCHAR(64) REFERENCES allowed_keys(key) ON DELETE CASCADE,
|
|
timestamp TIMESTAMP,
|
|
prompt_tokens INTEGER,
|
|
completion_tokens INTEGER,
|
|
total_tokens INTEGER
|
|
)
|
|
""")
|
|
|
|
conn.commit()
|
|
cur.close()
|
|
conn.close()
|
|
print("PostgreSQL tables created successfully")
|
|
|
|
if __name__ == "__main__":
|
|
print("Creating PostgreSQL tables...")
|
|
create_database_tables()
|
|
print("Database initialization complete!") |