quillsql 2.1.6__py3-none-any.whl → 2.2.1__py3-none-any.whl

quillsql/core.py

@@ -3,15 +3,27 @@ from dotenv import load_dotenv
 
 import requests
 import redis
+from .utils import Filter, convert_custom_filter
+import json
+from enum import Enum
 
 
 from quillsql.db.cached_connection import CachedConnection
-from quillsql.db.db_helper import get_db_credentials, get_schema_column_info_by_db, get_schema_tables_by_db
+from quillsql.db.db_helper import (
+    get_db_credentials,
+    get_schema_column_info_by_db,
+    get_schema_tables_by_db,
+)
 from quillsql.utils.schema_conversion import convert_type_to_postgres
 from quillsql.utils.run_query_processes import (
     array_to_map,
     remove_fields,
 )
+from quillsql.utils.tenants import extract_tenant_ids
+from quillsql.utils.pivot_template import (
+    parse_distinct_values,
+    hydrate_pivot_template,
+)
 
 load_dotenv()
 
@@ -20,8 +32,11 @@ DEV_HOST = "http://localhost:8080"
 PROD_HOST = "https://quill-344421.uc.r.appspot.com"
 HOST = DEV_HOST if ENV == "development" else PROD_HOST
 
+SINGLE_TENANT = "QUILL_SINGLE_TENANT"
+ALL_TENANTS = "QUILL_ALL_TENANTS"
+FLAG_TASKS = {'dashboard', 'report', 'item', 'report-info', 'filter-options'}
 
-## Quill - Fullstack API Platform for Dashboards and Reporting.
+# Quill - Fullstack API Platform for Dashboards and Reporting.
 class Quill:
     def __init__(
         self,
@@ -32,13 +47,27 @@ class Quill:
         metadataServerURL=None,
         cache=None,
     ):
+        if private_key is None:
+            raise ValueError("Private key is required")
+        if database_type is None:
+            raise ValueError("Database type is required")
+        if database_connection_string is None and database_config is None:
+            raise ValueError("You must provide either DatabaseConnectionString or DatabaseConfig")
+
         # Handles both dsn-style connection strings (eg. "dbname=test password=secret" )
         # as well as url-style connection strings (eg. "postgres://foo@db.com")
-        self.baseUrl = metadataServerURL if metadataServerURL != None else HOST
-        if database_connection_string != None:
-          self.target_connection = CachedConnection(database_type, get_db_credentials(database_type, database_connection_string), cache, True)
+        self.baseUrl = metadataServerURL if metadataServerURL is not None else HOST
+        if database_connection_string is not None:
+            self.target_connection = CachedConnection(
+                database_type,
+                get_db_credentials(database_type, database_connection_string),
+                cache,
+                True,
+            )
         else:
-          self.target_connection = CachedConnection(database_type, database_config, cache, False)
+            self.target_connection = CachedConnection(
+                database_type, database_config, cache, False
+            )
         self.private_key = private_key
 
     def get_cache(self, cache_config):
@@ -52,99 +81,434 @@ class Quill:
             )
         return None
 
-    def query(self, org_id, data):
-        metadata = data.get("metadata")
+    def query(
+        self,
+        tenants,
+        metadata,
+        flags=None,
+        filters: list[Filter] = None,
+        admin_enabled: bool = None,
+    ):
+        if not tenants:
+            raise ValueError("You may not pass an empty tenants array.")
+        if flags and not flags:
+            raise ValueError("You may not pass an empty flags array.")
+
+        responseMetadata = {}
         if not metadata:
-            return {"error": "400", "errorMessage": "Missing metadata."}
+            return {"error": "Missing metadata.", "status": "error", "data": {}}
 
         task = metadata.get("task")
         if not task:
-            return {"error": "400", "errorMessage": "Missing task."}
+            return {"error": "Missing task.", "status": "error", "data": {}}
 
         try:
