conceptql 0.2.0 → 0.4.0

data/lib/conceptql/rdbms/postgres.rb

@@ -0,0 +1,51 @@
+# frozen_string_literal: true
+
+require_relative "generic"
+
+module ConceptQL
+  module Rdbms
+    class Postgres < Generic
+      # Returns from - to (see Generic#datediff).
+      def datediff(from, to)
+        Sequel.extract(:days, Sequel.cast(from, Time) - Sequel.cast(to, Time))
+      end
+
+      def post_create(db, table_name)
+        db.vacuum_table(table_name, analyze: true) if analyze_temp_tables?
+      end
+
+      def preferred_formatter
+        SqlFormatters::POSTGRES
+      end
+
+      # DISTINCT ON is Postgres-only. The other adapters inherit
+      # Generic#primary_concepts, which uses ROW_NUMBER() instead.
+      def primary_concepts(db, all_primary_ids)
+        db[Sequel[:clinical_codes].as(:pcc)]
+          .where(provenance_concept_id: all_primary_ids)
+          .select(
+            Sequel[:pcc][:collection_id].as(:collection_id),
+            Sequel[:pcc][:clinical_code_concept_id].as(:concept_id),
+            Sequel[:pcc][:clinical_code_source_value].as(:concept_code),
+            Sequel[:pcc][:clinical_code_vocabulary_id].as(:vocabulary_id),
+            Sequel[true].as(:is_primary)
+          )
+          .order(Sequel[:pcc][:collection_id], Sequel[:pcc][:clinical_code_concept_id])
+          .from_self
+          .distinct(:collection_id)
+      end
+
+      def analyze_temp_tables?
+        ENV["CONCEPTQL_PG_ANALYZE_TEMP_TABLES"] == "true"
+      end
+
+      def explain_temp_tables?
+        ENV["CONCEPTQL_PG_EXPLAIN_TEMP_TABLES"] == "true"
+      end
+
+      def supports_materialized?
+        true
+      end
+    end
+  end
+end

data/lib/conceptql/rdbms/spark.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require_relative "generic"
+
+module ConceptQL
+  module Rdbms
+    class Spark < Generic
+      def days_between(from_column, to_column)
+        Sequel.function(:datediff, cast_date(to_column), cast_date(from_column))
+      end
+
+      # <=> is Spark's NULL-safe equality operator (see Generic#null_safe_equal).
+      def null_safe_equal(lhs, rhs)
+        Sequel.lit("(? <=> ?)", lhs, rhs)
+      end
+
+      def preferred_formatter
+        SqlFormatters::PG_FORMAT
+      end
+    end
+  end
+end

data/lib/conceptql/rdbms.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require_relative "rdbms/postgres"
+require_relative "rdbms/spark"
+
+module ConceptQL
+  module Rdbms
+    def self.generate(database_type)
+      case database_type.to_sym
+      when :postgres
+        ConceptQL::Rdbms::Postgres.new
+      when :spark
+        ConceptQL::Rdbms::Spark.new
+      when :duckdb
+        require_relative "rdbms/duckdb"
+        ConceptQL::Rdbms::DuckDB.new
+      else
+        raise "Unknown database_type -- '#{database_type}'"
+      end
+    end
+  end
+end

data/lib/conceptql/scope.rb

@@ -1,6 +1,12 @@
-require_relative 'utils/temp_table'
+# frozen_string_literal: true
+
+require "English"
+require_relative "window"
+require "securerandom"
+require "sequel"
+
 module ConceptQL
-  # Scope coordinates the creation of any TempTables that might
+  # Scope coordinates the creation of any common table expressions that might
   # be used when a Recall operator is present in the statement.
   #
   # Any time an operator is given a label, it becomes a candidate
@@ -8,62 +14,506 @@ module ConceptQL
   # somewhere else in the statement.
   #
   # Scope keeps track of all labeled operators and provides an
