|
| 1 | +# encoding: UTF-8 |
| 2 | +# frozen_string_literal: true |
| 3 | +# |
| 4 | +# Copyright (c) 2021 GoodData Corporation. All rights reserved. |
| 5 | +# This source code is licensed under the BSD-style license found in the |
| 6 | +# LICENSE file in the root directory of this source tree. |
| 7 | + |
| 8 | +require 'securerandom' |
| 9 | +require 'java' |
| 10 | +require 'pathname' |
| 11 | +require_relative '../cloud_resource_client' |
| 12 | + |
| 13 | +base = Pathname(__FILE__).dirname.expand_path |
| 14 | +Dir.glob(base + 'drivers/*.jar').each do |file| |
| 15 | + require file unless file.start_with?('lcm-postgresql-driver') |
| 16 | +end |
| 17 | + |
| 18 | +module GoodData |
| 19 | + module CloudResources |
| 20 | + class PostgresClient < CloudResourceClient |
| 21 | + JDBC_POSTGRES_PATTERN = %r{jdbc:postgresql:\/\/([^:^\/]+)(:([0-9]+))?(\/)?} |
| 22 | + POSTGRES_DEFAULT_PORT = 5432 |
| 23 | + JDBC_POSTGRES_PROTOCOL = 'jdbc:postgresql://' |
| 24 | + SSL_JAVA_FACTORY = '&sslfactory=org.postgresql.ssl.DefaultJavaSSLFactory' |
| 25 | + VERIFY_FULL = 'verify-full' |
| 26 | + PREFER = 'prefer' |
| 27 | + REQUIRE = 'require' |
| 28 | + POSTGRES_SET_SCHEMA_COMMAND = "set search_path to" |
| 29 | + POSTGRES_FETCH_SIZE = 1000 |
| 30 | + |
| 31 | + class << self |
| 32 | + def accept?(type) |
| 33 | + type == 'postgresql' |
| 34 | + end |
| 35 | + end |
| 36 | + |
| 37 | + def initialize(options = {}) |
| 38 | + raise("Data Source needs a client to Postgres to be able to query the storage but 'postgresql_client' is empty.") unless options['postgresql_client'] |
| 39 | + |
| 40 | + if options['postgresql_client']['connection'].is_a?(Hash) |
| 41 | + @database = options['postgresql_client']['connection']['database'] |
| 42 | + @schema = options['postgresql_client']['connection']['schema'] || 'public' |
| 43 | + @authentication = options['postgresql_client']['connection']['authentication'] |
| 44 | + @ssl_mode = options['postgresql_client']['connection']['sslMode'] |
| 45 | + raise "SSL Mode should be prefer, require and verify-full" unless @ssl_mode == 'prefer' || @ssl_mode == 'require' || @ssl_mode == 'verify-full' |
| 46 | + |
| 47 | + @url = build_url(options['postgresql_client']['connection']['url']) |
| 48 | + else |
| 49 | + raise('Missing connection info for Postgres client') |
| 50 | + end |
| 51 | + |
| 52 | + Java.org.postgresql.Driver |
| 53 | + end |
| 54 | + |
| 55 | + def realize_query(query, _params) |
| 56 | + GoodData.gd_logger.info("Realize SQL query: type=postgresql status=started") |
| 57 | + |
| 58 | + connect |
| 59 | + filename = "#{SecureRandom.urlsafe_base64(6)}_#{Time.now.to_i}.csv" |
| 60 | + measure = Benchmark.measure do |
| 61 | + statement = @connection.create_statement |
| 62 | + statement.set_fetch_size(POSTGRES_FETCH_SIZE) |
| 63 | + has_result = statement.execute(query) |
| 64 | + if has_result |
| 65 | + result = statement.get_result_set |
| 66 | + metadata = result.get_meta_data |
| 67 | + col_count = metadata.column_count |
| 68 | + CSV.open(filename, 'wb') do |csv| |
| 69 | + csv << Array(1..col_count).map { |i| metadata.get_column_name(i) } # build the header |
| 70 | + csv << Array(1..col_count).map { |i| result.get_string(i)&.to_s } while result.next |
| 71 | + end |
| 72 | + end |
| 73 | + end |
| 74 | + GoodData.gd_logger.info("Realize SQL query: type=postgresql status=finished duration=#{measure.real}") |
| 75 | + filename |
| 76 | + ensure |
| 77 | + @connection&.close |
| 78 | + @connection = nil |
| 79 | + end |
| 80 | + |
| 81 | + def connect |
| 82 | + GoodData.logger.info "Setting up connection to Postgresql #{@url}" |
| 83 | + |
| 84 | + prop = java.util.Properties.new |
| 85 | + prop.setProperty('user', @authentication['basic']['userName']) |
| 86 | + prop.setProperty('password', @authentication['basic']['password']) |
| 87 | + prop.setProperty('schema', @schema) |
| 88 | + |
| 89 | + @connection = java.sql.DriverManager.getConnection(@url, prop) |
| 90 | + statement = @connection.create_statement |
| 91 | + statement.execute("#{POSTGRES_SET_SCHEMA_COMMAND} #{@schema}") |
| 92 | + @connection.set_auto_commit(false) |
| 93 | + end |
| 94 | + |
| 95 | + def build_url(url) |
| 96 | + matches = url.scan(JDBC_POSTGRES_PATTERN) |
| 97 | + raise 'Cannot reach the url' unless matches |
| 98 | + |
| 99 | + host = matches[0][0] |
| 100 | + port = matches[0][2]&.to_i || POSTGRES_DEFAULT_PORT |
| 101 | + raise "Custom port #{port} is not supported. Remove it or use the default port '5432'" if POSTGRES_DEFAULT_PORT != port |
| 102 | + |
| 103 | + "#{JDBC_POSTGRES_PROTOCOL}#{host}:#{port}/#{@database}?sslmode=#{@ssl_mode}#{VERIFY_FULL == @ssl_mode ? SSL_JAVA_FACTORY : ''}" |
| 104 | + end |
| 105 | + end |
| 106 | + end |
| 107 | +end |
0 commit comments