-            pre_query_results = self.run_queries(
-                metadata.get("preQueries"),
-                self.target_connection.database_type,
-                metadata.get('databaseType'),
-                metadata,
-                metadata.get("runQueryConfig"),
-            )  # used by the view task to get non-sensitive data
-            if metadata.get("runQueryConfig") and metadata.get("runQueryConfig").get('overridePost'):
+            # Set tenant IDs in the connection
+            self.target_connection.tenant_ids = extract_tenant_ids(tenants)
+
+            # Handle pivot-template task
+            if task == "pivot-template":
+                # Step 1: Get pivot template and queries from server
+                pivot_payload = {
+                    **metadata,
+                    "tenants": tenants,
+                    "flags": flags,
+                    "adminEnabled": admin_enabled,
+                }
+                if filters is not None:
+                    pivot_payload["sdkFilters"] = [
+                        convert_custom_filter(f) for f in filters
+                    ]
+                pivot_template_response = self.post_quill(
+                    "pivot-template",
+                    pivot_payload,
+                )
+
+                if pivot_template_response.get("error"):
+                    return {
+                        "status": "error",
+                        "error": pivot_template_response.get("error"),
+                        "data": pivot_template_response.get("metadata") or {},
+                    }
+
+                template = pivot_template_response.get("metadata", {}).get("template")
+                config = pivot_template_response.get("metadata", {}).get("config")
+                distinct_values_query = pivot_template_response.get("metadata", {}).get("distinctValuesQuery")
+                row_count_query = pivot_template_response.get("metadata", {}).get("rowCountQuery")
+
+                # Step 2: Run the distinct values query to get unique values
+                distinct_values = []
+                if distinct_values_query:
+                    distinct_value_results = self.run_queries(
+                        [distinct_values_query],
+                        self.target_connection.database_type,
+                        metadata.get("databaseType"),
+                        metadata,
+                        None
+                    )
+
+                    # Parse distinct values from database results
+                    distinct_values = parse_distinct_values(
+                        distinct_value_results["queryResults"][0],
+                        config.get("databaseType")
+                    )
+
+                # Step 3: Hydrate the template with the distinct values
+                try:
+                    final_query = hydrate_pivot_template(template, distinct_values, config)
+                except Exception as err:
+                    return {
+                        "status": "error",
+                        "error": f"Failed to hydrate pivot template: {str(err)}",
+                        "data": {},
+                    }
+
+                # Step 4: Run queries - pivot query and optional row count query
+                queries_to_run = [final_query]
+                if row_count_query:
+                    # Hydrate the rowCountQuery with the same distinct values
+                    hydrated_row_count_query = hydrate_pivot_template(
+                        row_count_query,
+                        distinct_values,
+                        config
+                    )
+                    queries_to_run.append(hydrated_row_count_query)
+
+                final_results = self.run_queries(
+                    queries_to_run,
+                    self.target_connection.database_type,
+                    metadata.get("databaseType"),
+                    metadata,
+                    pivot_template_response.get("metadata", {}).get("runQueryConfig")
+                )
+
+                responseMetadata = pivot_template_response.get("metadata") or {}
+                # Set rows and fields from first query result (the pivot query)
+                if final_results.get("queryResults") and len(final_results["queryResults"]) >= 1:
+                    query_results = final_results["queryResults"][0]
+                    if query_results.get("rows"):
+                        responseMetadata["rows"] = query_results["rows"]
+                    if query_results.get("fields"):
+                        responseMetadata["fields"] = query_results["fields"]
+
+                # Remove internal SDK fields before returning to frontend
+                if "template" in responseMetadata:
+                    del responseMetadata["template"]
+                if "distinctValuesQuery" in responseMetadata:
+                    del responseMetadata["distinctValuesQuery"]
+                if "rowCountQuery" in responseMetadata:
+                    del responseMetadata["rowCountQuery"]
+
                 return {
-                    "data": pre_query_results ,
-                    "status": "success"
+                    "data": responseMetadata,
+                    "queries": final_results,
+                    "status": "success",
                 }