-  # API for Recall operators to fetch the results/types from
+  # API for Recall operators to fetch the results/domains from
   # labeled operators.
   class Scope
+    # Raised when the extracted CTEs read each other in a cycle, so no order
+    # puts every CTE after the CTEs it reads.
+    class CteCycleError < StandardError
+      def initialize(names)
+        super("These CTEs read each other in a cycle, so they have no valid order: #{names.join(', ')}")
+      end
+    end
+
+    DEFAULT_COLUMN_FAMILY = "default"
+
+    DEFAULT_COLUMNS = {
+      person_id: :Bigint,
+      criterion_id: :Bigint,
+      criterion_table: :String,
+      criterion_domain: :String,
+      start_date: :Date,
+      end_date: :Date,
+      source_value: :String,
+      source_vocabulary_id: :String,
+      label: :String,
+      column_family: :String
+    }.freeze
+
+    ADDITIONAL_COLUMNS = {
+      value_as_number: :Float,
+      value_as_string: :String,
+      value_as_concept_id: :Bigint,
+      value_as_concept: :String,
+      unit_concept_id: :Bigint,
+      unit_concept: :String,
+      context_id: :Bigint,
+      visit_occurrence_id: :Bigint,
+      file_provenance_type: :Bigint,
+      code_provenance_type: :Bigint,
+      provider_id: :Bigint,
+      context_specialty_concept_id: :Bigint,
+      practitioner_specialty_concept_id: :Bigint,
+      role_type_concept_id: :Bigint,
+      visit_source_concept_id: :Bigint,
+      range_low: :Float,
+      range_high: :Float,
+      drug_name: :String,
+      drug_amount: :Float,
+      drug_amount_units: :String,
+      drug_strength_source_value: :String,
+      drug_days_supply: :Bigint,
+      drug_quantity: :Float,
+      admission_date: :Date,
+      discharge_date: :Date,
+      length_of_stay: :Bigint,
+      admission_source_value: :String,
+      admission_source_description: :String,
+      discharge_location_source_value: :String,
+      discharge_location_source_description: :String,
+      primary_concept_id: :Bigint,
+      primary_diagnosis_code: :String,
+      primary_diagnosis_vocabulary: :String,
+      uuid: :String,
+      window_id: :Bigint
+    }.freeze
+
+    COLUMN_TYPES = DEFAULT_COLUMNS.merge(ADDITIONAL_COLUMNS).freeze
+
+    # Marks the WITH entry of a list that the Database's share_utilization
+    # option keeps in a shared table (Behaviors::Utilizable#evaluate).
+    SHARED_TABLE = :conceptql_shared_table
+
     attr_accessor :person_ids
-    attr :known_operators
-    def initialize
+
+    attr_reader :known_operators, :recall_stack, :recall_dependencies, :annotation, :opts, :query_columns
+
+    def initialize(opts = {})
       @known_operators = {}
-      @flagged = {}
+      @recall_dependencies = {}
+      @recall_stack = []
+      @label_cte_names = {}
+      @annotation = {}
+      @opts = opts.dup
+      @annotation[:errors] = @errors = {}
+      @annotation[:warnings] = @warnings = {}
+      @annotation[:counts] = @counts = {}
+      @query_columns = DEFAULT_COLUMNS.keys
+      @query_columns << :window_id if opts.dig(:window_opts, :window_table)
+
+      @i = 0
+      @mutex = Mutex.new
+      @cte_name_next = -> { @mutex.synchronize { @i += 1 } }
+
+      return unless force_temp_tables? && scratch_database.to_s.empty?
+
+      raise ArgumentError,
+            "You must set the DOCKER_SCRATCH_DATABASE environment variable to the name of the scratch database if using the CONCEPTQL_FORCE_TEMP_TABLES environment variable"
     end
 
