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

quillsql/core.py

@@ -3,15 +3,23 @@ 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
 
 load_dotenv()
 
@@ -20,8 +28,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 +43,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,47 +77,126 @@ 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):
+        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'):
-                return {
-                    "data": pre_query_results ,
-                    "status": "success"
-                }
+            # Set tenant IDs in the connection
+            self.target_connection.tenant_ids = extract_tenant_ids(tenants)
+
+            # 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,
+                })
+                
+                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]
             payload = {
                 **metadata,
-                "orgId": org_id,
-                "preQueryResults": pre_query_results,
-                "viewQuery": view_query
+                **pre_query_results,
+                "tenants": tenants,
+                "flags": tenant_flags,
+                "viewQuery": view_query,
             }
+            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")
+            responseMetadata = metadata
             final_query_results = self.run_queries(
-                quill_results.get("queries"),  self.target_connection.database_type, metadata.get('databaseType'), metadata, metadata.get("runQueryConfig")
+                quill_results.get("queries"),
+                self.target_connection.database_type,
+                metadata.get("databaseType"),
+                metadata,
+                metadata.get("runQueryConfig"),
             )
+
+            if final_query_results.get("mapped_array") and metadata.get("runQueryConfig", {}).get("arrayToMap"):
+                array_to_map = metadata["runQueryConfig"]["arrayToMap"]
+                for array, index in zip(final_query_results["mapped_array"], range(len(final_query_results["mapped_array"]))):
+                    responseMetadata[array_to_map["arrayName"]][index][array_to_map["field"]] = array
+                del final_query_results["mapped_array"]
+
             # 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]
@@ -105,46 +209,158 @@ class Quill:
             }
 
         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),
+                })
+            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 +376,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

quillsql/db/bigquery.py

@@ -3,112 +3,146 @@ from google.oauth2 import service_account
 import json
 import re
 
+
 def format_bigquery_config(connection_string):
     # find the start of the json {
-    json_start = connection_string.find('{')
-    if (json_start == -1) :
-      raise Exception("Invalid input string. No JSON data found.");
-    
+    json_start = connection_string.find("{")
+    if json_start == -1:
+        raise Exception("Invalid input string. No JSON data found.")
+
     dataset_name = connection_string[0:json_start].strip()
     json_string = connection_string[json_start:]
     try:
-      service_account = json.loads(json_string)
-      if (not service_account.get('project_id') or not service_account.get('private_key')) :
-        raise Exception( "Invalid service account JSON. Required fields are missing.")
-      
-      return {
-        "dataset_id": dataset_name,
-        "project": service_account.get('project_id'),
-        "credentials": service_account,
-      }
-    except:
-       print('Invalid JSON string')
+        service_account = json.loads(json_string)
+        if not service_account.get("project_id") or not service_account.get(
+            "private_key"
+        ):
+            raise Exception(
+                "Invalid service account JSON. Required fields are missing."
+            )
+
+        return {
+            "dataset_id": dataset_name,
+            "project": service_account.get("project_id"),
+            "credentials": service_account,
+        }
+    except (ValueError, TypeError) as e:
+        print("Invalid JSON string: ", e)
     return connection_string
 
+
 def connect_to_bigquery(config, using_connection_string):
     if using_connection_string:
-      credentials = service_account.Credentials.from_service_account_info(config['credentials'])
+        credentials = service_account.Credentials.from_service_account_info(
+            config["credentials"]
+        )
     else:
-      credentials = service_account.Credentials.from_service_account_file(config['service_account_file_path'])
-    return bigquery.Client(project=config['project'], credentials=credentials)
+        credentials = service_account.Credentials.from_service_account_file(
+            config["service_account_file_path"]
+        )
+    return bigquery.Client(project=config["project"], credentials=credentials)
+
 
 def run_query_big_query(query, connection):
     query_job = connection.query(query)
     result = query_job.result()
     rows = [dict(row) for row in result]
-    fields = [{"name": field.name, "dataTypeID": convert_bigquery_to_postgres(field.field_type)} for field in result.schema]
+    fields = [
+        {
+            "name": field.name,
+            "dataTypeID": convert_bigquery_to_postgres(field.field_type),
+        }
+        for field in result.schema
+    ]
     # TODO CONVERT to postgres types
 
     return {"rows": rows, "fields": fields}
 
+
 def get_tables_by_schema_big_query(connection, schema_names):
     all_table = []
     for schema_name in schema_names:
-      dataset_ref = connection.dataset(schema_name)
-      tables = connection.list_tables(dataset_ref)
-      for table in tables:
-         cur_table = {}
-         cur_table['table_name'] = table.table_id
-         cur_table['schema_name'] = schema_name
-         all_table.append(cur_table)
+        dataset_ref = connection.dataset(schema_name)
+        tables = connection.list_tables(dataset_ref)
+        for table in tables:
+            cur_table = {}
+            cur_table["table_name"] = table.table_id
+            cur_table["schema_name"] = schema_name
+            all_table.append(cur_table)
     return all_table
 