+
+            # Handle tenant flags synthesis
+            tenant_flags = None
+            if (task in FLAG_TASKS and 
+                tenants[0] != ALL_TENANTS and 
+                tenants[0] != SINGLE_TENANT
+            ):
+                
+                response = self.post_quill(
+                    'tenant-mapped-flags',
+                    {
+                        'reportId': metadata.get('reportId') or metadata.get('dashboardItemId'),
+                        'clientId': metadata.get('clientId'),
+                        'dashboardName': metadata.get('name'),
+                        'tenants': tenants,
+                        'flags': flags,
+                        'adminEnabled': admin_enabled,
+                    },
+                )
+                
+                if response.get('error'):
+                    return {
+                        'status': 'error',
+                        'error': response.get('error'),
+                        'data': response.get('metadata') or {},
+                    }
+                
+                flag_query_results = self.run_queries(
+                    response.get('queries'),
+                    self.target_connection.database_type,
+                )
+                
+                tenant_flags = [
+                    {
+                        'tenantField': tenant_field,
+                        'flags': list(set(row['quill_flag'] for row in query_result['rows']))
+                    }
+                    for tenant_field, query_result in zip(
+                        response['metadata']['queryOrder'],
+                        flag_query_results['queryResults']
+                    )
+                ]
+            elif tenants[0] == SINGLE_TENANT and flags:
+                if flags and isinstance(flags[0], dict):
+                    tenant_flags = [{'tenantField': SINGLE_TENANT, 'flags': flags}]
+                else:
+                    tenant_flags = flags
+
+            pre_query_results = (
+                self.run_queries(
+                    metadata.get("preQueries"),
+                    self.target_connection.database_type,
+                    metadata.get("databaseType"),
+                    metadata,
+                    metadata.get("runQueryConfig"),
+                )
+                if metadata.get("preQueries")
+                else {}
+            )
+
+            if metadata.get("runQueryConfig") and metadata.get("runQueryConfig").get(
+                "overridePost"
+            ):
+                return {"data": pre_query_results, "status": "success"}
             view_query = None
-            if (metadata.get("preQueries")):
+            if metadata.get("preQueries"):
                 view_query = metadata.get("preQueries")[0]
+            pre_query_columns = (
+                pre_query_results.get("columns")
+                if metadata.get("runQueryConfig")
+                and metadata.get("runQueryConfig").get("getColumns")
+                else None
+            )
             payload = {
                 **metadata,
-                "orgId": org_id,
-                "preQueryResults": pre_query_results,
-                "viewQuery": view_query
+                "tenants": tenants,
+                "flags": tenant_flags,
+                "viewQuery": view_query,
+                "preQueryResultsColumns": pre_query_columns,
+                "adminEnabled": admin_enabled,
             }
+            if filters is not None:
+                payload["sdkFilters"] = [convert_custom_filter(f) for f in filters]
             quill_results = self.post_quill(metadata.get("task"), payload)
             if quill_results.get("error"):
-                return {"error": quill_results.get("error"), "status": "error"}
-            # If there is no metedata in the quill results, create one
+                responseMetadata = quill_results.get("metadata")
+                response = {
+                    "error": quill_results.get("error"),
+                    "status": "error",
+                    "data": {},
+                }
+                if responseMetadata:
+                    response["data"] = responseMetadata
+                return response
+
+            # If there is no metadata in the quill results, create one
             if not quill_results.get("metadata"):
                 quill_results["metadata"] = {}
             metadata = quill_results.get("metadata")
