quill-sql 0.1.5 → 0.1.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2b78c0e2b82a49c007a604467d4c30ee9d57f0df37b7aa7b316fcdee1249f6a9
-  data.tar.gz: 3c519df47600372ab72b735ce5f0baf7402ebfc629b631b83a895823b1788c21
+  metadata.gz: b3d503d9223878a246289559453c27d45e4393620f9b548433ea5c072fc4c396
+  data.tar.gz: 1820f9f817ab7a29c947555a564e6981b283486cb0cdd0f65396cd01aaaecffd
 SHA512:
-  metadata.gz: f2ae021a1248508de215150b14f8584dd4e2590e09551f1a104a40e42146cc5c77415ec8a481baed6929743f41e8e6aa469594159de2ab2ddc78af34518f4811
-  data.tar.gz: 0e91cf6e674aeb7db748f719d998bcaf3ad2fb7743723efb9588449ae63a77b3a53a1d6b25a7692f8d7ba71da23aac462b3fdc64840a9abfb4d345bbc46503d2
+  metadata.gz: b25af374393d2d9f6bc083f7cb426f631802df7582c0e4bc540c574e54123884e723030ca879dc9c937da21e40e6047be35c9ee59c001ee9bce78339db353ed9
+  data.tar.gz: 7a3251a382f48867d20e7600a4d00f69b95963a435328a4ae3348505ceeb1fe9671834b2715e0f3545f53cee507bad086d87f1acf75150dde218a9005d758d80

data/lib/db/databricks.rb

