File size: 5,127 Bytes
d965a0a
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
"""
Migration 007: Add user_id columns for multi-tenant support

This migration adds user_id to all tables to enable user-specific data isolation.
Each user will only see their own orders, drivers, assignments, etc.
"""

import sys
from pathlib import Path

# Add parent directory to path
sys.path.insert(0, str(Path(__file__).parent.parent.parent))

from database.connection import execute_write


def up():
    """Add user_id columns and indexes"""

    migrations = [
        # Add user_id column to orders table
        """
        ALTER TABLE orders
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to drivers table
        """
        ALTER TABLE drivers
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to assignments table
        """
        ALTER TABLE assignments
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to exceptions table
        """
        ALTER TABLE exceptions
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to agent_decisions table
        """
        ALTER TABLE agent_decisions
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Add user_id column to metrics table
        """
        ALTER TABLE metrics
        ADD COLUMN IF NOT EXISTS user_id VARCHAR(255);
        """,

        # Create indexes for fast user-based filtering
        """
        CREATE INDEX IF NOT EXISTS idx_orders_user_id ON orders(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_drivers_user_id ON drivers(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_assignments_user_id ON assignments(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_exceptions_user_id ON exceptions(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_agent_decisions_user_id ON agent_decisions(user_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_metrics_user_id ON metrics(user_id);
        """,

        # Create composite indexes for common queries
        """
        CREATE INDEX IF NOT EXISTS idx_orders_user_status ON orders(user_id, status);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_orders_user_created ON orders(user_id, created_at DESC);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_drivers_user_status ON drivers(user_id, status);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_assignments_user_driver ON assignments(user_id, driver_id);
        """,

        """
        CREATE INDEX IF NOT EXISTS idx_assignments_user_order ON assignments(user_id, order_id);
        """,
    ]

    print("Migration 007: Adding user_id columns...")

    for i, sql in enumerate(migrations, 1):
        try:
            print(f"  [{i}/{len(migrations)}] Executing: {sql.strip()[:60]}...")
            execute_write(sql)
            print(f"  Success")
        except Exception as e:
            print(f"  Warning: {e}")
            # Continue even if column already exists

    print("\nMigration 007 complete!")
    print("\nNext steps:")
    print("  1. Existing data will have NULL user_id (that's OK for now)")
    print("  2. New data will automatically get user_id from authentication")
    print("  3. You can optionally run a data migration to assign existing records to a test user")


def down():
    """Remove user_id columns and indexes (rollback)"""

    rollback_migrations = [
        # Drop indexes first
        "DROP INDEX IF EXISTS idx_assignments_user_order;",
        "DROP INDEX IF EXISTS idx_assignments_user_driver;",
        "DROP INDEX IF EXISTS idx_drivers_user_status;",
        "DROP INDEX IF EXISTS idx_orders_user_created;",
        "DROP INDEX IF EXISTS idx_orders_user_status;",
        "DROP INDEX IF EXISTS idx_metrics_user_id;",
        "DROP INDEX IF EXISTS idx_agent_decisions_user_id;",
        "DROP INDEX IF EXISTS idx_exceptions_user_id;",
        "DROP INDEX IF EXISTS idx_assignments_user_id;",
        "DROP INDEX IF EXISTS idx_drivers_user_id;",
        "DROP INDEX IF EXISTS idx_orders_user_id;",

        # Drop columns
        "ALTER TABLE metrics DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE agent_decisions DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE exceptions DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE assignments DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE drivers DROP COLUMN IF EXISTS user_id;",
        "ALTER TABLE orders DROP COLUMN IF EXISTS user_id;",
    ]

    print("Rolling back Migration 007...")

    for i, sql in enumerate(rollback_migrations, 1):
        try:
            print(f"  [{i}/{len(rollback_migrations)}] {sql[:60]}...")
            execute_write(sql)
            print(f"  Success")
        except Exception as e:
            print(f"  Warning: {e}")

    print("\nRollback complete!")


if __name__ == "__main__":
    import sys

    if len(sys.argv) > 1 and sys.argv[1] == "down":
        down()
    else:
        up()