Merge branch 'pinterest:master' into master

Dockerfile

@@ -40,7 +40,9 @@ RUN pip install -r requirements/base.txt \
     pip install -r requirements/${PACKAGE}; \
     done \
     fi \
-    && pip install -r requirements/local.txt || true
+    && if [ -f requirements/local.txt ]; then \
+    pip install -r requirements/local.txt; \
+    fi
 
 COPY package.json yarn.lock ./
 RUN yarn install --pure-lockfile

docs_website/docs/changelog/breaking_change.mdx

@@ -7,6 +7,13 @@ slug: /changelog
 
 Here are the list of breaking changes that you should be aware of when updating Querybook:
 
+## v3.31.0
+
+Upgraded langchain to [0.1.6](https://blog.langchain.dev/langchain-v0-1-0/).
+
+-   Some langchain packages are imported from different paths, e.g. `PromptTemplate` is now from `langchain.prompts`
+-   Removed `StreamingWebsocketCallbackHandler` to adopt the new streaming approach.
+
 ## v3.29.0
 
 Made below changes for `S3BaseExporter` (csv table uploader feature):

package.json

@@ -1,6 +1,6 @@
 {
     "name": "querybook",
-    "version": "3.29.0",
+    "version": "3.31.2",
     "description": "A Big Data Webapp",
     "private": true,
     "scripts": {
@@ -50,6 +50,7 @@
         "core-js": "^3.19.1",
         "cron-parser": "^4.7.0",
         "dagre": "^0.8.5",
+        "dompurify": "^3.0.9",
         "draft-js": "0.11.7",
         "draft-js-export-html": "^1.4.1",
         "draft-js-import-html": "^1.4.1",

querybook/config/querybook_default_config.yaml

@@ -4,7 +4,8 @@ FLASK_SECRET_KEY: ~
 # Url of the querybook site, used for auth callback and notifications
 PUBLIC_URL: ''
 # Use this config to set cache policy of flask, see https://pythonhosted.org/Flask-Cache/ for details
-FLASK_CACHE_CONFIG: '{"CACHE_TYPE": "simple"}'
+FLASK_CACHE_CONFIG:
+    CACHE_TYPE: 'simple'
 
 # --------------- Celery ---------------
 REDIS_URL: ~
@@ -92,6 +93,7 @@ AI_ASSISTANT_CONFIG:
         model_args:
             model_name: ~
             temperature: ~
+            streaming: ~
         reserved_tokens: ~
 
 EMBEDDINGS_PROVIDER: ~

querybook/migrations/versions/299e24dcfd29_allow_unique_table_ownership_per_type.py

@@ -0,0 +1,87 @@
+"""allow unique_table_ownership per type
+
+Revision ID: 299e24dcfd29
+Revises: c00f08f16065
+Create Date: 2024-01-31 14:39:13.601013
+
+"""
+
+from alembic import op
+import sqlalchemy as sa
+from sqlalchemy.dialects import mysql
+
+# revision identifiers, used by Alembic.
+revision = "299e24dcfd29"
+down_revision = "c00f08f16065"
+branch_labels = None
+depends_on = None
+
+
+def upgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    ownership_table = sa.Table(
+        "data_table_ownership", sa.MetaData(op.get_bind()), autoload=True
+    )
+
+    for constraint in ownership_table.foreign_key_constraints:
+        op.drop_constraint(constraint.name, "data_table_ownership", type_="foreignkey")
+
+    op.drop_constraint("unique_table_ownership", "data_table_ownership", type_="unique")
+
+    op.create_foreign_key(
+        "data_table_ownership_data_table_id_fk",
+        "data_table_ownership",
+        "data_table",
+        ["data_table_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.create_foreign_key(
+        "data_table_ownership_uid_fk",
+        "data_table_ownership",
+        "user",
+        ["uid"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.create_unique_constraint(
+        "unique_table_ownership",
+        "data_table_ownership",
+        ["data_table_id", "uid", "type"],
+    )
+    # ### end Alembic commands ###
+
+
+def downgrade():
+    # ### commands auto generated by Alembic - please adjust! ###
+    ownership_table = sa.Table(
+        "data_table_ownership", sa.MetaData(op.get_bind()), autoload=True
+    )
+
+    for constraint in ownership_table.foreign_key_constraints:
+        op.drop_constraint(constraint.name, "data_table_ownership", type_="foreignkey")
+
+    op.drop_constraint("unique_table_ownership", "data_table_ownership", type_="unique")
+
+    op.create_foreign_key(
+        "data_table_ownership_data_table_id_fk",
+        "data_table_ownership",
+        "data_table",
+        ["data_table_id"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.create_foreign_key(
+        "data_table_ownership_uid_fk",
+        "data_table_ownership",
+        "user",
+        ["uid"],
+        ["id"],
+        ondelete="CASCADE",
+    )
+    op.create_unique_constraint(
+        "unique_table_ownership",
+        "data_table_ownership",
+        ["data_table_id", "uid"],
+    )
+    # ### end Alembic commands ###

querybook/server/env.py

@@ -14,6 +14,17 @@ class MissingConfigException(Exception):
     pass
 
 
+def is_json(value):
+    """
+    Check if a given value is a valid JSON serialized dict or list
+    """
+    try:
+        parsed = json.loads(value)
+        return isinstance(parsed, (dict, list))
+    except ValueError:
+        return False
+
+
 def get_env_config(name, optional=True):
     found = True
     val = None
@@ -32,6 +43,9 @@ def get_env_config(name, optional=True):
         raise MissingConfigException(
             "{} is required to start the process.".format(name)
         )
+    # Check for string-serialized JSON dicts/lists
+    if isinstance(val, str) and is_json(val):
+        val = json.loads(val)
     return val
 
 
@@ -40,7 +54,8 @@ class QuerybookSettings(object):
     PRODUCTION = os.environ.get("production", "false") == "true"
     PUBLIC_URL = get_env_config("PUBLIC_URL")
     FLASK_SECRET_KEY = get_env_config("FLASK_SECRET_KEY", optional=False)
-    FLASK_CACHE_CONFIG = json.loads(get_env_config("FLASK_CACHE_CONFIG"))
+    FLASK_CACHE_CONFIG = get_env_config("FLASK_CACHE_CONFIG")
+
     # Celery
     REDIS_URL = get_env_config("REDIS_URL", optional=False)
 
@@ -117,7 +132,7 @@ class QuerybookSettings(object):
 
     DB_MAX_UPLOAD_SIZE = int(get_env_config("DB_MAX_UPLOAD_SIZE"))
 
-    GOOGLE_CREDS = json.loads(get_env_config("GOOGLE_CREDS") or "null")
+    GOOGLE_CREDS = get_env_config("GOOGLE_CREDS")
 
     # Logging
     LOG_LOCATION = get_env_config("LOG_LOCATION")

querybook/server/lib/ai_assistant/ai_socket.py

@@ -26,12 +26,6 @@ def _send(self, event_type, payload: dict = None):
     def send_data(self, data: dict):
         self._send("data", data)
 
-    def send_delta_data(self, data: str):
-        self._send("delta_data", data)
-
-    def send_delta_end(self):
-        self._send("delta_end")
-
     def send_tables_for_sql_gen(self, data: list[str]):
         self._send("tables", data)
 

querybook/server/lib/ai_assistant/assistants/openai_assistant.py

@@ -1,7 +1,6 @@
 import openai
 import tiktoken
-from langchain.callbacks.manager import CallbackManager
-from langchain.chat_models import ChatOpenAI
+from langchain_openai import ChatOpenAI
 
 from lib.ai_assistant.base_ai_assistant import BaseAIAssistant
 from lib.logger import get_logger
@@ -46,19 +45,12 @@ def _get_token_count(self, ai_command: str, prompt: str) -> int:
         return len(encoding.encode(prompt))
 
     def _get_error_msg(self, error) -> str:
-        if isinstance(error, openai.error.AuthenticationError):
+        if isinstance(error, openai.AuthenticationError):
             return "Invalid OpenAI API key"
 
         return super()._get_error_msg(error)
 
-    def _get_llm(self, ai_command: str, prompt_length: int, callback_handler=None):
+    def _get_llm(self, ai_command: str, prompt_length: int):
         config = self._get_llm_config(ai_command)
-        if not callback_handler:
-            # non-streaming
-            return ChatOpenAI(**config)
 
-        return ChatOpenAI(
-            **config,
-            streaming=True,
-            callback_manager=CallbackManager([callback_handler])
-        )
+        return ChatOpenAI(**config)