-    def add_operator(operator)
-      # Recall operators respond to source, so when we add such an
-      # operator, we want to flag it so that we can come back to it later
-      # and create the appropriate temp tables
-      if operator.respond_to?(:source)
-        flagged[operator.source] = true
+    def add_errors(key, errors)
+      @errors[key] = errors
+    end
+
+    def add_warnings(key, errors)
+      @warnings[key] = errors
+    end
+
+    def add_counts(key, domain, counts)
+      c = @counts[key] ||= {}
+      c[domain] = counts
+    end
+
+    def add_required_columns(op)
+      @query_columns |= op.required_columns if op.required_columns
+    end
+
+    def harvest_person_ids(op)
+      if (person_ids = op.options[:person_ids])
+        opts[:person_ids] = person_ids
+      end
+    end
+
+    def nest(op)
+      add_required_columns(op)
+      harvest_person_ids(op)
+      return yield unless (label = op.is_a?(Operators::Recall) ? op.source : op.label)
+
+      unless label.is_a?(String)
+        op.instance_eval do
+          @errors = []
+          add_error("invalid label")
+        end
+        return
       end
-      return unless operator.label
+
+      recall_dependencies[label] ||= []
+
+      if nested_recall?(label)
+        op.instance_eval do
+          @errors = []
+          add_error("nested recall")
+        end
+        return
+      end
+
+      if duplicate_label?(label) && !op.is_a?(Operators::Recall)
+        op.instance_eval do
+          @errors = []
+          add_error("duplicate label")
+        end
+      end
+
+      if (last = recall_stack.last)
+        recall_dependencies[last] << label
+      end
+
+      begin
+        recall_stack.push(label)
+        yield
+      ensure
+        recall_stack.pop if recall_stack.last == label
+      end
+    end
+
+    def add_operator(operator)
       known_operators[operator.label] = operator
     end
 
+    def duplicate_label?(label)
+      known_operators.keys.map(&:downcase).include?(label.downcase)
+    end
+
+    def nested_recall?(label)
+      recall_stack.map(&:downcase).include?(label.downcase)
+    end
+
     def from(db, label)
-      temp_table(db, label).from(db)
+      if ConceptQL.avoid_ctes?
+        ds = db.from(fetch_operator(label).evaluate(db).as(label_cte_name(label)))
+      else
+        ds = db.from(label_cte_name(label))
+
+        if ENV["CONCEPTQL_CHECK_COLUMNS"]
+          # Work around requests for columns by operators.  These
+          # would fail because the CTE would not be defined.  You
+          # don't want to define the CTE normally, but to allow the
+          # columns to still work, send the columns request to the
+          # underlying operator.
+          op = fetch_operator(label)
+          ds = ds.with_extend do
+            define_method(:columns) do
+              (@main_op ||= op.evaluate(db)).columns
+            end
+          end
+        end
+      end
+
+      ds
     end
 
-    def types(label)
-      fetch_operator(label).types
+    def domains(label, db)
+      fetch_operator(label).domains(db)
+    rescue StandardError
+      [ :invalid ]
     end
 
-    def sql(db)
-      temp_tables(db).values.map do |temp|
-        temp.sql(db)
+    def sort_ctes(sorted, unsorted, deps)
+      return sorted if unsorted.empty?
+
+      add, unsorted = unsorted.partition do |label, _|
+        deps[label].empty?
       end
+
+      sorted += add
+
+      new_deps = {}
+      deps.map do |label, dps|
+        new_deps[label] = dps - sorted.map(&:first)
+      end
+
+      sort_ctes(sorted, unsorted, new_deps)
     end
 
-    def prep(db)
-      temp_tables(db).values.map do |temp|
-        temp.build(db)
+    def valid?
+      recall_dependencies.each_value do |deps|
+        return false unless (deps - known_operators.keys).empty?
       end
+      true
     end
 
-    private
-    attr :flagged
+    class CteExtractor < Sequel::ASTTransformer
+      def initialize(scope, ctes)
+        super()
+        @scope = scope
+        @ctes = ctes
+      end
+
+      private
 
