diff --git a/projects/extension/ai/secrets.py b/projects/extension/ai/secrets.py index a52fa8902..1decd2e8c 100644 --- a/projects/extension/ai/secrets.py +++ b/projects/extension/ai/secrets.py @@ -6,6 +6,7 @@ from backoff._typing import Details GUC_SECRETS_MANAGER_URL = "ai.external_functions_executor_url" +GUC_SECRET_ENV_ENABLED = "ai.secret_env_enabled" DEFAULT_SECRETS_MANAGER_PATH = "/api/v1/projects/secrets" @@ -54,9 +55,15 @@ def reveal_secret(plpy, secret_name: str) -> str | None: if secret != "": return secret - env_secret = os.environ.get(secret_name.upper()) - if env_secret is not None: - return env_secret + if not check_secret_permissions(plpy, secret_name): + plpy.error(f"user does not have access to secret '{secret_name}'") + return None + + # check the env var, unless disabled by guc + if get_guc_value(plpy, GUC_SECRET_ENV_ENABLED, "true") == "true": + env_secret = os.environ.get(secret_name.upper()) + if env_secret is not None: + return env_secret if secret_manager_enabled(plpy): secret_optional = fetch_secret(plpy, secret_name) @@ -75,10 +82,6 @@ def fetch_secret(plpy, secret_name: str) -> str | None: plpy.error("secrets manager is not enabled") return None - if not check_secret_permissions(plpy, secret_name): - plpy.error(f"user does not have access to secret '{secret_name}'") - return None - the_url = urljoin( get_guc_value(plpy, GUC_SECRETS_MANAGER_URL, ""), DEFAULT_SECRETS_MANAGER_PATH, diff --git a/projects/extension/tests/privileges/test_privileges.py b/projects/extension/tests/privileges/test_privileges.py index eb3621b91..14d0d79d8 100644 --- a/projects/extension/tests/privileges/test_privileges.py +++ b/projects/extension/tests/privileges/test_privileges.py @@ -115,6 +115,15 @@ def test_secret_privileges(): with pytest.raises(Exception, match="permission denied for table"): cur.execute("select * from ai._secret_permissions") + # jill cannot access the env var + with psycopg.connect(db_url("jill", "privs")) as con: + with con.cursor() as cur: + cur.execute( + "SET ai.external_functions_executor_url='http://localhost:8000'" + ) + with pytest.raises(Exception, match="user does not have access"): + cur.execute("select ai.reveal_secret('TEST_ENV_SECRET')") + # alice can access all the secrets and grant them to fred and joey with psycopg.connect(db_url("alice", "privs")) as con: with con.cursor() as cur: @@ -122,6 +131,7 @@ def test_secret_privileges(): "SET ai.external_functions_executor_url='http://localhost:8000'" ) cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + cur.execute("select ai.reveal_secret('TEST_ENV_SECRET')") cur.execute("select ai.grant_secret('OPENAI_API_KEY', 'joey')") cur.execute("select ai.grant_secret('*', 'joey2')") cur.execute("select ai.grant_secret('OPENAI_API_KEY', 'fred')") @@ -142,6 +152,7 @@ def test_secret_privileges(): "SET ai.external_functions_executor_url='http://localhost:8000'" ) cur.execute("select ai.grant_secret('OPENAI_API_KEY', 'jill')") + cur.execute("select ai.grant_secret('TEST_ENV_SECRET', 'jill')") # jill can access the secret granted to her but not the other one with psycopg.connect(db_url("jill", "privs")) as con: @@ -150,6 +161,7 @@ def test_secret_privileges(): "SET ai.external_functions_executor_url='http://localhost:8000'" ) cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") + cur.execute("select ai.reveal_secret('TEST_ENV_SECRET')") with pytest.raises(Exception, match="user does not have access"): cur.execute("select ai.reveal_secret('OPENAI_API_KEY_2')") @@ -195,6 +207,7 @@ def test_secret_privileges(): ) cur.execute("select ai.reveal_secret('OPENAI_API_KEY')") cur.execute("select ai.reveal_secret('OPENAI_API_KEY_2')") + cur.execute("select ai.reveal_secret('TEST_ENV_SECRET')") # alice can revoke the * privilege from jill with psycopg.connect(db_url("alice", "privs")) as con: