refactor: unify duckdb local connection flow for uploads and warehouse

Author: nextchamp-saqib

insights/api/__init__.py

@@ -1,18 +1,14 @@
 # Copyright (c) 2022, Frappe Technologies Pvt. Ltd. and contributors
 # For license information, please see license.txt
 
-import os
-
 import frappe
 from frappe.defaults import get_user_default, set_user_default
 from frappe.handler import is_valid_http_method, is_whitelisted
 from frappe.monitor import add_data_to_monitor
 
 from insights.api.shared import is_public
 from insights.decorators import insights_whitelist, validate_type
-from insights.insights.doctype.insights_data_source_v3.connectors.duckdb import (
-    get_duckdb_connection,
-)
+from insights.insights.doctype.insights_data_source_v3.connectors.duckdb import get_duckdb_connection
 from insights.insights.doctype.insights_data_source_v3.ibis_utils import (
     get_columns_from_schema,
 )
@@ -126,30 +122,23 @@ def get_file_data(filename: str):
 
     create_uploads_if_not_exists()
     ds = frappe.get_doc("Insights Data Source v3", "uploads")
-    private_folder = frappe.utils.get_files_path(is_private=1)
-    private_folder = os.path.realpath(private_folder)
-    db = get_duckdb_connection(ds, read_only=True, allowed_dir=private_folder)
-    try:
-        if ext in ["xlsx"]:
-            table = db.read_xlsx(file_path)
-        elif ext in ["json", "jsonl"]:
-            table = db.read_json(file_path)
-        else:
-            table = db.read_csv(file_path, table_name=file_name)
+    db = get_duckdb_connection(ds, read_only=False, allow_private_files=True)
 
+    try:
+        table = _read_uploaded_table(db, file_path, ext)
         columns = get_columns_from_schema(table.schema())
         rows = table.head(50).execute().fillna("").to_dict(orient="records")
         row_count = table.count().execute()
+
+        return {
+            "tablename": file_name,
+            "rows": rows,
+            "columns": columns,
+            "total_rows": int(row_count),
+        }
     finally:
         db.disconnect()
 
-    return {
-        "tablename": file_name,
-        "rows": rows,
-        "columns": columns,
-        "total_rows": int(row_count),
-    }
-
 
 @insights_whitelist()
 @validate_type
@@ -162,33 +151,46 @@ def import_csv_data(filename: str, tablename: str = ""):
 
     create_uploads_if_not_exists()
     ds = frappe.get_doc("Insights Data Source v3", "uploads")
-    private_folder = os.path.realpath(frappe.utils.get_files_path(is_private=1))
+    db = get_duckdb_connection(ds, read_only=False, allow_private_files=True)
 
-    db = get_duckdb_connection(ds, read_only=False, allowed_dir=private_folder)
     try:
-        if ext in ["xlsx"]:
-            table = db.read_xlsx(file_path)
-        elif ext in ["json", "jsonl"]:
-            table = db.read_json(file_path)
-        else:
-            table = db.read_csv(file_path, table_name=table_name)
+        table = _read_uploaded_table(db, file_path, ext)
         db.create_table(table_name, table, overwrite=True)
+    except frappe.ValidationError:
+        raise
     except Exception as e:
         frappe.log_error(e)
-        if ext in ["xlsx"]:
+        frappe.throw("Failed to import uploaded file data into Insights uploads table. Please try again.")
+    finally:
+        db.disconnect()
+
+    InsightsTablev3.bulk_create(ds.name, [table_name])
+
+
+def _read_uploaded_table(db, file_path: str, ext: str):
+    try:
+        if ext == "xlsx":
+            return db.read_xlsx(file_path)
+
+        if ext in ["json", "jsonl"]:
+            return db.read_json(file_path)
+
+        return db.read_csv(file_path)
+
+    except Exception as e:
+        frappe.log_error(e)
+
+        if ext == "xlsx":
             frappe.throw(
                 "Failed to read Excel data from uploaded file. Please ensure the file is a valid Excel format and try again."
             )
-        elif ext in ["json", "jsonl"]:
+
+        if ext in ["json", "jsonl"]:
             frappe.throw(
                 "Failed to read JSON data from uploaded file. Please ensure the file is a valid JSON or JSONL format and try again."
             )
-        else:
-            frappe.throw("Failed to read CSV data from uploaded file. Please try again.")
-    finally:
-        db.disconnect()
 
-    InsightsTablev3.bulk_create(ds.name, [table_name])
+        frappe.throw("Failed to read CSV data from uploaded file. Please try again.")
 
 
 @frappe.whitelist(allow_guest=True)

