File: generate-selects-for-missing-sharding-key

package info (click to toggle)
gitlab 17.6.5-19
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 629,368 kB
  • sloc: ruby: 1,915,304; javascript: 557,307; sql: 60,639; xml: 6,509; sh: 4,567; makefile: 1,239; python: 406
file content (119 lines) | stat: -rwxr-xr-x 2,692 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
#!/usr/bin/env ruby
# frozen_string_literal: true

require 'optparse'
require 'open3'

class GenerateSelectsForMissingShardingKey
  attr_reader :psql_command

  def self.for_cli
    options = parse_options

    new(**options)
  end

  def self.parse_options
    options = {}

    OptionParser.new do |opts|
      opts.banner = "Usage: generate-selects-for-missing-sharding-key [options]"

      opts.on('-p', '--psql COMMAND', 'Default: gdk psql -c') do |value|
        options[:psql_command] = value
      end
      opts.on('--debug', 'Enable debugging output') { options[:debug] = true }
    end.parse!

    options
  end

  def initialize(psql_command: 'gdk psql -c', debug: false)
    @psql_command = psql_command
    @debug = debug
  end

  def execute
    print_selects_for_tables_without_sharding_key
  end

  def print_selects_for_tables_without_sharding_key
    table_names = get_table_names(without_columns: %w[organization_id namespace_id project_id])
    table_names.each do |table_name|
      copy_select_sql = <<~SQL
        COPY (
          SELECT #{table_name}.*
          FROM #{table_name}
        )
        TO STDOUT
        WITH (FORMAT CSV, HEADER)
      SQL

      puts copy_select_sql.split.join(' ')
    end
  end

  def get_table_names(with_column: nil, without_columns: nil)
    command = <<-SQL
      COPY (
        SELECT DISTINCT table_name
        FROM information_schema.columns
        WHERE
          table_schema = 'public'
          #{with_column_clause(with_column)}
          #{without_columns_clause(without_columns)}
      ) TO STDOUT
    SQL

    table_names_str, status = psql(command)
    raise "Failed to get tables with #{column_name} column" unless status.success?

    table_names_str.split("\n")
  end

  def with_column_clause(with_column)
    return unless with_column

    "AND column_name = '#{with_column}'"
  end

  def without_columns_clause(without_columns)
    return unless without_columns

    <<-SQL
      AND table_name NOT IN (
        SELECT table_name
        FROM information_schema.columns
        WHERE
          table_schema = 'public'
          AND column_name IN ('#{without_columns.join("','")}')
      )
    SQL
  end

  def psql(query)
    command = psql_command.split << query.split.join(' ')

    capture2(command)
  end

  def capture2(command)
    debug %(Run command: #{command.inspect})

    stdout_str, status = Open3.capture2(*command)

    debug %(Run command: stdout: "#{stdout_str}", exitstatus: "#{status.exitstatus}")

    [stdout_str, status]
  end

  def debug(output)
    puts "[DEBUG] #{output}" if debug?
  end

  def debug?
    !!@debug
  end
end

GenerateSelectsForMissingShardingKey.for_cli.execute