quill-sql 0.1.6 → 0.1.7

checksums.yaml

@@ -1,7 +1,7 @@
 ---
 SHA256:
-  metadata.gz: 2a0a484270d6c53d3aace8649706262abc4025f036b2fcbe8b65e9f0e15d017e
-  data.tar.gz: 350c28eb59616ef2f23e23015fbb5c4909d01f90e218730e20b3cae48cd613ad
+  metadata.gz: b3d503d9223878a246289559453c27d45e4393620f9b548433ea5c072fc4c396
+  data.tar.gz: 1820f9f817ab7a29c947555a564e6981b283486cb0cdd0f65396cd01aaaecffd
 SHA512:
-  metadata.gz: c7952deda10956bebd8bbb43e6f5a38f1f58f2a488e8ca2d250a2b68b5bde1701f869f27e2e661461a19a7538994f5b3ca6104f6df820fb72ef607c8016afeeb
-  data.tar.gz: 94fddf190d5903b22b90d21f3f2f886592df35440d8babf4c499f59c8e57033834294acb02a8f6f4f1ec25484713981b489042540607cbba39cfd148a2691d0d
+  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

@@ -41,25 +41,26 @@ module PostgresHelper
           }
         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|
-            # First convert all values to strings
-            string_value = value.to_s
-
-            # Then try to parse JSON if it looks like JSON
-            parsed_value = begin
-              if string_value.start_with?('[') || string_value.start_with?('{')
-                json_value = JSON.parse(string_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
-              else
-                string_value
+              rescue JSON::ParserError
+                value
               end
-            rescue JSON::ParserError
-              string_value
+            else
+              value
             end
             [field_name, parsed_value]
           end.to_h

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.6
+  version: 0.1.7
 platform: ruby
 authors:
 - Shawn Magee, Albert Yan
 - Sam Bishop
 bindir: bin
 cert_chain: []
-date: 2025-04-17 00:00:00.000000000 Z
+date: 2025-05-29 00:00:00.000000000 Z
 dependencies:
 - !ruby/object:Gem::Dependency
   name: json
@@ -162,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