structured_data_to_sql 0.1.0

data/lib/structured_data_to_sql/mysql_dump_xml/table_discovery.rb

@@ -0,0 +1,651 @@
+# frozen_string_literal: true
+
+require "zlib"
+require_relative "../errors"
+require_relative "../io_support"
+require_relative "../conversion_result"
+
+module StructuredDataToSql
+  module MysqlDumpXml
+    class TableDiscovery
+      FileResult =
+        Struct.new(
+          :path,
+          :tables,
+          :table_metadata,
+          :bytes_scanned,
+          :truncated,
+          keyword_init: true
+        ) do
+          def count
+            tables.length
+          end
+        end
+
+      Result =
+        Struct.new(
+          :tables,
+          :files,
+          :files_found,
+          :files_skipped,
+          :skipped,
+          :table_metadata,
+          :bytes_scanned,
+          :truncated,
+          keyword_init: true
+        ) do
+          def sorted_tables
+            sized, unsized =
+              tables.partition do |table|
+                TableDiscovery.size_metadata?(table_metadata[table])
+              end
+            sized.sort_by do |table|
+              metadata = table_metadata[table] || {}
+              [
+                -TableDiscovery.total_size(metadata),
+                -metadata[:rows].to_i,
+                table
+              ]
+            end + unsized.sort
+          end
+        end
+
+      CDATA_START = "<![CDATA[".b
+      CDATA_END = "]]>".b
+      COMMENT_START = "<!--".b
+      COMMENT_END = "-->".b
+      PROCESSING_START = "<?".b
+      PROCESSING_END = "?>".b
+      TABLE_DATA_END_START = "</table_data".b
+      DOCUMENT_MARKER =
+        /<\s*(?:table_structure|table_data|options)\b|<!\[CDATA\[|<!--|<\?/n
+      TABLE_DATA_MARKER = %r{<!\[CDATA\[|<!--|<\?|</table_data\b}n
+      TAG_NAME = /\A<\s*([A-Za-z_:][\w:.-]*)/n
+      ATTR = /\b([A-Za-z_:][\w:.-]*)\s*=\s*(['"])(.*?)\2/mn
+      SCAN_TAIL_BYTES = 64
+
+      def initialize(
+        include_files: nil,
+        exclude_files: nil,
+        chunk_size: 1024 * 1024,
+        max_bytes: nil
+      )
+        @include_files = include_files&.to_set
+        @exclude_files = Array(exclude_files).to_set
+        @chunk_size = chunk_size
+        @max_bytes = max_bytes
+      end
+
+      def discover(
+        source,
+        file_pattern: "*.xml",
+        progress_callback: nil,
+        on_progress: nil,
+        progress_interval: 5,
+        input_gzip: false
+      )
+        files, temporary_inputs =
+          IOSupport.discover(
+            source,
+            patterns: [file_pattern, "#{file_pattern}.gz"].uniq,
+            stream_extension: ".xml",
+            input_gzip:
+          )
+        raise UsageError, "No XML files found in #{source}" if files.empty?
+
+        selected, skipped = files.partition { |file| process_file?(file) }
+        if selected.empty?
+          raise UsageError,
+                "No XML files to scan after filtering. All #{files.length} files were excluded."
+        end
+
+        progress =
+          DiscoveryProgress.new(
+            callback: progress_callback,
+            on_progress: on_progress,
+            interval: progress_interval,
+            total_input_bytes:
+              selected.sum do |file|
+                begin
+                  file.size
+                rescue StandardError
+                  0
+                end
+              end
+          )
+        progress.report(:start, force: true, file_count: selected.length)
+
+        file_results =
+          selected.each_with_index.map do |file, index|
+            progress.start_file(file, index: index + 1, count: selected.length)
+            tables, metadata, bytes_scanned, truncated =
+              discover_file(file, progress: progress)
+            progress.finish_file(force: true)
+            FileResult.new(
+              path: file,
+              tables: tables.sort,
+              table_metadata: metadata,
+              bytes_scanned: bytes_scanned,
+              truncated: truncated
+            )
+          end
+        result =
+          Result.new(
+            tables: file_results.flat_map(&:tables).uniq.sort,
+            files: file_results,
+            files_found: files.length,
+            files_skipped: skipped.length,
+            skipped: skipped,
+            table_metadata: merge_table_metadata(file_results),
+            bytes_scanned: file_results.sum(&:bytes_scanned),
+            truncated: file_results.any?(&:truncated)
+          )
+        progress.report(:complete, force: true, tables: result.tables.length)
+        result
+      rescue StructuredDataToSql::InputError => e
+        raise MysqlDumpXml::InputError, e.message
+      ensure
+        temporary_inputs&.each(&:unlink)
+      end
+
+      def self.total_size(metadata)
+        return 0 unless metadata
+
+        metadata.values_at(:data_length, :index_length).compact.sum
+      end
+
+      def self.size_metadata?(metadata)
+        metadata &&
+          (metadata.key?(:data_length) || metadata.key?(:index_length))
+      end
+
+      private
+
+      class DiscoveryProgress
+        def initialize(callback:, on_progress:, interval:, total_input_bytes:)
+          @callback = callback
+          @on_progress = on_progress
+          @interval = interval
+          @total_input_bytes = total_input_bytes
+          @started_at = Time.now
+          @last_progress_at = nil
+          @bytes_scanned = 0
+          @current_file = nil
+        end
+
+        def start_file(file, index:, count:)
+          @current_file = {
+            path: file.to_s,
+            name: file.basename.to_s,
+            index: index,
+            count: count,
+            size:
+              (
+                begin
+                  file.size
+                rescue StandardError
+                  0
+                end
+              ),
+            bytes_scanned: 0
+          }
+          report(:file_start, force: true)
+        end
+
+        def advance(bytes, tables:)
+          @bytes_scanned += bytes
+          @current_file[:bytes_scanned] += bytes if @current_file
+          report(:bytes, tables: tables)
+        end
+
+        def finish_file(force:)
+          report(:file_complete, force: force)
+          @current_file = nil
+        end
+
+        def report(event, force: false, **extra)
+          return unless @callback || @on_progress
+
+          now = Time.now
+          if !force && @last_progress_at && now - @last_progress_at < @interval
+            return
+          end
+
+          @last_progress_at = now
+          payload = {
+            event: event,
+            elapsed: now - @started_at,
+            bytes_scanned: @bytes_scanned,
+            total_input_bytes: @total_input_bytes,
+            current_file: @current_file&.dup
+          }.merge(extra)
+          @callback&.call(payload)
+          @on_progress&.call(
+            ProgressEvent.new(event, payload.reject { |key, _| key == :event })
+          )
+        end
+      end
+
+      def discover_file(path, progress:)
+        tables = Set.new
+        metadata = {}
+        bytes_scanned = 0
+        truncated = false
+        scanner =
+          TableTagScanner.new(chunk_size: @chunk_size) do |event, table, info|
+            case event
+            when :table
+              tables << table
+            when :metadata
+              tables << table
+              metadata[table] = merge_metadata(metadata[table], info)
+            when :xml_size
+              tables << table
+              metadata[table] = merge_metadata(metadata[table], info)
+            end
+          end
+        open_input(path) do |input|
+          while (chunk = input.read(@chunk_size))
+            if @max_bytes && bytes_scanned + chunk.bytesize > @max_bytes
+              remaining = @max_bytes - bytes_scanned
+              if remaining.positive?
+                scanner.feed(chunk.byteslice(0, remaining))
+                bytes_scanned += remaining
+                progress.advance(remaining, tables: tables.length)
+              end
+              truncated = true
+              break
+            end
+
+            scanner.feed(chunk)
+            bytes_scanned += chunk.bytesize
+            progress.advance(chunk.bytesize, tables: tables.length)
+          end
+        end
+        scanner.finish
+        [tables.to_a, metadata, bytes_scanned, truncated]
+      end
+
+      def merge_table_metadata(file_results)
+        file_results.each_with_object({}) do |file, merged|
+          file.table_metadata.each do |table, info|
+            merged[table] = merge_metadata(merged[table], info)
+          end
+        end
+      end
+
+      def merge_metadata(existing, incoming)
+        return incoming if existing.nil?
+
+        existing.merge(incoming) do |key, old_value, new_value|
+          if key == :xml_data_bytes
+            old_value.to_i + new_value.to_i
+          elsif old_value.is_a?(Integer) && new_value.is_a?(Integer)
+            [old_value, new_value].max
+          else
+            old_value || new_value
+          end
+        end
+      end
+
+      class TableTagScanner
+        def initialize(chunk_size:, &event_handler)
+          @event_handler = event_handler
+          @chunk_size = chunk_size
+          @buffer = +"".b
+          @cursor = 0
+          @skip_until = nil
+          @skip_table_data_body = false
+          @table_data_skip_until = nil
+          @table_data_table = nil
+          @table_data_bytes = 0
+          @current_structure_table = nil
+        end
+
+        def feed(chunk)
+          @buffer << chunk.to_s.b
+          scan(final: false)
+        end
+
+        def finish
+          scan(final: true)
+        end
+
+        private
+
+        def scan(final:)
+          loop do
+            if @skip_table_data_body
+              break unless drain_table_data_body(final: final)
+
+              next
+            end
+
+            if @skip_until
+              break unless drain_skip(final: final)
+
+              next
+            end
+
+            marker = next_marker
+            unless marker
+              advance_document_cursor(final: final)
+              break
+            end
+
+            index, type = marker
+            if type == :tag
+              tag_end = start_tag_end(index)
+              unless tag_end
+                @cursor = index
+                discard_oversized_partial_tag
+                break
+              end
+
+              process_tag(@buffer.byteslice(index, tag_end - index + 1))
+              @cursor = tag_end + 1
+            else
+              start, ending = skip_markers(type)
+              @cursor = index + start.bytesize
+              @skip_until = ending
+              drain_skip(final: final)
+            end
+          end
+        ensure
+          compact_buffer
+        end
+
+        def next_marker
+          match = DOCUMENT_MARKER.match(@buffer, @cursor)
+          return unless match
+
+          index = match.begin(0)
+          [index, marker_type(index)]
+        end
+
+        def start_tag_end(index)
+          quote = nil
+          i = index
+          while i < @buffer.bytesize
+            char = @buffer.getbyte(i)
+            if quote
+              quote = nil if char == quote
+            elsif char == 34 || char == 39
+              quote = char
+            elsif char == 62
+              return i
+            end
+            i += 1
+          end
+          nil
+        end
+
+        def process_tag(tag)
+          tag_name = tag[TAG_NAME, 1]
+          attrs = parse_attributes(tag)
+
+          case tag_name
+          when "table_structure"
+            table = decoded_attribute(attrs["name"])
+            return if table.empty?
+
+            @current_structure_table = table
+            @event_handler.call(:table, table, nil)
+          when "table_data"
+            table = decoded_attribute(attrs["name"])
+            @event_handler.call(:table, table, nil) unless table.empty?
+            if tag.end_with?("/>".b)
+              unless table.empty?
+                @event_handler.call(
+                  :xml_size,
+                  table,
+                  { xml_data_bytes: tag.bytesize }
+                )
+              end
+            else
+              @skip_table_data_body = true
+              @table_data_table = table
+              @table_data_bytes = tag.bytesize
+            end
+          when "options"
+            table =
+              decoded_attribute(attrs["Name"] || attrs["name"]) ||
+                @current_structure_table
+            table = @current_structure_table if table.to_s.empty?
+            metadata = options_metadata(attrs)
+            if table && !table.empty? && metadata.any?
+              @event_handler.call(:metadata, table, metadata)
+            end
+          end
+        end
+
+        def drain_table_data_body(final:)
+          return drain_table_data_special(final:) if @table_data_skip_until
+
+          marker = next_table_data_marker
+          unless marker
+            keep =
+              (
+                if final
+                  0
+                else
+                  [SCAN_TAIL_BYTES, TABLE_DATA_END_START.bytesize - 1].max
+                end
+              )
+            consume_table_data_to([@buffer.bytesize - keep, @cursor].max)
+            return false
+          end
+
+          index, type = marker
+          if type == :table_data_end
+            tag_end = start_tag_end(index)
+            unless tag_end
+              consume_table_data_to(index)
+              return false
+            end
+
+            consume_table_data_to(tag_end + 1)
+            finish_table_data
+            true
+          else
+            start, ending = skip_markers(type)
+            consume_table_data_to(index + start.bytesize)
+            @table_data_skip_until = ending
+            true
+          end
+        end
+
+        def drain_table_data_special(final:)
+          index = @buffer.index(@table_data_skip_until, @cursor)
+          if index
+            consume_table_data_to(index + @table_data_skip_until.bytesize)
+            @table_data_skip_until = nil
+            true
+          else
+            keep =
+              (
+                if final
+                  0
+                else
+                  [SCAN_TAIL_BYTES, @table_data_skip_until.bytesize - 1].max
+                end
+              )
+            consume_table_data_to([@buffer.bytesize - keep, @cursor].max)
+            false
+          end
+        end
+
+        def next_table_data_marker
+          match = TABLE_DATA_MARKER.match(@buffer, @cursor)
+          return unless match
+
+          index = match.begin(0)
+          [index, marker_type(index)]
+        end
+
+        def marker_type(index)
+          case @buffer.getbyte(index + 1)
+          when 33
+            @buffer.getbyte(index + 2) == 45 ? :comment : :cdata
+          when 63
+            :processing
+          when 47
+            :table_data_end
+          else
+            :tag
+          end
+        end
+
+        def consume_table_data_to(position)
+          bytes = position - @cursor
+          return unless bytes.positive?
+
+          @table_data_bytes += bytes
+          @cursor = position
+        end
+
+        def finish_table_data
+          unless @table_data_table.to_s.empty?
+            @event_handler.call(
+              :xml_size,
+              @table_data_table,
+              { xml_data_bytes: @table_data_bytes }
+            )
+          end
+          @skip_table_data_body = false
+          @table_data_skip_until = nil
+          @table_data_table = nil
+          @table_data_bytes = 0
+        end
+
+        def parse_attributes(tag)
+          tag
+            .scan(ATTR)
+            .each_with_object({}) do |(name, _quote, value), attrs|
+              attrs[name] = value
+            end
+        end
+
+        def decoded_attribute(value)
+          return "" if value.nil?
+
+          value =
+            value
+              .dup
+              .force_encoding("UTF-8")
+              .encode("UTF-8", invalid: :replace, undef: :replace)
+          unescape_xml_attribute(value)
+        end
+
+        def options_metadata(attrs)
+          {
+            rows: integer_attribute(attrs["Rows"] || attrs["rows"]),
+            data_length:
+              integer_attribute(attrs["Data_length"] || attrs["data_length"]),
+            index_length:
+              integer_attribute(attrs["Index_length"] || attrs["index_length"])
+          }.compact
+        end
+
+        def integer_attribute(value)
+          return nil if value.nil? || value.empty?
+
+          Integer(value, 10)
+        rescue ArgumentError
+          nil
+        end
+
+        def unescape_xml_attribute(value)
+          value.gsub(
+            /&(?:amp|lt|gt|quot|apos|#\d+|#x[0-9a-fA-F]+);/
+          ) do |entity|
+            case entity
+            when "&amp;"
+              "&"
+            when "&lt;"
+              "<"
+            when "&gt;"
+              ">"
+            when "&quot;"
+              '"'
+            when "&apos;"
+              "'"
+            when /\A&#(\d+);\z/
+              Regexp.last_match(1).to_i.chr(Encoding::UTF_8)
+            when /\A&#x([0-9a-fA-F]+);\z/
+              Regexp.last_match(1).to_i(16).chr(Encoding::UTF_8)
+            else
+              entity
+            end
+          rescue RangeError
+            entity
+          end
+        end
+
+        def skip_markers(type)
+          case type
+          when :cdata
+            [CDATA_START, CDATA_END]
+          when :comment
+            [COMMENT_START, COMMENT_END]
+          when :processing
+            [PROCESSING_START, PROCESSING_END]
+          end
+        end
+
+        def drain_skip(final:)
+          index = @buffer.index(@skip_until, @cursor)
+          if index
+            @cursor = index + @skip_until.bytesize
+            @skip_until = nil
+            true
+          else
+            keep = final ? 0 : [SCAN_TAIL_BYTES, @skip_until.bytesize - 1].max
+            @cursor = [@buffer.bytesize - keep, @cursor].max
+            false
+          end
+        end
+
+        def advance_document_cursor(final:)
+          keep = final ? 0 : SCAN_TAIL_BYTES
+          @cursor = [@buffer.bytesize - keep, @cursor].max
+        end
+
+        def discard_oversized_partial_tag
+          return if @buffer.bytesize - @cursor <= @chunk_size + SCAN_TAIL_BYTES
+
+          @cursor += 1
+        end
+
+        def compact_buffer
+          return if @cursor.zero?
+
+          @buffer.slice!(0, @cursor)
+          @cursor = 0
+        end
+      end
+
+      def open_input(path)
+        if path.to_s.end_with?(".gz")
+          Zlib::GzipReader.open(path.to_s) { |gz| yield gz }
+        else
+          File.open(path, "rb") { |file| yield file }
+        end
+      end
+
+      def process_file?(file)
+        file_name = file.basename.to_s
+        base = file_name.sub(/\.gz\z/, "").sub(/\.xml\z/, "")
+        if @include_files && !@include_files.include?(file_name) &&
+             !@include_files.include?(base)
+          return false
+        end
+        if @exclude_files.include?(file_name) || @exclude_files.include?(base)
+          return false
+        end
+
+        true
+      end
+    end
+  end
+end