-            final_query_results = self.run_queries(
-                quill_results.get("queries"),  self.target_connection.database_type, metadata.get('databaseType'), metadata, metadata.get("runQueryConfig")
+            responseMetadata = metadata
+            results = self.run_queries(
+                quill_results.get("queries"),
+                self.target_connection.database_type,
+                metadata.get("databaseType"),
+                metadata,
+                metadata.get("runQueryConfig"),
             )
-            # Quick hack to make the sdk work with the Frontend
-            if len(final_query_results.get("queryResults")) == 1:
-                query_result = final_query_results.get("queryResults")[0]
+
+            should_wrap_results = isinstance(results, list) or not results
+            if should_wrap_results:
+                normalized_results = {
+                    "queryResults": results if isinstance(results, list) else []
+                }
+            else:
+                normalized_results = results
+
+            if (
+                should_wrap_results
+                and not normalized_results.get("queryResults")
+                and quill_results.get("queries")
+            ):
+                normalized_results["queryResults"] = (
+                    normalized_results.get("queryResults") or []
+                )
+
+            if (
+                normalized_results.get("mapped_array")
+                and metadata.get("runQueryConfig", {}).get("arrayToMap")
+            ):
+                array_to_map = metadata["runQueryConfig"]["arrayToMap"]
+                for array, index in zip(
+                    normalized_results["mapped_array"],
+                    range(len(normalized_results["mapped_array"])),
+                ):
+                    responseMetadata[array_to_map["arrayName"]][index][array_to_map["field"]] = array
+                del normalized_results["mapped_array"]
+
+            query_results_list = normalized_results.get("queryResults") or []
+            if len(query_results_list) == 1:
+                query_result = query_results_list[0]
                 quill_results["metadata"]["rows"] = query_result.get("rows")
                 quill_results["metadata"]["fields"] = query_result.get("fields")
             return {
                 "data": quill_results.get("metadata"),
-                "queries": final_query_results,
+                "queries": normalized_results,
                 "status": "success",
             }
 
         except Exception as err:
-            return {"error": str(err), "status": "error"}
+            if task == "update-view":
+                self.post_quill(
+                    "set-broken-view",
+                    {
+                        "table": metadata.get("name"),
+                        "clientId": metadata.get("clientId"),
+                        "error": str(err),
+                        "adminEnabled": admin_enabled,
+                    },
+                )
+            return {
+                "error": str(err).splitlines()[0],
+                "status": "error",
+                "data": responseMetadata,
+            }
+
+    def apply_limit(self, query, limit):
+        # Simple logic: if query already has a limit, don't add another
+        if getattr(self.target_connection, 'database_type', '').lower() == 'mssql':
+            import re
+            if re.search(r'SELECT TOP \\d+', query, re.IGNORECASE):
+                return query
+            return re.sub(r'select', f'SELECT TOP {limit}', query, flags=re.IGNORECASE)
+        else:
+            if 'limit ' in query.lower():
+                return query
+            return f"{query.rstrip(';')} limit {limit}"
 
-    def run_queries(self, queries, pkDatabaseType, databaseType, metadata=None, runQueryConfig=None):
+    def run_queries(
+        self, queries, pkDatabaseType, databaseType=None, metadata=None, runQueryConfig=None
+    ):
         results = {}
         if not queries:
             return {"queryResults": []}
         if databaseType and databaseType.lower() != pkDatabaseType.lower():
             return {"dbMismatched": True, "backendDatabaseType": pkDatabaseType}
-        if runQueryConfig and runQueryConfig.get("getColumnsForSchema"):
-            return {"queryResults": []}
         if runQueryConfig and runQueryConfig.get("arrayToMap"):
-            array_to_map(
-                queries, runQueryConfig.get("arrayToMap"), metadata, self.target_connection
+            mapped_array = array_to_map(
+                queries,
+                runQueryConfig.get("arrayToMap"),
+                metadata,
+                self.target_connection,
             )
-            return {"queryResults": []}
+
+            return {"queryResults": [], "mapped_array": mapped_array}
         elif runQueryConfig and runQueryConfig.get("getColumns"):
-            query_results = self.target_connection.query(queries[0].strip().rstrip(";") + ' limit 1')
-            results["columns"] = [{'name': result['name'], 'dataTypeID': convert_type_to_postgres(result['dataTypeID'])} for result in query_results['fields']]
+            query_results = self.target_connection.query(
+                queries[0].strip().rstrip(";") + " limit 1000"
+            )
+            results["columns"] = [
+                {
+                    "fieldType": convert_type_to_postgres(result["dataTypeID"]),
+                    "name": result["name"],
+                    "displayName": result["name"],
+                    "isVisible": True,
+                    "field": result["name"],
+                }
+                for result in query_results["fields"]
+            ]
+        elif runQueryConfig and runQueryConfig.get("getColumnsForSchema"):
+            query_results = []
+            for table in queries:
+                if not table.get("viewQuery") or (
+                    not table.get("isSelectStar") and not table.get("customFieldInfo")
+                ):
+                    query_results.append(table)
+                    continue
+
+                limit = ""
+                if runQueryConfig.get("limitBy"):
+                    limit = f" limit {runQueryConfig.get('limitBy')}"
+
+                try:
+                    query_result = self.target_connection.query(
+                        f"{table['viewQuery'].strip().rstrip(';')} {limit}"
+                    )
+                    columns = [
+                        {
+                            "fieldType": convert_type_to_postgres(field["dataTypeID"]),
+                            "name": field["name"],
+                            "displayName": field["name"],
+                            "isVisible": True,
+                            "field": field["name"],
+                        }
+                        for field in query_result["fields"]
+                    ]
+                    query_results.append(
+                        {**table, "columns": columns, "rows": query_result["rows"]}
+                    )
+                except Exception as e:
+                    query_results.append(
+                        {**table, "error": f"Error fetching columns {e}"}
+                    )
+
+            results["queryResults"] = query_results
+            if runQueryConfig.get("fieldsToRemove"):
+                results["queryResults"] = [
+                    {
+                        **table,
+                        "columns": [
+                            column
+                            for column in table.get("columns", [])
+                            if column["name"] not in runQueryConfig["fieldsToRemove"]
+                        ],
+                    }
+                    for table in query_results
+                ]
+            return results
         elif runQueryConfig and runQueryConfig.get("getTables"):
             tables = get_schema_tables_by_db(
-              self.target_connection.database_type,
-              self.target_connection.connection,
-              runQueryConfig['schemaNames']
+                self.target_connection.database_type,
+                self.target_connection.connection,
+                runQueryConfig["schemaNames"],
             )
             schema = get_schema_column_info_by_db(
-              self.target_connection.database_type,
-              self.target_connection.connection,
-              runQueryConfig['schemaNames'],
-              tables
+                self.target_connection.database_type,
+                self.target_connection.connection,
+                runQueryConfig["schemaNames"],
+                tables,
             )
             results["queryResults"] = schema
+        elif runQueryConfig and runQueryConfig.get("runIndividualQueries"):
+            # so that one query doesn't fail the whole thing
+            # the only reason this isn't the default behavior is for backwards compatibility
+            query_results = []
+            for query in queries:
+                try:
+                    run_query = query
+                    if runQueryConfig.get("limitBy"):
+                        run_query = self.apply_limit(query, runQueryConfig["limitBy"])
+                    query_result = self.target_connection.query(run_query)
+                    query_results.append(query_result)
+                except Exception as e:
+                    query_results.append({
+                        "query": query,
+                        "error": str(e),
+                    })
+            results["queryResults"] = query_results
         else:
             if runQueryConfig and runQueryConfig.get("limitThousand"):
-                queries = [query.strip().rstrip(";") + " limit 1000" for query in queries]
+                queries = [
+                    query.strip().rstrip(";") + " limit 1000" for query in queries
+                ]
             elif runQueryConfig and runQueryConfig.get("limitBy"):
-                queries = [query.strip().rstrip(";") + " limit " + runQueryConfig.get("limitBy") for query in queries]
+                queries = [
+                    query.strip().rstrip(";")
+                    + f" limit {runQueryConfig.get('limitBy')}"
+                    for query in queries
+                ]
             query_results = [self.target_connection.query(query) for query in queries]
             results["queryResults"] = query_results
             if runQueryConfig and runQueryConfig.get("fieldsToRemove"):
-                results["queryResults"] = [ remove_fields(query_result, runQueryConfig.get("fieldsToRemove")) for query_result in results["queryResults"]]
+                results["queryResults"] = [
+                    remove_fields(query_result, runQueryConfig.get("fieldsToRemove"))
+                    for query_result in results["queryResults"]
+                ]
             if runQueryConfig and runQueryConfig.get("convertDatatypes"):
                 for query_result in results["queryResults"]:
                     query_result["fields"] = [
@@ -160,9 +524,18 @@ class Quill:
                     ]
 
         return results
-
+    
     def post_quill(self, path, payload):
+        # Custom JSON Encoder to handle Enums
+        class EnumEncoder(json.JSONEncoder):
+            def default(self, obj):
+                if isinstance(obj, Enum):
+                    return obj.value  # Convert enum to its value (string in this case)
+                return super().default(obj)
+        
         url = f"{self.baseUrl}/sdk/{path}"
-        headers = {"Authorization": f"Bearer {self.private_key}"}
-        response = requests.post(url, json=payload, headers=headers)
+        # Set content type to application/json
+        headers = {"Authorization": f"Bearer {self.private_key}", "Content-Type": "application/json"}
+        encoded = json.dumps(payload, cls=EnumEncoder)
+        response = requests.post(url, data=encoded, headers=headers)
         return response.json()

quillsql/db/__init__.py

@@ -1,3 +1,3 @@
 # __init__.py
 
-from .cached_connection import CachedConnection
+from .cached_connection import CachedConnection