@@ -0,0 +1,207 @@
+require 'uri'
+require 'json'
+require 'net/http'
+require 'active_support/inflector'
+
+module DatabricksHelper
+  DATABRICKS_DATATYPE_NAME_MAP = {
+    # Numeric Types
+    'BIGINT' => 20,      # int8
+    'BINARY' => 17,      # bytea
+    'BOOLEAN' => 16,     # bool
+    'DATE' => 1082,      # date
+    'DECIMAL' => 1700,   # numeric
+    'DOUBLE' => 701,     # float8
+    'FLOAT' => 701,      # float8
+    'INT' => 23,         # int4
+    'INTERVAL' => 1186,  # interval
+    'VOID' => 2278,      # void
+    'SMALLINT' => 21,    # int2
+    'STRING' => 1043,    # varchar
+    'TIMESTAMP' => 1114, # timestamp
+    'TIMESTAMP_NTZ' => 1114, # timestamp without timezone
+    'TINYINT' => 21,     # mapped to smallint (int2) as postgres has no tinyint
+
+    # Complex Types - mapping to closest PostgreSQL equivalents
+    'ARRAY' => 2277,     # anyarray
+    'MAP' => 3802,       # jsonb
+    'STRUCT' => 3802,    # jsonb
+    'VARIANT' => 3802,   # jsonb
+    'OBJECT' => 3802     # jsonb
+  }.freeze
+
+  class << self
+    def connect_to_databricks(config)
+      warehouse_id = config[:path].split('/').last
+
+      {
+        host: config[:host],
+        warehouse_id: warehouse_id,
+        token: config[:token]
+      }
+    end
+
+    def disconnect_from_databricks(client)
+      # No-op as we're using REST API
+    end
+
+    def run_query_databricks(sql, client)
+      uri = URI("https://#{client[:host]}/api/2.0/sql/statements")
+      http = Net::HTTP.new(uri.host, uri.port)
+      http.use_ssl = true
+
+      request = Net::HTTP::Post.new(uri)
+      request['Authorization'] = "Bearer #{client[:token]}"
+      request['Content-Type'] = 'application/json'
+      request.body = JSON.generate({
+        statement: sql,
+        warehouse_id: client[:warehouse_id],
+        wait_timeout: '50s'
+      })
+
+      response = http.request(request)
+
+      raise "Query failed: #{response.body}" unless response.is_a?(Net::HTTPSuccess)
+
+      result = JSON.parse(response.body)
+      
+      raise "Query state: #{result['status']['state']}" unless result['status']['state'] == 'SUCCEEDED'
+
+      columns = result['manifest']['schema']['columns']
+      column_names = columns.map { |col| col['name'] }
+      column_types = columns.map { |col| col['type_name'] }
+
+      json_like_types = ['ARRAY', 'MAP', 'STRUCT', 'JSON']
+
+
+      fields = columns.map do |column|
+        {
+          name: column['name'],
+          dataTypeID: DATABRICKS_DATATYPE_NAME_MAP[column['type_name']] || 1043 # default to varchar if type not found
+        }
+      end
+
+      # Transform the data array into rows of hashes
+      rows = result['result']['data_array'].map do |row_array|
+        row_hash = column_names.zip(row_array).to_h
+        transformed_row = {}
+      
+        column_names.each_with_index do |key, idx|
+          value = row_hash[key]
+          type = column_types[idx]
+          # parse JSON if the field is a JSON type
+          transformed_value = if value.is_a?(String) && json_like_types.include?(type.upcase)
+            begin
+              parsed = JSON.parse(value)
+              parsed.is_a?(Array) ? parsed.map(&:to_s) : parsed
+            rescue JSON::ParserError
+              value
+            end
+          else
+            value
+          end
+      
+          transformed_row[key] = transformed_value
+        end
+      
+        transformed_row
+      end
+
+      {
+        fields: fields,
+        rows: rows
+      }
+    end
+
+    def get_schemas_databricks(client)
+      sql = <<~SQL
+        SELECT schema_name 
+        FROM system.information_schema.schemata 
+        WHERE schema_name != "information_schema"
+      SQL
+      
+      results = run_query_databricks(sql, client)
+      results[:rows].map { |row| row['schema_name'] }
+    end
+
+    def get_tables_by_schema_databricks(client, schema_names)
+      all_tables = schema_names.flat_map do |schema|
+        sql = <<~SQL
+          SELECT 
+            table_name as tableName,
+            table_schema as schemaName,
+            table_catalog as catalogName
+          FROM system.information_schema.tables
+          WHERE table_schema = '#{schema}'
+        SQL
+        
+        results = run_query_databricks(sql, client)
+        results[:rows].map do |row|
+          {
+            tableName: row['tableName'],
+            schemaName: row['schemaName'],
+            catalogName: row['catalogName']
+          }
+        end
+      end
+
+      all_tables
+    end
+
+    def get_columns_by_table_databricks(client, schema_name, table_name)
+      sql = <<~SQL
+        SELECT 
+          column_name as columnName
+        FROM system.information_schema.columns
+        WHERE table_schema = '#{schema_name}'
+          AND table_name = '#{table_name}'
+        ORDER BY ordinal_position
+      SQL
+      
+      results = run_query_databricks(sql, client)
+      results[:rows].map { |row| row['columnName'] }
+    end
+
+    def get_foreign_keys_databricks(client, schema_name, table_name, primary_key)
+      # Databricks doesn't support foreign keys
+      []
+    end
+
+    def get_schema_column_info_databricks(client, schema_name, table_names)
+      table_names.map do |table|
+        sql = <<~SQL
+          SELECT 
+            column_name as columnName,
+            data_type as dataType
+          FROM system.information_schema.columns
+          WHERE table_schema = '#{table[:schemaName]}'
+            AND table_name = '#{table[:tableName]}'
+          ORDER BY ordinal_position
+        SQL
+
+        results = run_query_databricks(sql, client)
+        {
+          tableName: "#{table[:schemaName]}.#{table[:tableName]}",
+          columns: results[:rows].map do |row|
+            data_type = row['dataType'].split('<')[0].upcase
+            {
+              columnName: row['columnName'],
+              displayName: row['columnName'],
+              fieldType: row['dataType'],
+              dataTypeID: DATABRICKS_DATATYPE_ID_MAP[data_type] || 1043
+            }
+          end
+        }
+      end
+    end
+
+    def format_databricks_config(connection_string)
+      parsed = URI(connection_string)
+      {
+        host: parsed.host,
+        path: parsed.path[1..],  # Remove leading slash
+        token: parsed.password
+      }
+    end
+  end
+end

data/lib/db/db_helper.rb

@@ -2,9 +2,9 @@ require 'json'
 require 'uri'
 require_relative 'clickhouse'
 require_relative 'postgres'
-
+require_relative 'databricks'
 module DatabaseHelper