-    def temp_tables(db)
-      Hash[flagged.keys.map do |label|
-        [label, TempTable.new(fetch_operator(label), db)]
-      end]
+      def v(o)
+        if o.is_a?(Sequel::Dataset)
+          @scope.recursive_extract_ctes(o, @ctes)
+        else
+          super
+        end
+      end
     end
 
-    def temp_table(db, label)
-      temp_tables(db)[label]
+    def recursive_extract_cte_expr(t, ctes)
+      case t
+      when Sequel::Dataset
+        recursive_extract_ctes(t, ctes)
+      when Sequel::SQL::AliasedExpression
+        if t.expression.is_a?(Sequel::Dataset)
+          Sequel.as(recursive_extract_ctes(t.expression, ctes), t.alias)
+        else
+          t
+        end
+      else
+        t
+      end
+    end
+
+    def recursive_extract_ctes(query, ctes)
+      # puts
+      # p [:rec, ctes, query.opts]
+
+      if (with = query.opts[:with])
+        keep, remove = with.partition { |w| w[:no_temp_table] }
+        ctes.concat(remove.map do |w|
+          opts = w.except(:name, :dataset, :no_temp_table)
+          # A shared table is created from its whole list, nested CTEs included.
+          ds = w[SHARED_TABLE] ? w[:dataset] : recursive_extract_ctes(w[:dataset], ctes)
+          [ w[:name], ds, opts ]
+        end)
+        # p [:rec_with, ctes.map(&:first), with]
+        query = query.clone(with: keep.empty? ? nil : keep)
+      end
+
+      if (from = query.opts[:from])
+        query = query.clone(from: from.map { |t| recursive_extract_cte_expr(t, ctes) })
+        # p [:rec_from, ctes.map(&:first), from]
+      end
+
+      if (joins = query.opts[:join])
+        query = query.clone(join: joins.map do |jc|
+          table_expr = recursive_extract_cte_expr(jc.table_expr, ctes)
+          # A CROSS JOIN is a plain JoinClause, with no ON condition.
+          if jc.respond_to?(:on)
+            jc.class.new(jc.on, jc.join_type, table_expr)
+          else
+            jc.class.new(jc.join_type, table_expr)
+          end
+        end)
+        # p [:rec_join, ctes.map(&:first), joins]
+      end
+
+      if (compounds = query.opts[:compounds])
+        query = query.clone(compounds: compounds.map { |t, ds, a| [ t, recursive_extract_ctes(ds, ctes), a ] })
+        # p [:rec_compounds, ctes.map(&:first), compounds]
+      end
+
+      if (where = query.opts[:where])
+        query = query.clone(where: CteExtractor.new(self, ctes).transform(where))
+      end
+
+      query
+    end
+
+    # Removes duplicate CTE names, then moves each CTE after the CTEs it reads.
+    #
+    # A CTE reads another when its SQL contains the other's name as the
+    # adapter quotes it (ds.db.literal): "name" on Postgres and DuckDB,
+    # `name` on Spark, and "scratch"."name" for a temp table in a scratch
+    # database. The closing quote stops a name from matching the start of a
+    # longer name. Each name also ends with the scope's counter and 32 random
+    # hex digits, so a name does not appear in other SQL by accident.
+    #
+    # The sort is stable: each step takes the first CTE whose dependencies are
+    # already placed. recursive_extract_ctes appends the CTEs that a dataset
+    # reads before the dataset itself, so an order that is already correct
+    # does not change. (sort_ctes sorts by layers, and would move CTEs.)
+    def sort_extracted_ctes(ctes)
+      unsorted = ctes.uniq(&:first)
+      literals = unsorted.to_h { |name, ds| [ name, ds.db.literal(name) ] }
+      deps = unsorted.to_h do |name, ds|
+        sql = ds.sql
+        [ name, literals.filter_map { |other, literal| other if other != name && sql.include?(literal) } ]
+      end
+
+      sorted = []
+      until unsorted.empty?
+        placed = sorted.map(&:first)
+        index = unsorted.index { |name, _| (deps[name] - placed).empty? }
+        raise(CteCycleError, unsorted.map { |name, _| literals[name] }) unless index
+
+        sorted << unsorted.delete_at(index)
+      end
+      sorted
+    end
+
+    def with_ctes(op, db, options = {})
+      raise "recall operator use without matching label" unless valid?
+
+      query = op.evaluate(db)
+      rdbms = op.rdbms
+
+      query = query.from_self
+      temp_tables = []
+      ctes.each do |label, operator|
+        temp_tables << [ label_cte_name(label), recursive_extract_ctes(operator.evaluate(db), temp_tables) ]
+      end
+      query = recursive_extract_ctes(query, temp_tables)
+      shared_tables, temp_tables = temp_tables.partition { |_, _, with_opts| with_opts&.fetch(SHARED_TABLE, false) }
+      shared_tables = shared_tables.uniq(&:first)
+      temp_tables = sort_extracted_ctes(temp_tables)
+
+      if force_temp_tables?(options)
+        scope = self
+        query = query.with_extend do
+          # Create temp tables for each CTE
+          #
+          # Need to override multiple methods if sequel_pg is in use, as in
+          # that case not every method calls each.  For other adapters or
+          # when sequel_pg is not in use, it is probably safe to override just each.
+          # There are other methods that may need to be overridden in order to handle
+          # all cases when sequel_pg is in use.
+          %i[each to_hash_groups to_hash].each do |meth|
+            define_method(meth) do |*args, &block|
+              if !temp_tables.empty? && !opts[:conceptql_temp_tables_created]
+                begin
+                  temp_tables.each do |table_name, ds, _with_opts|
+                    db.create_table(table_name, rdbms.create_options(scope, ds).merge(options).merge(as: ds))
+                    rdbms.post_create(db, table_name)
+                  end
+
+                  clone(conceptql_temp_tables_created: true).send(meth, *args, &block)
+                ensure
+                  drop_temp_tables
+                end
+              else
+                super(*args, &block)
+              end
+            end
+          end
+
+          define_method(:columns) do
+            if self.db.respond_to?(:record_table)
+              temp_tables.each do |table_name, ds, _with_opts|
+                self.db.record_table(table_name, self.db.columns_from_sql(ds))
+              end
+            end
+            super()
+          end
+
+          define_method(:sql_statements) do |*args, &block|
+            sql_statements = temp_tables.map do |table_name, ds|
+              [ table_name, ds.sql ]
+            end.compact.push([ :query, sql(*args, &block) ])
+            Hash[scope.shared_table_statements(shared_tables, op) + sql_statements]
+          end
+
+          define_method(:drop_temp_tables) do |opts = {}|
+            opts ||= {}
+            table_names = temp_tables.reverse_each.map(&:first)
+            begin
+              db.drop_table?(*table_names, rdbms.drop_options.merge(opts))
+            rescue Sequel::DatabaseError
+              warn("Unable to drop scratch table: #{literal(*table_names)}")
+            end
+          end
+        end
+      else
+        unless ConceptQL.avoid_ctes?
+          temp_tables.each do |table_name, ds, with_opts|
+            query = if with_opts.nil? || with_opts.empty?
+                      query.with(table_name, ds)
+            else
+                      query.with(table_name, ds, with_opts)
+            end
+          end
+        end
+
+        scope = self
+        query = query.with_extend do
+          define_method(:sql_statements) do |*args, &block|
+            Hash[scope.shared_table_statements(shared_tables, op).push([ :query, sql(*args, &block) ])]
+          end
+
+          define_method(:drop_temp_tables) do |opts = {}|
+          end
+        end
+      end
+
+      shared_tables.empty? ? query : create_shared_tables_on_read(query, shared_tables, op)
+    end
+
+    # The sql_statements entries that create the shared tables that do not
+    # exist yet. They come first, because the other entries read them.
+    def shared_table_statements(shared_tables, op)
+      shared_tables.filter_map { |name, key, _with_opts| op.nodifier.cdb.shared_utilization_statement(name, key) }
+    end
+
+    # Creates the shared tables that do not exist yet before the query reads
+    # its first rows, for a caller that reads rows through conceptql. This
+    # hook runs before the force_temp_tables hook, so the temp tables can read
+    # them. The shared tables are never in drop_temp_tables.
+    def create_shared_tables_on_read(query, shared_tables, op)
+      scope = self
+      cdb = op.nodifier.cdb
+      rdbms = op.rdbms
+      query.with_extend do
+        %i[each to_hash_groups to_hash].each do |meth|
+          define_method(meth) do |*args, &block|
+            shared_tables.each do |name, key, _with_opts|
+              cdb.create_shared_utilization_table(name, key, rdbms, scope)
+            end
+            super(*args, &block)
+          end
+        end
+      end
+    end
+
+    def label_cte_name(label)
+      @label_cte_names[label] ||= cte_name(label)
+    end
+
+    def ctes
+      @ctes ||= sort_ctes([], known_operators, recall_dependencies)
     end
 
     def fetch_operator(label)
