strata-cli 0.1.0.beta

data/lib/strata/cli/sub_commands/datasource.rb

@@ -0,0 +1,353 @@
+require_relative "../guard"
+require_relative "../credentials"
+require_relative "../terminal"
+require "tty-prompt"
+require_relative "../helpers/datasource_helper"
+require_relative "../helpers/description_helper"
+
+module Strata
+  module CLI
+    module SubCommands
+      class Datasource < Thor
+        include Thor::Actions
+        include Guard
+        include Terminal
+        include DatasourceHelper
+        extend Helpers::DescriptionHelper
+
+        desc "adapters", "Lists supported data warehouse adapters"
+        def adapters
+          say "\n\tSupported Adapters\n\n", :yellow
+          DWH.adapters.keys.each do
+            say "\t\t● #{it}", :magenta
+          end
+        end
+
+        desc "list", "List current configured datasources by key and name"
+        def list
+          ds = begin
+            YAML.safe_load_file("datasources.yml", permitted_classes: [Date, Time], aliases: true) || {}
+          rescue Errno::ENOENT
+            {}
+          end
+
+          if ds.empty?
+            say "No datasources configured. Run 'strata datasource add' to add one.", :yellow
+            return
+          end
+
+          names = ds.keys.map { "#{it} => #{ds[it]["name"]}" }
+          out = "\n  #{names.join("\n  ")}"
+          say out, :magenta
+        end
+
+        desc "add [ADAPTER]", "Add a new datasource interactively"
+        long_desc_from_file "datasource/add"
+        def add(adapter_name = nil)
+          prompt = TTY::Prompt.new
+
+          if adapter_name && !DWH.adapters.keys.map(&:to_s).include?(adapter_name)
+            say "Error: '#{adapter_name}' is not a supported adapter", :red
+            say "Supported adapters: #{DWH.adapters.keys.join(", ")}", :yellow
+            return
+          end
+
+          adapter = adapter_name || prompt.select("Choose adapter:", DWH.adapters.keys.map(&:to_s))
+          default_key = generate_default_ds_key(adapter)
+          ds_key = prompt.ask("Datasource key (unique identifier):", default: default_key) do |q|
+            q.required true
+            q.validate(/\A[a-z_][a-z0-9_]*\z/i,
+              "Key must be alphanumeric with underscores, starting with a letter or underscore")
+          end
+
+          say "\n  Configure #{adapter} datasource (press Enter to accept defaults):\n", :yellow
+
+          # Collect common fields
+          config = {"adapter" => adapter}
+          config["name"] = prompt.ask("  Display name:", default: ds_key.upcase)
+          config["description"] = prompt.ask("  Description:", default: "#{adapter.capitalize} datasource")
+          config["tier"] = prompt.select("  Tier:", %w[hot warm cold], default: "warm")
+          config["query_timeout"] = prompt.ask("  Query timeout (seconds):", default: "3600", convert: :int)
+
+          # Collect adapter-specific fields
+          adapter_fields(adapter).each do |field, default_value|
+            config[field] = if field == "auth_mode"
+              prompt.select("  Authentication mode:", %w[pat kp oauth], default: default_value)
+            elsif %w[ssl azure].include?(field)
+              prompt.yes?("  #{field.tr("_", " ").capitalize}?", default: default_value)
+            elsif field == "port"
+              prompt.ask("  #{field.capitalize}:", default: default_value.to_s, convert: :int)
+            else
+              prompt.ask("  #{humanize_field(field)}:", default: default_value)
+            end
+          end
+
+          require_relative "../generators/datasource"
+          generator = Generators::Datasource.new([adapter, ds_key], options.merge(config: config))
+          generator.invoke_all
+
+          say "\n✔ Added #{adapter} config to datasources.yml", :green
+
+          # Automatically collect credentials if required
+          creds = Credentials.new(adapter)
+          if creds.required?
+            say "\n  Now let's set up credentials:\n", :yellow
+            say "  Note: Credentials are stored securely in the local .strata file", :cyan
+            say "  and are NOT committed to the repository (ensured by .gitignore).", :cyan
+            say ""
+            creds.collect
+            creds.write_local(ds_key, self)
+            say "\n✔ Credentials saved to .strata file", :green
+          else
+            say "\n  No credentials required for #{adapter}.", :yellow
+          end
+
+          # AI Setup
+          if ai_not_configured? && prompt.yes?("\n  Enable AI-powered features?", default: true)
+            collect_ai_config(prompt)
+          elsif ai_not_configured?
+            say "\n  AI features skipped. You can configure them later with 'strata ds auth' or manually in .strata",
+              :yellow
+          end
+
+          say "\n✔ Datasource '#{ds_key}' is ready!", :green
+        end
+
+        desc "auth DS_KEY", "Set credentials for the given datasource key (DS_KEY)."
+        long_desc_from_file "datasource/auth"
+        method_option :remote, aliases: ["r"], type: :boolean, desc: "Set credentials on remote server"
+        def auth(ds_key)
+          unless datasources[ds_key]
+            say "Error: Datasource '#{ds_key}' not found in datasources.yml", :red
+            return
+          end
+
+          adapter = datasources[ds_key]["adapter"]
+          creds = Credentials.new(adapter)
+
+          unless creds.required?
+            say "Credentials not required for #{adapter} adapter.", :yellow
+            return
+          end
+
+          say "\nEnter credentials for #{ds_key}", :red
+          say "  Note: Credentials are stored securely in the local .strata file", :cyan
+          say "  and are NOT committed to the repository (ensured by .gitignore).", :cyan
+          say ""
+          creds.collect
+          creds.write_local(ds_key, self)
+
+          say "Credentials saved successfully to .strata file.", :green
+        end
+
+        desc "test DS_KEY", "Test connect to the given datasource."
+        long_desc_from_file "datasource/test"
+        def test(ds_key)
+          adapter = create_adapter(ds_key)
+          with_spinner("Testing #{ds_key} connection...", success_message: "Connected!",
+            failed_message: "Failed to connect.") do
+            adapter.test_connection(raise_exception: true)
+          end
+        rescue => e
+          say "\t!! Failed to connect: \n\t#{e.message}", :red
+        end
+
+        desc "tables DS_KEY", "List tables from DS_KEY datasource"
+        long_desc_from_file "datasource/tables"
+        method_option :pattern, aliases: "p", type: :string, desc: "Regex pattern to filter table list"
+        method_option :catalog, aliases: "c", type: :string, desc: "Change the catalog from the configured one."
+        method_option :schema, aliases: "s", type: :string, desc: "Change the schema from the configured one."
+        def tables(ds_key = nil)
+          prompt = TTY::Prompt.new
+          ds_key = resolve_datasource(ds_key, prompt: prompt)
+          return unless ds_key
+
+          say "\nListing #{ds_key} tables...\n\n", :yellow
+          adapter = create_adapter(ds_key)
+          tables = adapter.tables(**options)
+          tables = tables.select { it =~ /#{options[:pattern]}/ } if options[:pattern]
+
+          if tables.empty?
+            say "No tables found.", :yellow
+            return
+          end
+
+          # Use interactive list for browsing
+          prompt.select("Tables in #{ds_key} (Type to filter):", tables, per_page: 20, filter: true)
+        rescue => e
+          say "\n\t!!Failed: #{e.message}", :red
+        end
+
+        desc "meta DS_KEY TABLE_NAME", "Show the structure of TABLE_NAME in datasource DS_KEY."
+        long_desc_from_file "datasource/meta"
+        method_option :catalog, aliases: "c", type: :string, desc: "Change the catalog from the configured one."
+        method_option :schema, aliases: "s", type: :string, desc: "Change the schema from the configured one."
+        def meta(ds_key, table_name)
+          say "\n● Schema for table: #{table_name} (#{ds_key}):\n", :yellow
+          adapter = create_adapter(ds_key)
+          md = adapter.metadata(table_name, **options.transform_keys { it.to_sym })
+
+          headings = md.columns.first.to_h.keys
+          rows = md.columns.map(&:to_h).map(&:values)
+
+          say print_table(rows, headers: headings, color: :yellow)
+        rescue => e
+          say "\n\t!!Failed: #{e.message}", :red
+        end
+
+        desc 'exec DS_KEY -q "select * from my_table"', "Run the given query or queries on DS_KEY"
+        long_desc_from_file "datasource/exec"
+        method_option :query, aliases: "q", type: :string, desc: "Inline SQL query"
+        method_option :file, aliases: "f", type: :string, desc: "SQL query from file"
+        def exec(ds_key)
+          adapter = create_adapter(ds_key)
+          if options[:file]
+            file_path = validate_file_path(options[:file])
+            queries = File.read(file_path).split(";").reject { it.nil? || it.strip == "" }
+          elsif options[:query]
+            queries = options[:query].split(";").reject { it.nil? || it.strip == "" }
+          else
+            raise StrataError, "Either a file (-f) or a query (-q) should b submitted."
+          end
+
+          queries.each_with_index do |query, index|
+            puts ""
+            res = with_spinner("running #{index + 1}/#{queries.length} queries...") do
+              adapter.execute(query, format: :object)
+            end
+            print_table(res.map(&:values), headers: res.first.keys)
+          end
+        rescue => e
+          say "\n\t!!Failed: #{e.message}", :red
+        end
+
+        private
+
+        def validate_file_path(file_path, project_path = Dir.pwd)
+          expanded = File.expand_path(file_path, project_path)
+          project_root = File.expand_path(project_path)
+          unless expanded.start_with?(project_root)
+            raise Strata::CommandError, "File path must be within project directory"
+          end
+          expanded
+        end
+
+        def generate_default_ds_key(adapter)
+          existing_ds = begin
+            YAML.safe_load_file("datasources.yml", permitted_classes: [Date, Time], aliases: true) || {}
+          rescue Errno::ENOENT
+            {}
+          end
+          base_key = adapter.downcase
+          key_id = 1
+          ds_key = base_key
+          while existing_ds.key?(ds_key)
+            ds_key = "#{base_key}_#{key_id}"
+            key_id += 1
+          end
+          ds_key
+        end
+
+        def adapter_fields(adapter)
+          case adapter
+          when "snowflake"
+            {
+              "account_identifier" => "myorg-myaccount",
+              "database" => "ANALYTICS_DB",
+              "warehouse" => "COMPUTE_WH",
+              "schema" => "PUBLIC",
+              "role" => "ACCOUNTADMIN",
+              "auth_mode" => "pat"
+            }
+          when "postgres"
+            {
+              "host" => "localhost",
+              "port" => 5432,
+              "database" => "mydb",
+              "username" => "postgres",
+              "schema" => "public",
+              "ssl" => false
+            }
+          when "mysql"
+            {
+              "host" => "127.0.0.1",
+              "port" => 3306,
+              "database" => "mydb",
+              "username" => "root",
+              "ssl" => false
+            }
+          when "trino"
+            {
+              "host" => "localhost",
+              "port" => 8080,
+              "catalog" => "native",
+              "username" => "strata_user",
+              "ssl" => false
+            }
+          when "athena"
+            {
+              "region" => "us-east-1",
+              "s3_output_location" => "s3://your-athena-results-bucket/queries/",
+              "database" => "default",
+              "catalog" => "awsdatacatalog"
+            }
+          when "duckdb"
+            {
+              "file" => "./data/warehouse.db"
+            }
+          when "sqlserver"
+            {
+              "host" => "localhost",
+              "port" => 1433,
+              "database" => "mydb",
+              "username" => "sa",
+              "azure" => false
+            }
+          when "druid"
+            {
+              "protocol" => "http",
+              "host" => "localhost",
+              "port" => 7103
+            }
+          else
+            {}
+          end
+        end
+
+        def humanize_field(field)
+          {
+            "account_identifier" => "Account identifier",
+            "s3_output_location" => "S3 output location",
+            "auth_mode" => "Auth mode"
+          }.fetch(field, field.split("_").map(&:capitalize).join(" "))
+        end
+
+        def ai_not_configured?
+          ai_key = CLI.config["ai_api_key"]
+          ai_key.nil? || ai_key.empty?
+        end
+
+        def collect_ai_config(prompt)
+          require_relative "../ai/configuration"
+
+          provider = prompt.select("  AI Provider:", AI::Configuration::PROVIDERS)
+
+          api_key = prompt.mask("  #{provider.capitalize} API Key:")
+
+          return if api_key.nil? || api_key.empty?
+
+          # Append AI config to .strata file
+          ai_config = "\nai_provider: #{provider}\nai_api_key: #{api_key}\n"
+          append_to_file Configuration::STRATA_CONFIG_FILE, ai_config
+
+          # Set restrictive permissions (read/write for owner only)
+          strata_file = Configuration::STRATA_CONFIG_FILE
+          File.chmod(0o600, strata_file) if File.exist?(strata_file)
+
+          say "\n✔ AI configured with #{provider}", :green
+          say "  Note: API key is stored securely in .strata (not committed to repo)", :cyan
+        end
+      end
+    end
+  end
+end