-  SUPPORTED_DATABASES = ['clickhouse', 'postgresql'].freeze
+  SUPPORTED_DATABASES = ['clickhouse', 'postgresql', 'databricks'].freeze
 
   class QuillQueryResults
     attr_reader :fields, :rows
@@ -23,6 +23,8 @@ module DatabaseHelper
       ClickHouseHelper.format_clickhouse_config(connection_string)
     when 'postgresql'
       PostgresHelper.format_postgres_config(connection_string)
+    when 'databricks'
+      DatabricksHelper.format_databricks_config(connection_string)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end
@@ -34,6 +36,8 @@ module DatabaseHelper
       ClickHouseHelper.connect_to_clickhouse(config)
     when 'postgresql'
       PostgresHelper.connect_to_postgres(config)
+    when 'databricks'
+      DatabricksHelper.connect_to_databricks(config)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end
@@ -57,6 +61,8 @@ module DatabaseHelper
       ClickHouseHelper.run_query_clickhouse(sql, connection)
     when 'postgresql'
       PostgresHelper.run_query_postgres(sql, connection)
+    when 'databricks'
+      DatabricksHelper.run_query_databricks(sql, connection)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end
@@ -74,6 +80,8 @@ module DatabaseHelper
       ClickHouseHelper.disconnect_from_clickhouse(database)
     when 'postgresql'
       PostgresHelper.disconnect_from_postgres(database)
+    when 'databricks'
+      DatabricksHelper.disconnect_from_databricks(database)
     end
   end
 
@@ -83,6 +91,8 @@ module DatabaseHelper
       ClickHouseHelper.get_schemas_clickhouse(connection)
     when 'postgresql'
       PostgresHelper.get_schemas_postgres(connection)
+    when 'databricks'
+      DatabricksHelper.get_schemas_databricks(connection)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end
@@ -94,6 +104,8 @@ module DatabaseHelper
       ClickHouseHelper.get_tables_by_schema_clickhouse(connection, schema_name)
     when 'postgresql'
       PostgresHelper.get_tables_by_schema_postgres(connection, schema_name)
+    when 'databricks'
+      DatabricksHelper.get_tables_by_schema_databricks(connection, schema_name)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end
@@ -105,6 +117,8 @@ module DatabaseHelper
       ClickHouseHelper.get_columns_by_table_clickhouse(connection, schema_name, table_name)
     when 'postgresql'
       PostgresHelper.get_columns_by_table_postgres(connection, schema_name, table_name)
+    when 'databricks'
+      DatabricksHelper.get_columns_by_table_databricks(connection, schema_name, table_name)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end
@@ -116,6 +130,8 @@ module DatabaseHelper
       ClickHouseHelper.get_foreign_keys_clickhouse(connection, schema_name, table_name, primary_key)
     when 'postgresql'
       PostgresHelper.get_foreign_keys_postgres(connection, schema_name, table_name, primary_key)
+    when 'databricks'
+      DatabricksHelper.get_foreign_keys_databricks(connection, schema_name, table_name, primary_key)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end
@@ -127,6 +143,8 @@ module DatabaseHelper
       ClickHouseHelper.get_schema_column_info_clickhouse(connection, schema_name, tables)
     when 'postgresql'
       PostgresHelper.get_schema_column_info_postgres(connection, schema_name, tables)
+    when 'databricks'
+      DatabricksHelper.get_schema_column_info_databricks(connection, schema_name, tables)
     else
       raise DatabaseError, "Invalid database type: #{database_type}"
     end

data/lib/db/postgres.rb

@@ -2,42 +2,77 @@ require 'pg'
 require 'json'
 require 'uri'
 require 'active_support/inflector'
+require 'connection_pool'
 
 module PostgresHelper
   class << self
     def connect_to_postgres(config)
-      conn = PG.connect(
-        host: config[:host],
-        port: config[:port],
-        dbname: config[:database],
-        user: config[:username],
-        password: config[:password],
-        sslmode: config[:sslmode] || 'prefer',
-        gssencmode: 'disable',  # Disable GSSAPI encryption
-        krbsrvname: nil,        # Disable Kerberos service name
-        target_session_attrs: 'read-write'
-      )
-      conn
+      # Create a connection pool that matches Puma's thread count
+      ConnectionPool.new(size: 5, timeout: 5) do
+        PG.connect(
+          host: config[:host],
+          port: config[:port],
+          dbname: config[:database],
+          user: config[:username],
+          password: config[:password],
+          sslmode: config[:sslmode] || 'prefer',
+          gssencmode: 'disable',  # Disable GSSAPI encryption
+          krbsrvname: nil,        # Disable Kerberos service name
+          target_session_attrs: 'read-write'
+        )
+      end
     end
 
