databricks-tellr-app 0.3.13.dev4__tar.gz → 0.3.13.dev11__tar.gz

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-tellr-app
-Version: 0.3.13.dev4
+Version: 0.3.13.dev11
 Summary: Tellr application package for Databricks Apps
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -24,7 +24,7 @@ Requires-Dist: tinycss2==1.3.0
 Requires-Dist: esprima==4.0.1
 Requires-Dist: sqlalchemy==2.0.36
 Requires-Dist: psycopg2-binary==2.9.10
-Requires-Dist: databricks-sql-connector==3.0.1
+Requires-Dist: databricks-sql-connector>=3.0.0
 Requires-Dist: httpx==0.28.1
 Requires-Dist: jinja2==3.1.2
 Requires-Dist: python-dotenv==1.2.1
@@ -32,12 +32,12 @@ Requires-Dist: pyyaml==6.0.1
 Requires-Dist: pandas==2.3.3
 Requires-Dist: python-pptx==0.6.23
 Requires-Dist: playwright==1.48.0
-Requires-Dist: Pillow==11.3.0
-Requires-Dist: google-api-python-client==2.190.0
-Requires-Dist: google-auth-oauthlib==1.2.4
-Requires-Dist: google-auth-httplib2==0.3.0
-Requires-Dist: cryptography==46.0.3
-Requires-Dist: sse-starlette==3.0.3
+Requires-Dist: Pillow>=10.0.0
+Requires-Dist: google-api-python-client>=2.100.0
+Requires-Dist: google-auth-oauthlib>=1.2.0
+Requires-Dist: google-auth-httplib2>=0.2.0
+Requires-Dist: cryptography>=42.0.0
+Requires-Dist: sse-starlette>=3.0.0
 Requires-Dist: databricks-sdk==0.120.0
 Requires-Dist: mlflow[databricks]==3.14.0
 Requires-Dist: mlflow-skinny==3.14.0

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/databricks_tellr_app/run.py

@@ -62,18 +62,6 @@ def init_database(seed_databricks_defaults: bool = False) -> None:
         logger.error(f"Failed to seed defaults: {e}\n{tb}")
         raise SystemExit(1) from e
 
-    # Seed/load the Fernet master key pre-fork (SDR-4437 CRITICAL-3) so all
-    # uvicorn workers find the encryption_keys row already present.
-    logger.info("Ensuring encryption key...")
-    try:
-        from src.core.encryption import ensure_encryption_key
-        ensure_encryption_key()
-        logger.info("Encryption key ready")
-    except Exception as e:
-        tb = traceback.format_exc()
-        logger.error(f"Failed to ensure encryption key: {e}\n{tb}")
-        raise SystemExit(1) from e
-
 
 def main() -> None:
     """Start the uvicorn server."""

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/databricks_tellr_app.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: databricks-tellr-app
-Version: 0.3.13.dev4
+Version: 0.3.13.dev11
 Summary: Tellr application package for Databricks Apps
 Requires-Python: >=3.10
 Description-Content-Type: text/markdown
@@ -24,7 +24,7 @@ Requires-Dist: tinycss2==1.3.0
 Requires-Dist: esprima==4.0.1
 Requires-Dist: sqlalchemy==2.0.36
 Requires-Dist: psycopg2-binary==2.9.10
-Requires-Dist: databricks-sql-connector==3.0.1
+Requires-Dist: databricks-sql-connector>=3.0.0
 Requires-Dist: httpx==0.28.1
 Requires-Dist: jinja2==3.1.2
 Requires-Dist: python-dotenv==1.2.1
@@ -32,12 +32,12 @@ Requires-Dist: pyyaml==6.0.1
 Requires-Dist: pandas==2.3.3
 Requires-Dist: python-pptx==0.6.23
 Requires-Dist: playwright==1.48.0
