|
11 | 11 | from alembic import op |
12 | 12 | from onyx.configs.app_configs import DB_READONLY_PASSWORD |
13 | 13 | from onyx.configs.app_configs import DB_READONLY_USER |
14 | | -from shared_configs.configs import MULTI_TENANT |
15 | 14 |
|
16 | 15 |
|
17 | 16 | # revision identifiers, used by Alembic. |
|
22 | 21 |
|
23 | 22 |
|
24 | 23 | def upgrade() -> None: |
25 | | - if MULTI_TENANT: |
| 24 | + # Enable pg_trgm extension if not already enabled |
| 25 | + op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") |
26 | 26 |
|
27 | | - # Enable pg_trgm extension if not already enabled |
28 | | - op.execute("CREATE EXTENSION IF NOT EXISTS pg_trgm") |
| 27 | + # Create the read-only db user if it does not already exist. |
| 28 | + if not (DB_READONLY_USER and DB_READONLY_PASSWORD): |
| 29 | + raise Exception("DB_READONLY_USER or DB_READONLY_PASSWORD is not set") |
29 | 30 |
|
30 | | - # Create read-only db user here only in multi-tenant mode. For single-tenant mode, |
31 | | - # the user is created in the standard migration. |
32 | | - if not (DB_READONLY_USER and DB_READONLY_PASSWORD): |
33 | | - raise Exception("DB_READONLY_USER or DB_READONLY_PASSWORD is not set") |
34 | | - |
35 | | - op.execute( |
36 | | - text( |
37 | | - f""" |
38 | | - DO $$ |
39 | | - BEGIN |
40 | | - -- Check if the read-only user already exists |
41 | | - IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN |
42 | | - -- Create the read-only user with the specified password |
43 | | - EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}'); |
44 | | - -- First revoke all privileges to ensure a clean slate |
45 | | - EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}'); |
46 | | - -- Grant only the CONNECT privilege to allow the user to connect to the database |
47 | | - -- but not perform any operations without additional specific grants |
48 | | - EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}'); |
49 | | - END IF; |
50 | | - END |
51 | | - $$; |
52 | | - """ |
53 | | - ) |
54 | | - ) |
55 | | - |
56 | | - |
57 | | -def downgrade() -> None: |
58 | | - if MULTI_TENANT: |
59 | | - # Drop read-only db user here only in single tenant mode. For multi-tenant mode, |
60 | | - # the user is dropped in the alembic_tenants migration. |
61 | | - |
62 | | - op.execute( |
63 | | - text( |
64 | | - f""" |
| 31 | + op.execute( |
| 32 | + text( |
| 33 | + f""" |
65 | 34 | DO $$ |
66 | 35 | BEGIN |
67 | | - IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN |
68 | | - -- First revoke all privileges from the database |
| 36 | + -- Check if the read-only user already exists |
| 37 | + IF NOT EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN |
| 38 | + -- Create the read-only user with the specified password |
| 39 | + EXECUTE format('CREATE USER %I WITH PASSWORD %L', '{DB_READONLY_USER}', '{DB_READONLY_PASSWORD}'); |
| 40 | + -- First revoke all privileges to ensure a clean slate |
69 | 41 | EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}'); |
70 | | - -- Then revoke all privileges from the public schema |
71 | | - EXECUTE format('REVOKE ALL ON SCHEMA public FROM %I', '{DB_READONLY_USER}'); |
72 | | - -- Then drop the user |
73 | | - EXECUTE format('DROP USER %I', '{DB_READONLY_USER}'); |
| 42 | + -- Grant only the CONNECT privilege to allow the user to connect to the database |
| 43 | + -- but not perform any operations without additional specific grants |
| 44 | + EXECUTE format('GRANT CONNECT ON DATABASE %I TO %I', current_database(), '{DB_READONLY_USER}'); |
74 | 45 | END IF; |
75 | 46 | END |
76 | 47 | $$; |
77 | | - """ |
78 | | - ) |
| 48 | + """ |
| 49 | + ) |
| 50 | + ) |
| 51 | + |
| 52 | + |
| 53 | +def downgrade() -> None: |
| 54 | + op.execute( |
| 55 | + text( |
| 56 | + f""" |
| 57 | + DO $$ |
| 58 | + BEGIN |
| 59 | + IF EXISTS (SELECT FROM pg_catalog.pg_roles WHERE rolname = '{DB_READONLY_USER}') THEN |
| 60 | + -- First revoke all privileges from the database |
| 61 | + EXECUTE format('REVOKE ALL ON DATABASE %I FROM %I', current_database(), '{DB_READONLY_USER}'); |
| 62 | + -- Then revoke all privileges from the public schema |
| 63 | + EXECUTE format('REVOKE ALL ON SCHEMA public FROM %I', '{DB_READONLY_USER}'); |
| 64 | + -- Then drop the user |
| 65 | + EXECUTE format('DROP USER %I', '{DB_READONLY_USER}'); |
| 66 | + END IF; |
| 67 | + END |
| 68 | + $$; |
| 69 | + """ |
79 | 70 | ) |
80 | | - op.execute(text("DROP EXTENSION IF EXISTS pg_trgm")) |
| 71 | + ) |
| 72 | + op.execute(text("DROP EXTENSION IF EXISTS pg_trgm")) |
0 commit comments