insights/insights/doctype/insights_data_source_v3/connectors/duckdb.py

@@ -10,36 +10,39 @@
 from frappe.utils import get_files_path
 
 
-def get_duckdb_connection(data_source, read_only=True, allowed_dir=None):
+def get_duckdb_connection(data_source, read_only=True, allowed_dir=None, allow_private_files=False):
     name = data_source.name or frappe.scrub(data_source.title)
     db_name = data_source.database_name
 
     if db_name.startswith("http"):
         return get_http_duckdb_connection(data_source, name, db_name)
 
-    return get_local_duckdb_connection(db_name, read_only=read_only, allowed_dir=allowed_dir)
+    path = os.path.join(get_files_path(is_private=1), f"{db_name}.duckdb")
+    return get_local_duckdb_connection(
+        path, read_only=read_only, allowed_dir=allowed_dir, allow_private_files=allow_private_files
+    )
 
 
-def get_local_duckdb_connection(db_name, read_only=True, allowed_dir=None):
-    private_folder = os.path.realpath(get_files_path(is_private=1))
-    path = os.path.join(private_folder, f"{db_name}.duckdb")
-
+def get_local_duckdb_connection(path, read_only=True, allowed_dir=None, allow_private_files=False):
     if not os.path.exists(path):
         db = ibis.duckdb.connect(path)
         db.disconnect()
 
     db = ibis.duckdb.connect(path, read_only=read_only)
+
+    private_folder = os.path.realpath(get_files_path(is_private=1))
+    private_folder = _escape_sql_path(private_folder)
     db.raw_sql(f"SET home_directory='{private_folder}'")
 
-    if allowed_dir:
-        escaped_dir = allowed_dir.replace("'", "''")
-        # Some environments start with external access already disabled.
-        # Best effort: try enabling it first, then configure directory allowlist.
+    if not read_only and (allowed_dir or allow_private_files):
+        allowed_dir = _escape_sql_path(allowed_dir) if allowed_dir else private_folder
+
         with suppress(Exception):
             db.raw_sql("SET enable_external_access = true")
-        db.raw_sql(f"SET allowed_directories = ['{escaped_dir}']")
 
-    db.raw_sql("SET enable_external_access = false")
+        db.raw_sql(f"SET allowed_directories = ['{allowed_dir}']")
+    else:
+        db.raw_sql("SET enable_external_access = false")
 
     return db
 
@@ -80,3 +83,7 @@ def get_http_secret(data_source, name, db_name):
     except Exception as e:
         frappe.log_error(title="Error creating HTTP Secret for DuckDB", message=str(e))
         return
+
+
+def _escape_sql_path(path: str) -> str:
+    return path.replace("'", "''")

insights/insights/doctype/insights_data_source_v3/data_warehouse.py

@@ -20,6 +20,7 @@
 from ibis.expr.types import Expr
 
 import insights
+from insights.insights.doctype.insights_data_source_v3.connectors.duckdb import get_local_duckdb_connection
 from insights.utils import InsightsDataSourcev3, InsightsTablev3
 
 WAREHOUSE_DB_NAME = "insights"
@@ -39,31 +40,17 @@ def get_db_path(self) -> str:
     def get_connection(self, database: str | None = None, read_only: bool = True) -> DuckDBBackend:
         path = self.get_db_path()
 
-        if not os.path.exists(path):
-            db = ibis.duckdb.connect(path)
-            db.disconnect()
-
-        db = ibis.duckdb.connect(path, read_only=read_only)
-
-        if not read_only:
-            self._configure_temp_directory_access(db)
+        db = get_local_duckdb_connection(
+            path,
+            read_only=read_only,
+            allowed_dir=str(Path(tempfile.gettempdir())) if not read_only else None,
+        )
 
         if database:
             db.raw_sql(f"USE '{database}'")
 
         return db
 
-    def _configure_temp_directory_access(self, db: DuckDBBackend) -> None:
-        tmp_dir = str(Path(tempfile.gettempdir())).replace("'", "''")
-
-        # Some environments start with external access already disabled.
-        # Best effort: try enabling it first, then configure directory allowlist.
-        with suppress(Exception):
-            db.raw_sql("SET enable_external_access = true")
-
-        db.raw_sql(f"SET allowed_directories = ['{tmp_dir}']")
-        db.raw_sql("SET enable_external_access = false")
-
     def create_database(self, database: str):
         with self.get_write_connection() as db:
             with suppress(CatalogException):