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

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

quillsql/db/cached_connection.py

@@ -1,11 +1,10 @@
 import psycopg2
-from psycopg2.extras import RealDictCursor
 import json
 import redis
 from quillsql.error import PgQueryError
 from quillsql.db.db_helper import connect_to_db, run_query_by_db
 
-## The TTL for new cache entries (default: 1h)
+# The TTL for new cache entries (default: 1h)
 DEFAULT_CACHE_TTL = 24 * 60 * 60
 
 
@@ -33,7 +32,9 @@ class CachedConnection:
     def exec_with_reconnect(self, sql):
         reconnect_count = 0
         while True:
-            self.connection = connect_to_db(self.database_type, self.config, self.using_connection_string)
+            self.connection = connect_to_db(
+                self.database_type, self.config, self.using_connection_string
+            )
             try:
                 return run_query_by_db(self.database_type, sql, self.connection)
             except psycopg2.Error as err:
@@ -57,9 +58,9 @@ class CachedConnection:
         cached_result = self.cache.get(key)
         if cached_result:
             cached = json.loads(cached_result)
-            return {"rows": cached, "fields": cached['fields']}
+            return {"rows": cached, "fields": cached["fields"]}
         else:
             new_result = self.exec(sql)
             new_result_string = json.dumps(new_result)
             self.cache.set(key, new_result_string, "EX", DEFAULT_CACHE_TTL)
-            return {"rows": new_result, "fields": new_result['fields']}
+            return {"rows": new_result, "fields": new_result["fields"]}

quillsql/db/db_helper.py

@@ -1,42 +1,61 @@
-from quillsql.db.postgres import format_postgres, connect_to_postgres, get_schema_column_info_postgres, get_tables_by_schema_postgres, run_query_postgres, disconnect_from_postgres
-from quillsql.db.bigquery import format_bigquery_config, connect_to_bigquery, get_schema_column_info_big_query, get_tables_by_schema_big_query, run_query_big_query
+from quillsql.db.postgres import (
+    format_postgres,
+    connect_to_postgres,
+    get_schema_column_info_postgres,
+    get_tables_by_schema_postgres,
+    run_query_postgres,
+    disconnect_from_postgres,
+)
+from quillsql.db.bigquery import (
+    format_bigquery_config,
+    connect_to_bigquery,
+    get_schema_column_info_big_query,
+    get_tables_by_schema_big_query,
+    run_query_big_query,
+)
+
 
 def get_db_credentials(database_type, connection_string):
-  if(database_type.lower() == 'postgresql'):
-    return format_postgres(connection_string)
-  elif(database_type.lower() == 'bigquery'):
-    return format_bigquery_config(connection_string)
-  return {}
+    if database_type.lower() == "postgresql":
+        return format_postgres(connection_string)
+    elif database_type.lower() == "bigquery":
+        return format_bigquery_config(connection_string)
+    return {}
+
 
 def connect_to_db(database_type, config, using_connection_string):
-    if(database_type.lower() == 'postgresql'):
+    if database_type.lower() == "postgresql":
         return connect_to_postgres(config, using_connection_string)
-    elif(database_type.lower() == 'bigquery'):
+    elif database_type.lower() == "bigquery":
         return connect_to_bigquery(config, using_connection_string)
     return None
 
+
 def run_query_by_db(database_type, query, connection):
-    if (database_type.lower() == 'postgresql'):
+    if database_type.lower() == "postgresql":
         return run_query_postgres(query, connection)
-    elif (database_type.lower() == 'bigquery'):
+    elif database_type.lower() == "bigquery":
         return run_query_big_query(query, connection)
     return None
 
+
 def disconnect_from_db(database_type, connection):
-    if (database_type.lower() == 'postgresql'):
+    if database_type.lower() == "postgresql":
         return disconnect_from_postgres(connection)
     return None
 
+
 def get_schema_tables_by_db(database_type, connection, schema_name):
-    if (database_type.lower() == 'postgresql'):
+    if database_type.lower() == "postgresql":
         return get_tables_by_schema_postgres(connection, schema_name)
-    elif (database_type.lower() == 'bigquery'):
+    elif database_type.lower() == "bigquery":
         return get_tables_by_schema_big_query(connection, schema_name)
     return None
 
+
 def get_schema_column_info_by_db(database_type, connection, schema_name, table_names):
-    if (database_type.lower() == 'postgresql'):
+    if database_type.lower() == "postgresql":
         return get_schema_column_info_postgres(connection, schema_name, table_names)
-    elif (database_type.lower() == 'bigquery'):
+    elif database_type.lower() == "bigquery":
         return get_schema_column_info_big_query(connection, schema_name, table_names)
-    return None
+    return None

quillsql/db/postgres.py

@@ -1,72 +1,127 @@
+import os
 import psycopg2
 from psycopg2.extensions import make_dsn
 from quillsql.assets.pgtypes import PG_TYPES
 
-def format_postgres(connection_string ):
-    to_dsn = lambda conn: make_dsn(conn) if "://" in conn else conn
+
+def format_postgres(connection_string):
+    def to_dsn(conn):
+        return make_dsn(conn) if "://" in conn else conn
+
     return to_dsn(connection_string)
 
+
 def connect_to_postgres(config, usingConnectionString):
+    os.environ['PGGSSENCMODE'] = 'disable' # https://github.com/psycopg/psycopg2/issues/1084
     if usingConnectionString:
-      return psycopg2.connect(config)
+        return psycopg2.connect(config)
     else:
-      return psycopg2.connect(
-        database=config['dbname'],
-        user=config['user'],
-        password=config['password'],
-        host=config['host'],
-        port=config['port']
-      )
+        return psycopg2.connect(
+            database=config["dbname"],
+            user=config["user"],
+            password=config["password"],
+            host=config["host"],
+            port=config["port"],
+        )
+
 
 def run_query_postgres(query, connection):
     cursor = connection.cursor()
     cursor.execute(query)
     result = cursor.fetchall()
-    fields = [
-            {"name": desc[0], "dataTypeID": desc[1]} for desc in cursor.description
-        ]
+    fields = [{"name": desc[0], "dataTypeID": desc[1]} for desc in cursor.description]
     cursor.close()
-    rows_dict = [dict(zip([field['name'] for field in fields], row)) for row in result]
+    rows_dict = [dict(zip([field["name"] for field in fields], row)) for row in result]
     return {"rows": rows_dict, "fields": fields}
 
+
 def disconnect_from_postgres(connection):
     connection.close()
     return
 
+
 # getTablesBySchemaPostgres
+
+
 def get_tables_by_schema_postgres(connection, schema_names):
     all_tables = []
     for schema_name in schema_names:
-      query = f"SELECT table_name, table_schema FROM information_schema.tables WHERE table_schema = '{schema_name}'"
-      results = run_query_postgres(query, connection)
-      for row in results['rows']:
-        cur_table = {}
-        cur_table['table_name'] = row['table_name']
-        cur_table['schema_name'] = row['table_schema']
-        all_tables.append(cur_table)
+        query = f"""
+            SELECT table_name, table_schema
+            FROM information_schema.tables
+            WHERE table_schema = '{schema_name}'
+
+            UNION
+
+            SELECT c.relname as table_name, n.nspname as table_schema
+            FROM pg_class c
+            JOIN pg_namespace n ON c.relnamespace = n.oid
+            WHERE n.nspname = '{schema_name}' 
+            AND c.relkind = 'm';
+            """
+        results = run_query_postgres(query, connection)
+        for row in results["rows"]:
+            cur_table = {}
+            cur_table["table_name"] = row["table_name"]
+            cur_table["schema_name"] = row["table_schema"]
+            all_tables.append(cur_table)
     return all_tables
 
-# getSchemaColumnInfoPostgress
+
+# getSchemaColumnInfoPostgres
+
+
 def get_schema_column_info_postgres(connection, schema_name, table_names):
     all_columns = []
     for table_name in table_names:
-        query = f"SELECT column_name, udt_name FROM information_schema.columns WHERE table_schema = '{table_name['schema_name']}' AND table_name = '{table_name['table_name']}' ORDER BY ordinal_position"
+        query = f"""
+            SELECT column_name as "column_name", udt_name as "field_type", ordinal_position as "sort_number"
+            FROM information_schema.columns
+            WHERE table_schema = '{table_name['schema_name']}' 
+            AND table_name = '{table_name['table_name']}'
+            
+            UNION
+
+            SELECT a.attname as "column_name", t.typname as "field_type", a.attnum as "sort_number"
+            FROM pg_attribute a
+            JOIN pg_class c ON a.attrelid = c.oid
+            JOIN pg_namespace n ON c.relnamespace = n.oid
+            JOIN pg_type t ON a.atttypid = t.oid
+            WHERE n.nspname = '{table_name['schema_name']}'
+                AND c.relname = '{table_name['table_name']}'
+                AND c.relkind = 'm'
+                AND a.attnum > 0
+                AND NOT a.attisdropped
+            ORDER BY "sort_number"
+            """
         results = run_query_postgres(query, connection)
         columns = []
-        for row in results['rows']:
-            # Convert row['udt_name'] to postgresql oid
-            pg_type = next((pg_type for pg_type in PG_TYPES if pg_type['typname'] == row['udt_name']), None)
-            if pg_type == None:
+        for row in results["rows"]:
+            pg_type = next(
+                (
+                    pg_type
+                    for pg_type in PG_TYPES
+                    if pg_type["typname"] == row["field_type"]
+                ),
+                None,
+            )
+            if pg_type is None:
                 pg_type = 1043
-            columns.append({
-                'columnName': row['column_name'],
-                'displayName': row['column_name'],
-                'dataTypeID': pg_type['oid'],
-                'fieldType': row['udt_name'],
-            })
-        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
+            columns.append(
+                {
+                    "columnName": row["column_name"],
+                    "displayName": row["column_name"],
+                    "dataTypeID": pg_type["oid"],
+                    "fieldType": row["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

quillsql/error.py

@@ -1,5 +1,5 @@
 class PgQueryError(Exception):
-  def __init__(self, message, query, position):
-    super().__init__(message)
-    self.query = query
-    self.position = position
+    def __init__(self, message, query, position):
+        super().__init__(message)
+        self.query = query
+        self.position = position

quillsql/utils/__init__.py

@@ -1,3 +1,4 @@
 # __init__.py
 
-from .run_query_processes import remove_fields, array_to_map
+from .run_query_processes import remove_fields, array_to_map
+from .filters import Filter, FilterType, FieldType, StringOperator, NumberOperator, NullOperator, DateOperator, convert_custom_filter