-    def disconnect_from_postgres(client)
-      client.close if client.respond_to?(:close)
+    def disconnect_from_postgres(pool)
+      return unless pool.respond_to?(:shutdown)
+      pool.shutdown { |conn| conn.close if conn.respond_to?(:close) }
+    rescue PG::Error => e
+      puts "Error closing connection pool: #{e.message}"
     end
 
-    def run_query_postgres(sql, client)
-      result = client.exec(sql)
-      
-      fields = result.fields.map do |field|
+    def run_query_postgres(sql, pool)
+      pool.with do |client|
+        result = client.exec(sql)
+        
+        fields = result.fields.map do |field|
+          {
+            name: field,
+            dataTypeID: result.ftype(result.fields.index(field))
+          }
+        end
+
+        # Create a map of field names to their types
+        field_types = result.fields.map.with_index { |field, i| [field, result.ftype(i)] }.to_h
+        json_types = [114, 3802] # JSON and JSONB OIDs
+
+        rows = result.values.map do |row|
+          result.fields.zip(row).map do |field_name, value|
+            # only parse JSON if the field is a JSON type
+            parsed_value = if json_types.include?(field_types[field_name])
+              begin
+                json_value = JSON.parse(value.to_s)
+                if json_value.is_a?(Array)
+                  json_value.map(&:to_s)
+                else
+                  json_value
+                end
+              rescue JSON::ParserError
+                value
+              end
+            else
+              value
+            end
+            [field_name, parsed_value]
+          end.to_h
+        end
+
         {
-          name: field,
-          dataTypeID: result.ftype(result.fields.index(field))
+          fields: fields,
+          rows: rows
         }
+      ensure
+        result&.clear if result.respond_to?(:clear)
       end
-
-      {
-        fields: fields,
-        rows: result.values.map { |row| result.fields.zip(row).to_h }
-      }
     end
 
     def get_schemas_postgres(client)

data/lib/quill-sql.rb

@@ -14,10 +14,8 @@ Dotenv.load(File.expand_path('../.env', __dir__))
 
 module DatabaseType
   POSTGRESQL = 'postgresql'.freeze
-  SNOWFLAKE = 'snowflake'.freeze
-  BIGQUERY = 'bigquery'.freeze
-  MYSQL = 'mysql'.freeze
   CLICKHOUSE = 'clickhouse'.freeze
+  DATABRICKS = 'databricks'.freeze
 
   def self.valid?(type)
     constants.map { |c| const_get(c) }.include?(type)

metadata

@@ -1,14 +1,14 @@
 --- !ruby/object:Gem::Specification
 name: quill-sql
 version: !ruby/object:Gem::Version
-  version: 0.1.5
+  version: 0.1.7
 platform: ruby
 authors:
 - Shawn Magee, Albert Yan
 - Sam Bishop
 bindir: bin
 cert_chain: []
-date: 2025-04-15 00:00:00.000000000 Z
+date: 2025-05-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: json
@@ -94,6 +94,20 @@ dependencies:
     - - ">="
       - !ruby/object:Gem::Version
         version: '0'
+- !ruby/object:Gem::Dependency
+  name: connection_pool
+  requirement: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
+  type: :runtime
+  prerelease: false
+  version_requirements: !ruby/object:Gem::Requirement
+    requirements:
+    - - ">="
+      - !ruby/object:Gem::Version
+        version: '0'
 - !ruby/object:Gem::Dependency
   name: rubocop
   requirement: !ruby/object:Gem::Requirement
@@ -148,6 +162,7 @@ files:
 - lib/assets/pg_types.rb
 - lib/db/cached_connection.rb
 - lib/db/clickhouse.rb
+- lib/db/databricks.rb
 - lib/db/db_helper.rb
 - lib/db/postgres.rb
 - lib/models/filters.rb