-Requires-Dist: Pillow==11.3.0
-Requires-Dist: google-api-python-client==2.190.0
-Requires-Dist: google-auth-oauthlib==1.2.4
-Requires-Dist: google-auth-httplib2==0.3.0
-Requires-Dist: cryptography==46.0.3
-Requires-Dist: sse-starlette==3.0.3
+Requires-Dist: Pillow>=10.0.0
+Requires-Dist: google-api-python-client>=2.100.0
+Requires-Dist: google-auth-oauthlib>=1.2.0
+Requires-Dist: google-auth-httplib2>=0.2.0
+Requires-Dist: cryptography>=42.0.0
+Requires-Dist: sse-starlette>=3.0.0
 Requires-Dist: databricks-sdk==0.120.0
 Requires-Dist: mlflow[databricks]==3.14.0
 Requires-Dist: mlflow-skinny==3.14.0

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/databricks_tellr_app.egg-info/SOURCES.txt

@@ -76,7 +76,6 @@ setup.py
 ./src/database/__init__.py
 ./src/database/models/__init__.py
 ./src/database/models/deck_contributor.py
-./src/database/models/encryption_key.py
 ./src/database/models/feedback.py
 ./src/database/models/genie_space.py
 ./src/database/models/google_global_credentials.py
@@ -220,7 +219,6 @@ src/core/user_context.py
 src/database/__init__.py
 src/database/models/__init__.py
 src/database/models/deck_contributor.py
-src/database/models/encryption_key.py
 src/database/models/feedback.py
 src/database/models/genie_space.py
 src/database/models/google_global_credentials.py

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/databricks_tellr_app.egg-info/requires.txt

@@ -18,7 +18,7 @@ tinycss2==1.3.0
 esprima==4.0.1
 sqlalchemy==2.0.36
 psycopg2-binary==2.9.10
-databricks-sql-connector==3.0.1
+databricks-sql-connector>=3.0.0
 httpx==0.28.1
 jinja2==3.1.2
 python-dotenv==1.2.1
@@ -26,12 +26,12 @@ pyyaml==6.0.1
 pandas==2.3.3
 python-pptx==0.6.23
 playwright==1.48.0
-Pillow==11.3.0
-google-api-python-client==2.190.0
-google-auth-oauthlib==1.2.4
-google-auth-httplib2==0.3.0
-cryptography==46.0.3
-sse-starlette==3.0.3
+Pillow>=10.0.0
+google-api-python-client>=2.100.0
+google-auth-oauthlib>=1.2.0
+google-auth-httplib2>=0.2.0
+cryptography>=42.0.0
+sse-starlette>=3.0.0
 databricks-sdk==0.120.0
 mlflow[databricks]==3.14.0
 mlflow-skinny==3.14.0

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "databricks-tellr-app"
-version = "0.3.13.dev4"
+version = "0.3.13.dev11"
 description = "Tellr application package for Databricks Apps"
 readme = "README.md"
 requires-python = ">=3.10"