+
 def get_schema_column_info_big_query(connection, schema_name, table_names):
     all_columns = []
     for table_name in table_names:
-        table_ref = connection.dataset(table_name['schema_name']).table(table_name['table_name'])
+        table_ref = connection.dataset(table_name["schema_name"]).table(
+            table_name["table_name"]
+        )
         table = connection.get_table(table_ref)
         columns = []
         for field in table.schema:
-            columns.append({
-                'columnName': field.name,
-                'displayName': field.name,
-                'dataTypeId': convert_bigquery_to_postgres(field.field_type),
-                'fieldType': field.field_type
-            })
-        all_columns.append({
-            'tableName': table_name['schema_name']+'.'+table_name['table_name'],
-            'displayName': table_name['schema_name']+'.'+table_name['table_name'],
-            'columns': columns
-        })
+            columns.append(
+                {
+                    "columnName": field.name,
+                    "displayName": field.name,
+                    "dataTypeId": convert_bigquery_to_postgres(field.field_type),
+                    "fieldType": field.field_type,
+                }
+            )
+        all_columns.append(
+            {
+                "tableName": table_name["schema_name"] + "." + table_name["table_name"],
+                "displayName": table_name["schema_name"]
+                + "."
+                + table_name["table_name"],
+                "columns": columns,
+            }
+        )
     return all_columns
 
+
 def infer_schema_big_query(elem):
-  # compare elem with regex
-  if isinstance(elem, list):
-     return 23
-  if isinstance(elem, object):
-    if re.match(r"/^\d{4}-\d{2}-\d{2}$/", elem.get('value')):
-      return 1082
-    elif re.match(r"/^\d{2}\/\d{2}\/\d{2,4}$/", elem.get('value')):
-      return 1082
-    elif re.match(r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?Z$/", elem.get('value')):
-      return 1184
-    elif re.match(r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?$/", elem.get('value')):
-      return 1114
-    elif re.match(r"/^\d{2}:\d{2}:\d{2}$/", elem.get('value')):
-      return 1083
-  if isinstance(elem, str):
-    if re.match(r"/^\d{4}-\d{2}-\d{2}$/", elem):
-      return 1082
-    elif re.match(r"/^\d{2}\/\d{2}\/\d{2,4}$/", elem):
-      return 1082
-    elif re.match(r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?Z$/", elem):
-      return 1184
-    elif re.match(r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?$/", elem):
-      return 1114
-    elif re.match(r"/^\d{2}:\d{2}:\d{2}$/", elem):
-      return 1083
-    else:
-      return 1043
-  return 1043
+    # compare elem with regex
+    if isinstance(elem, list):
+        return 23
+    if isinstance(elem, object):
+        if re.match(r"/^\d{4}-\d{2}-\d{2}$/", elem.get("value")):
+            return 1082
+        elif re.match(r"/^\d{2}\/\d{2}\/\d{2,4}$/", elem.get("value")):
+            return 1082
+        elif re.match(
+            r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?Z$/", elem.get("value")
+        ):
+            return 1184
+        elif re.match(
+            r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?$/", elem.get("value")
+        ):
+            return 1114
+        elif re.match(r"/^\d{2}:\d{2}:\d{2}$/", elem.get("value")):
+            return 1083
+    if isinstance(elem, str):
+        if re.match(r"/^\d{4}-\d{2}-\d{2}$/", elem):
+            return 1082
+        elif re.match(r"/^\d{2}\/\d{2}\/\d{2,4}$/", elem):
+            return 1082
+        elif re.match(r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?Z$/", elem):
+            return 1184
+        elif re.match(r"/^\d{4}-\d{2}-\d{2}T\d{2}:\d{2}:\d{2}(\.\d+)?$/", elem):
+            return 1114
+        elif re.match(r"/^\d{2}:\d{2}:\d{2}$/", elem):
+            return 1083
+        else:
+            return 1043
+    return 1043
+
 
 def convert_bigquery_to_postgres(value):
-  type_to_oid = {
-    "VARCHAR": 1043,
-    "INTEGER": 23,
-    "FLOAT": 700,
-    "TIMESTAMP": 1114,
-    "DATE": 1082,
-  }
-  return type_to_oid.get(value.upper()) or 1043
+    type_to_oid = {
+        "VARCHAR": 1043,
+        "INTEGER": 23,
+        "FLOAT": 700,
+        "TIMESTAMP": 1114,
+        "DATE": 1082,
+        "BOOL": 16,
+    }
+    return type_to_oid.get(value.upper()) or 1043