-      known_operators[label] || raise("No operator with label: '#{label}'")
+      known_operators[label]
+    end
+
+    def windows
+      @windows ||= Window.from(opts.merge(opts[:window_opts] || {}))
+    end
+
+    def force_temp_tables?(options = {})
+      # opts.fetch(:force_temp_tables, opts[:force_temp_tables])
+      options.fetch(:force_temp_tables, opts[:force_temp_tables])
+    end
+
+    def scratch_database
+      opts[:scratch_database]
+    end
+
+    def cte_name(name, options = {})
+      cte_parts = [
+        opts[:table_prefix],
+        name.to_s.gsub(/\W+/, "_"),
+        $PROCESS_ID,
+        @cte_name_next.call,
+        SecureRandom.hex(16)
+      ].compact
+
+      name = Sequel.identifier(cte_parts.join("_"))
+
+      name = name.qualify(scratch_database) if force_temp_tables?(options) && scratch_database
+
+      name
     end
   end
 end

data/lib/conceptql/sql_formatters/formatter.rb

@@ -0,0 +1,25 @@
+# frozen_string_literal: true
+
+module ConceptQL
+  module SqlFormatters
+    # An external SQL formatter, given as a command array that reads SQL on
+    # stdin, e.g. Formatter.new(%w[sql-formatter -l postgresql]).
+    class Formatter
+      attr_reader :command
+
+      def initialize(command)
+        @command = command
+      end
+
+      def available?
+        system("which", command.first, out: File::NULL, err: File::NULL)
+      end
+
+      def format(sql)
+        ConceptQL::Utils.timed_capture(*command, stdin_data: sql, timeout: 10)
+      rescue Timeout::Error
+        sql
+      end
+    end
+  end
+end

data/lib/conceptql/sql_formatters.rb

@@ -0,0 +1,22 @@
+# frozen_string_literal: true
+
+require_relative "sql_formatters/formatter"
+
+module ConceptQL
+  module SqlFormatters
+    SQL_FORMATTER = %w[sql-formatter].freeze
+    POSTGRES = %w[sql-formatter -l postgresql].freeze
+    SQLFORMAT = %w[sqlformat --reindent --use_space_around_operators --identifiers lower --keywords upper -].freeze
+    PG_FORMAT = %w[pg_format].freeze
+
+    # Formats with the rdbms's preferred formatter, else the first of the
+    # fallbacks that is installed, else returns the SQL untouched.
+    def self.format(sql, rdbms)
+      formatter = [ rdbms.preferred_formatter, SQL_FORMATTER, SQLFORMAT, PG_FORMAT ]
+                  .compact
+                  .map { |command| Formatter.new(command) }
+                  .detect(&:available?)
+      formatter ? formatter.format(sql) : sql
+    end
+  end
+end