@@ -29,7 +29,7 @@ dependencies = [
     "esprima==4.0.1",
     "sqlalchemy==2.0.36",
     "psycopg2-binary==2.9.10",
-    "databricks-sql-connector==3.0.1",
+    "databricks-sql-connector>=3.0.0",
     "httpx==0.28.1",
     "jinja2==3.1.2",
     "python-dotenv==1.2.1",
@@ -37,12 +37,12 @@ dependencies = [
     "pandas==2.3.3",
     "python-pptx==0.6.23",
     "playwright==1.48.0",
-    "Pillow==11.3.0",
-    "google-api-python-client==2.190.0",
-    "google-auth-oauthlib==1.2.4",
-    "google-auth-httplib2==0.3.0",
-    "cryptography==46.0.3",
-    "sse-starlette==3.0.3",
+    "Pillow>=10.0.0",
+    "google-api-python-client>=2.100.0",
+    "google-auth-oauthlib>=1.2.0",
+    "google-auth-httplib2>=0.2.0",
+    "cryptography>=42.0.0",
+    "sse-starlette>=3.0.0",
     # --- Resolver-speedup pins: the backtracking-driver cluster ---
     # Databricks Apps installs this wheel in a platform-managed BUILD phase with a
     # ~10-minute budget. Left ranged, the databricks-langchain ->
@@ -58,8 +58,6 @@ dependencies = [
     # via a devloop deploy if these constraints change. Formerly injected at
     # deploy time as _RESOLVER_SPEEDUP_PINS; now carried in the wheel so a bare
     # `pip install databricks-tellr-app` is safe.
-    # SDR-4437 MEDIUM-7 (July 2026): the 7 formerly-ranged direct deps above
-    # this block are now exact-pinned too; same refresh discipline.
     "databricks-sdk==0.120.0",
     "mlflow[databricks]==3.14.0",
     "mlflow-skinny==3.14.0",

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/src/api/services/chat_service.py

@@ -80,6 +80,11 @@ class ChatService:
         # This avoids re-parsing HTML on every request
         self._deck_cache: Dict[str, SlideDeck] = {}
 
+        # DB version each cached deck corresponds to. The cache is per-process
+        # while prod runs multiple uvicorn workers sharing one database, so a
+        # cache hit is only valid if the DB version hasn't moved on.
+        self._deck_cache_versions: Dict[str, int] = {}
+
         logger.info("ChatService initialized successfully")
 
     def _substitute_images_for_response(self, deck_dict, raw_html=None):
@@ -600,7 +605,7 @@ class ChatService:
                     slide_deck_dict = current_deck.to_dict()
 
                 try:
-                    session_manager.save_slide_deck(
+                    save_result = session_manager.save_slide_deck(
                         session_id=session_id,
                         title=current_deck.title,
                         html_content=current_deck.knit(),
@@ -610,6 +615,7 @@ class ChatService:
                         modified_by=_user,
                         expected_version=_deck_version_before_llm,
                     )
+                    self._record_deck_version(session_id, save_result)
                 except VersionConflictError:
                     logger.warning(
                         "Chat save rejected: deck was edited during LLM call, reloading",
@@ -1385,7 +1391,7 @@ class ChatService:
                 slide_deck_dict = current_deck.to_dict()
 
             try:
-                session_manager.save_slide_deck(
+                save_result = session_manager.save_slide_deck(
                     session_id=session_id,
                     title=current_deck.title,
                     html_content=current_deck.knit(),
@@ -1395,6 +1401,7 @@ class ChatService:
                     modified_by=_user,
                     expected_version=_deck_version_before_llm,
                 )
+                self._record_deck_version(session_id, save_result)
             except VersionConflictError:
                 logger.warning(
                     "Chat save rejected: deck was edited during LLM call, reloading",
@@ -1921,18 +1928,63 @@ class ChatService:
         context = "\n\n[Attached images]\n" + "\n".join(image_descriptions)
         return message + context
 
+    def _deck_versions(self) -> Dict[str, int]:
+        """Return the cached-deck version map, creating it if needed.
+
+        Lazy creation keeps instances built without __init__ (test harnesses
+        use ChatService.__new__) working. Callers must hold _cache_lock.
+        """
+        try:
+            return self._deck_cache_versions
+        except AttributeError:
+            self._deck_cache_versions = {}
+            return self._deck_cache_versions
+
+    def _record_deck_version(self, session_id: str, save_result: Optional[Dict[str, Any]]) -> None:
+        """Remember which DB version the cached deck now corresponds to.
+
+        Call after save_slide_deck so the next _get_or_load_deck cache hit
+        validates cleanly instead of reloading from the database.
+        """
+        version = (save_result or {}).get("version")
+        with self._cache_lock:
+            if version is not None:
+                self._deck_versions()[session_id] = version
+            else:
+                self._deck_versions().pop(session_id, None)
+
     def _get_or_load_deck(self, session_id: str) -> Optional[SlideDeck]:
         """Get deck from cache or load from database.
 
         Thread-safe access to deck cache using _cache_lock.
-        
+
+        The cache is per-process while prod runs multiple uvicorn workers
+        sharing one database, so a cached deck is only served if its recorded
+        version matches the current DB version; otherwise it is reloaded.
+
         Uses deck_dict (slides array with individual scripts) when available,
         falling back to from_html_string for legacy data.
         """
         # Check cache first (with lock)
         with self._cache_lock:
-            if session_id in self._deck_cache:
-                return self._deck_cache[session_id]
+            cached_deck = self._deck_cache.get(session_id)
+            cached_version = self._deck_versions().get(session_id)
+
+        if cached_deck is not None:
+            db_version = self._get_deck_version(session_id)
+            if db_version is None or cached_version == db_version:
+                return cached_deck
+            logger.info(
+                "Deck cache stale, reloading from database",
+                extra={
+                    "session_id": session_id,
+                    "cached_version": cached_version,
+                    "db_version": db_version,
+                },
+            )
+            with self._cache_lock:
+                self._deck_cache.pop(session_id, None)
+                self._deck_versions().pop(session_id, None)
 
         # Try to load from database (outside lock to avoid blocking)
         session_manager = get_session_manager()
@@ -1966,6 +2018,10 @@ class ChatService:
             # Store in cache (with lock)
             with self._cache_lock:
                 self._deck_cache[session_id] = deck
+                if deck_data.get("version") is not None:
+                    self._deck_versions()[session_id] = deck_data["version"]
+                else:
+                    self._deck_versions().pop(session_id, None)
             return deck
         except Exception as e:
             logger.warning(f"Failed to load deck from database: {e}")
@@ -1975,23 +2031,18 @@ class ChatService:
     def _get_deck_version(self, session_id: str) -> Optional[int]:
         """Read the current deck version from the database.
 
-        Called once before the LLM runs to capture the version for
-        optimistic locking. Uses the existing get_slide_deck which
-        returns the full deck dict including the version column.
-
-        On Lakebase/PostgreSQL this is a single-row indexed lookup (~1-2ms).
+        Used on every deck-cache hit (multi-worker cache validation) and
+        before LLM calls (optimistic locking), so it must stay cheap: a
+        single-column indexed lookup, never a full deck fetch.
 
         Returns:
             Current deck version number, or None if no deck exists.
         """
         session_manager = get_session_manager()
         try:
-            deck_data = session_manager.get_slide_deck(session_id)
-            if deck_data:
-                return deck_data.get("version")
+            return session_manager.get_slide_deck_version(session_id)
         except Exception:
-            pass
-        return None
+            return None
 
     @staticmethod
     def _reindex_slide_ids(deck: "SlideDeck") -> None:
@@ -2008,6 +2059,7 @@ class ChatService:
         """Remove the cached deck for a session so the next read hits the DB."""
         with self._cache_lock:
             self._deck_cache.pop(session_id, None)
+            self._deck_versions().pop(session_id, None)
 
     def _replace_slide_htmls_from_cache(self, session_id: str, slide_context: Dict[str, Any]) -> Dict[str, Any]:
         """Replace frontend-supplied slide_htmls with backend cache versions.
@@ -2540,7 +2592,7 @@ class ChatService:
         # Persist to database
         deck_dict = current_deck.to_dict()
         session_manager = get_session_manager()
-        session_manager.save_slide_deck(
+        save_result = session_manager.save_slide_deck(
             session_id=session_id,
             title=current_deck.title,
             html_content=current_deck.knit(),
@@ -2549,6 +2601,7 @@ class ChatService:
             deck_dict=deck_dict,
             expected_version=expected_version,
         )
+        self._record_deck_version(session_id, save_result)
 
         # Create save point
         try:
@@ -2616,7 +2669,7 @@ class ChatService:
         # Persist to database
         deck_dict = current_deck.to_dict()
         session_manager = get_session_manager()
-        session_manager.save_slide_deck(
+        save_result = session_manager.save_slide_deck(
             session_id=session_id,
             title=current_deck.title,
             html_content=current_deck.knit(),
@@ -2625,6 +2678,7 @@ class ChatService:
             deck_dict=deck_dict,
             expected_version=expected_version,
         )
+        self._record_deck_version(session_id, save_result)
 
         # Create save point immediately after persisting
         try:
@@ -2681,7 +2735,7 @@ class ChatService:
         # Persist to database
         deck_dict = current_deck.to_dict()
         session_manager = get_session_manager()
-        session_manager.save_slide_deck(
+        save_result = session_manager.save_slide_deck(
             session_id=session_id,
             title=current_deck.title,
             html_content=current_deck.knit(),
@@ -2690,6 +2744,7 @@ class ChatService:
             deck_dict=deck_dict,
             expected_version=expected_version,
         )
+        self._record_deck_version(session_id, save_result)
 
         # Create save point
         try:
@@ -2744,7 +2799,7 @@ class ChatService:
         # Persist to database
         deck_dict = current_deck.to_dict()
         session_manager = get_session_manager()
-        session_manager.save_slide_deck(
+        save_result = session_manager.save_slide_deck(
             session_id=session_id,
             title=current_deck.title,
             html_content=current_deck.knit(),
@@ -2753,6 +2808,7 @@ class ChatService:
             deck_dict=deck_dict,
             expected_version=expected_version,
         )
+        self._record_deck_version(session_id, save_result)
 
         # Create save point
         try:

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/src/api/services/session_manager.py

@@ -723,6 +723,33 @@ class SessionManager:
                 "version": deck.version,
             }
 
+    def get_slide_deck_version(self, session_id: str) -> Optional[int]:
+        """Return only the deck's current version number.
+
+        Selects a single column — no deck_json parse, verification merge, or
+        display-name resolution — so it is cheap enough to run on every
+        deck-cache hit (multi-worker cache validation) and before LLM calls
+        (optimistic locking).
+
+        Args:
+            session_id: Session to check (contributor sessions resolve to the
+                owner's deck)
+
+        Returns:
+            Current version number, or None if the session/deck doesn't exist
+        """
+        with get_db_session() as db:
+            try:
+                session = self._get_session_or_raise(db, session_id)
+            except SessionNotFoundError:
+                return None
+            deck_owner = self._get_deck_owner_session(db, session)
+            return (
+                db.query(SessionSlideDeck.version)
+                .filter(SessionSlideDeck.session_id == deck_owner.id)
+                .scalar()
+            )
+
     def get_slide_deck(self, session_id: str) -> Optional[Dict[str, Any]]:
         """Get slide deck for a session with verification merged by content hash.
 

databricks_tellr_app-0.3.13.dev11/src/core/encryption.py

@@ -0,0 +1,86 @@
+"""Fernet symmetric encryption for sensitive data (OAuth credentials, tokens).
+
+Uses AES-128-CBC with HMAC-SHA256 authentication via the ``cryptography`` library.
+The encryption key is read from the ``GOOGLE_OAUTH_ENCRYPTION_KEY`` environment
+variable.  If the variable is not set, a key is auto-generated on first access,
+persisted to ``.encryption_key`` in the project root, and a warning is logged.
+"""
+
+import logging
+import os
+from functools import lru_cache
+from pathlib import Path
+
+from cryptography.fernet import Fernet, InvalidToken
+
+logger = logging.getLogger(__name__)
+
+_ENV_KEY_NAME = "GOOGLE_OAUTH_ENCRYPTION_KEY"
+_KEY_FILE = Path(__file__).resolve().parents[2] / ".encryption_key"
+
+
+@lru_cache(maxsize=1)
+def get_encryption_key() -> bytes:
+    """Return the Fernet encryption key.
+
+    Resolution order:
+    1. ``GOOGLE_OAUTH_ENCRYPTION_KEY`` env var (production).
+    2. ``.encryption_key`` file in project root (persisted across restarts).
+    3. Generate a new key → write to the file → set the env var.
+    """
+    # 1. Env var takes precedence
+    key = os.getenv(_ENV_KEY_NAME)
+    if key:
+        return key.encode()
+
+    # Production guard: fail fast if the env var is not set
+    if os.getenv("ENVIRONMENT") == "production":
+        raise RuntimeError(
+            "GOOGLE_OAUTH_ENCRYPTION_KEY environment variable is required in production. "
+            'Generate one with: python -c "from cryptography.fernet import Fernet; '
+            'print(Fernet.generate_key().decode())"'
+        )
+
+    # 2. Read from persisted key file
+    if _KEY_FILE.exists():
+        key = _KEY_FILE.read_text().strip()
+        if key:
+            os.environ[_ENV_KEY_NAME] = key
+            logger.info("Loaded encryption key from %s", _KEY_FILE)
+            return key.encode()
+
+    # 3. Generate, persist, and warn
+    generated = Fernet.generate_key()
+    try:
+        _KEY_FILE.write_text(generated.decode())
+        _KEY_FILE.chmod(0o600)  # owner-only read/write
+        logger.warning(
+            "GOOGLE_OAUTH_ENCRYPTION_KEY not set — auto-generated and saved to %s. "
+            "Set the env var in production.",
+            _KEY_FILE,
+        )
+    except OSError:
+        logger.warning(
+            "GOOGLE_OAUTH_ENCRYPTION_KEY not set — auto-generated a key but "
+            "could not persist to %s. Key will be lost on restart.",
+            _KEY_FILE,
+        )
+
+    os.environ[_ENV_KEY_NAME] = generated.decode()
+    return generated
+
+
+def encrypt_data(plaintext: str) -> str:
+    """Encrypt *plaintext* and return a base64-encoded ciphertext string."""
+    f = Fernet(get_encryption_key())
+    return f.encrypt(plaintext.encode()).decode()
+
+
+def decrypt_data(ciphertext: str) -> str:
+    """Decrypt a base64-encoded *ciphertext* string and return the plaintext.
+
+    Raises:
+        cryptography.fernet.InvalidToken: If the key does not match or data is corrupt.
+    """
+    f = Fernet(get_encryption_key())
+    return f.decrypt(ciphertext.encode()).decode()

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/src/core/lakebase.py

@@ -280,6 +280,120 @@ def get_lakebase_connection_url(
     return url
 
 
+def setup_lakebase_schema(
+    instance_name: str,
+    schema: str,
+    client_id: str,
+    user: Optional[str] = None,
+    client: Optional[WorkspaceClient] = None,
+    endpoint_name: Optional[str] = None,
+    host: Optional[str] = None,
+    lakebase_type: str = "provisioned",
+) -> None:
+    """
+    Set up schema and grant permissions in Lakebase for an app.
+
+    Supports both provisioned and autoscaling Lakebase.
+
+    This function connects to the Lakebase instance and:
+    1. Creates the schema if it doesn't exist
+    2. Grants necessary permissions to the app's service principal (client_id)
+
+    Note: This should be called after the app is created and the database
+    resource is attached, as the client_id comes from the app's service principal.
+
+    Args:
+        instance_name: Name of the Lakebase instance
+        schema: Schema name to create
+        client_id: App's service principal client ID (becomes Postgres role)
+        user: Postgres username for connection (uses default if not provided)
+        client: Optional WorkspaceClient
+        endpoint_name: Full endpoint resource name (autoscaling only)
+        host: Pre-resolved host (autoscaling only)
+        lakebase_type: 'autoscaling' or 'provisioned'
+
+    Raises:
+        LakebaseError: If schema setup fails
+    """
+    logger.info(
+        "Setting up Lakebase schema",
+        extra={
+            "instance_name": instance_name, "schema": schema,
+            "client_id": client_id, "lakebase_type": lakebase_type,
+        },
+    )
+
+    try:
+        # Get connection info
+        conn_info = get_lakebase_connection_info(
+            instance_name=instance_name,
+            user=user,
+            client=client,
+            endpoint_name=endpoint_name,
+            host=host,
+            lakebase_type=lakebase_type,
+        )
+
+        # Use psycopg2 directly for DDL operations
+        import psycopg2
+
+        conn = psycopg2.connect(
+            host=conn_info["host"],
+            port=conn_info["port"],
+            user=conn_info["user"],
+            password=conn_info["password"],
+            dbname=conn_info["database"],
+            sslmode="require",
+        )
+        conn.autocommit = True
+
+        with conn.cursor() as cur:
+            # Create schema if not exists
+            cur.execute(f'CREATE SCHEMA IF NOT EXISTS "{schema}"')
+            logger.info(f"Schema created or verified: {schema}")
+
+            # Grant permissions to app's service principal
+            # The client_id is automatically available as a Postgres role when
+            # the database is added as an app resource
+            cur.execute(f'GRANT USAGE ON SCHEMA "{schema}" TO "{client_id}"')
+            cur.execute(f'GRANT CREATE ON SCHEMA "{schema}" TO "{client_id}"')
+
+            # Grant permissions on all tables in schema (for existing tables)
+            cur.execute(
+                f'GRANT SELECT, INSERT, UPDATE, DELETE ON ALL TABLES IN SCHEMA "{schema}" TO "{client_id}"'
+            )
+
+            # Grant permissions on all sequences in schema (for auto-increment columns)
+            cur.execute(
+                f'GRANT USAGE, SELECT ON ALL SEQUENCES IN SCHEMA "{schema}" TO "{client_id}"'
+            )
+
+            # Set default privileges for future tables
+            cur.execute(
+                f'ALTER DEFAULT PRIVILEGES IN SCHEMA "{schema}" '
+                f'GRANT SELECT, INSERT, UPDATE, DELETE ON TABLES TO "{client_id}"'
+            )
+
+            # Set default privileges for future sequences
+            cur.execute(
+                f'ALTER DEFAULT PRIVILEGES IN SCHEMA "{schema}" '
+                f'GRANT USAGE, SELECT ON SEQUENCES TO "{client_id}"'
+            )
+
+            logger.info(f"Permissions granted to {client_id} on schema {schema}")
+
+        conn.close()
+
+    except ImportError:
+        raise LakebaseError(
+            "psycopg2 is required for Lakebase schema setup. "
+            "Install with: pip install psycopg2-binary"
+        )
+    except Exception as e:
+        logger.error(f"Failed to setup Lakebase schema: {e}", exc_info=True)
+        raise LakebaseError(f"Lakebase schema setup failed: {e}") from e
+
+
 def initialize_lakebase_tables(
     instance_name: str,
     schema: str,

{databricks_tellr_app-0.3.13.dev4 → databricks_tellr_app-0.3.13.dev11}/src/database/models/__init__.py

@@ -1,7 +1,6 @@
 """Database models."""
 
 from src.database.models.deck_contributor import DeckContributor
-from src.database.models.encryption_key import EncryptionKey
 from src.database.models.feedback import FeedbackConversation, SurveyResponse
 from src.database.models.genie_space import ConfigGenieSpace
 from src.database.models.google_global_credentials import GoogleGlobalCredentials
@@ -35,7 +34,6 @@ __all__ = [
     "ChatRequest",
     "ConfigGenieSpace",
     "DeckContributor",
-    "EncryptionKey",
     "ConfigProfile",
     "ConfigProfileContributor",
     "ConfigPrompts",