Merge branch 'main' into fix/macro_structs_generator

Author: os-rss

Project.toml

@@ -1,7 +1,7 @@
 name = "Elasticsearch"
 uuid = "e586a49d-aa29-4ce5-8f91-fa4f824367bd"
 authors = ["Egor Shmorgun <egor.shmorgun@opensesame.com>"]
-version = "0.1.5"
+version = "0.2.1"
 
 [deps]
 CodecZlib = "944b1d66-785c-5afd-91f1-9de20f533193"

src/elastic_transport/ElasticTransport.jl

@@ -6,6 +6,7 @@ export perform_request
 include("transport/errors.jl")
 include("transport/connections/Connections.jl")
 include("transport/transport.jl")
+include("transport/sniffing.jl")
 include("client.jl")
 
 end

src/elastic_transport/client.jl

@@ -2,19 +2,24 @@ using HTTP
 using URIs
 using Mocking
 
-const DEFAULT_HOST = "localhost:9200"
-
-struct Client
+const DEFAULT_PORT = 9200
+const DEFAULT_PROTOCOL = "http"
+const DEFAULT_HOST = "localhost"
+const DEFAULT_URL = "$DEFAULT_PROTOCOL://$DEFAULT_HOST:$DEFAULT_PORT"
+const SECURITY_PRIVILEGES_VALIDATION_WARNING = "The client is unable to verify that the server is Elasticsearch due to security privileges on the server side. Some functionality may not be compatible if the server is running an unsupported product."
+const VALIDATION_WARNING = "The client is unable to verify that the server is Elasticsearch. Some functionality may not be compatible if the server is running an unsupported product."
+
+mutable struct Client
   arguments::Dict
   options::Dict
   hosts::Vector
   send_get_body_as::String
-  ca_fingerpring::Bool
+  verified::Bool
   transport::Transport
 end
 
-function Client(arguments::Dict{Symbol,Any}=Dict{Symbol,Any}(); http_client::Module=HTTP)
-  options = deepcopy(arguments)
+function Client(;http_client::Module=HTTP, kwargs...)
+  options = deepcopy(Dict{Symbol, Any}(kwargs))
   arguments = options
 
   get!(options, :reload_connections, false)
@@ -30,7 +35,7 @@ function Client(arguments::Dict{Symbol,Any}=Dict{Symbol,Any}(); http_client::Mod
   hosts_config = if !isnothing(host_key_index)
     arguments[host_keys[host_key_index]]
   else
-    get(ENV, "ELASTICSEARCH_URL", DEFAULT_HOST)
+    get(ENV, "ELASTICSEARCH_URL", DEFAULT_URL)
   end
   hosts = extract_hosts(hosts_config, options)
 
@@ -50,6 +55,37 @@ function Client(arguments::Dict{Symbol,Any}=Dict{Symbol,Any}(); http_client::Mod
   )
 end
 
+function verify_elasticsearch(client::Client)
+  response = nothing
+  try
+    response = elastisearch_validation_request(client)
+  catch exc
+    if typeof(exc) in [Forbidden, Unauthorized, RequestEntityTooLarge]
+      client.verified = true
+      @warn SECURITY_PRIVILEGES_VALIDATION_WARNING
+      return
+    else
+      @warn VALIDATION_WARNING
+      return
+    end
+  end
+
+  body = response.body
+  version = get(() -> Dict(), body, "version") |> version -> get(version, "number", nothing)
+
+  verify_with_version_and_headers(client, version, response.headers)
+end
+
+@warn "Version verification isn't implemented"
+function verify_with_version_and_headers(client::Client, _headers, _version)
+  @warn "Version verification isn't implemented"
+  client.verified = true
+end
+
+function elastisearch_validation_request(client::Client)
+  @mock perform_request(client.transport, "GET", "/")
+end
+
 function perform_request(
   client::Client,
   method::String,
@@ -62,23 +98,22 @@ function perform_request(
     method = client.send_get_body_as
   end
 
-  validate_ca_fingerprints(client)
+  if !client.verified
+    verify_elasticsearch(client)
+  end
 
   @mock perform_request(client.transport, method, path; params=params, body=body, headers=headers)
 end
 
-@warn "ca fingerprints validation is not implemented"
-function validate_ca_fingerprints(::Client)
-  @warn "ca fingerprints validation is not implemented"
-end
-
 function extract_hosts(hosts_config, options)
   hosts = if hosts_config isa String
     split(hosts_config, ",") .|> strip .|> String
   elseif hosts_config isa Vector
     hosts_config
   elseif hosts_config isa Dict || hosts_config isa URI
     [hosts_config]
+  elseif hosts_config isa NamedTuple
+    [Dict(zip(keys(hosts_config), values(hosts_config)))]
   else
     error("Can't extract hosts")
   end
@@ -121,8 +156,8 @@ function parse_host_parts(host::String)
     host_info = split(host, ":")
 
     Dict(
-      :host => get(host_info, 0, ""),
-      :port => get(host_info, 1, "")
+      :host => get(host_info, 0, DEFAULT_HOST),
+      :port => parse(Int16, get(host_info, 1, string(DEFAULT_PORT)))
     )
   end
 
@@ -131,17 +166,22 @@ end
 
 function parse_host_parts(host::URI)
   userinfo = split(host.userinfo, ":") .|> string
+  port = if isempty(host.port)
+    DEFAULT_PORT
+  else
+    parse(Int16, host.port)
+  end
 
   Dict(
     :scheme => host.scheme,
     :user => get(userinfo, 0, ""),
     :password => get(userinfo, 1, ""),
     :host => host.host,
     :path => host.path,
-    :port => host.port
+    :port => port
   )
 end
 
-function parse_host_parts(host::Dict{Symbol,Any})
+function parse_host_parts(host::Union{Dict{Symbol,Any}, NamedTuple})
   host
 end

src/elastic_transport/transport/connections/collection.jl

@@ -24,6 +24,20 @@ function get_connection(collection::Collection)
   conn
 end
 
-function dead(collection::Collection)
-  filter(conn -> conn.dead, collection.connections)
+function remove!(collection::Collection, conn::Connection)
+  index = findfirst(==(conn), collection.connections)
+
+  if !isnothing(index)
+    deleteat!(collection.connections, index)
+  end
 end
+
+dead(collection::Collection) = filter(conn -> conn.dead, collection)
+
+Base.length(collection::Collection) = length(collection.connections)
+Base.push!(collection::Collection, conns::Connection...) = push!(collection.connections, conns...)
+Base.filter(func::Function, collection::Collection) =
+  Collection(filter(func, collection.connections), collection.selector)
+Base.any(func::Function, collection::Collection) = any(func, collection.connections)
+Base.foreach(func, collection::Collection) = foreach(func, collection.connections)
+

src/elastic_transport/transport/connections/connection.jl

@@ -134,3 +134,9 @@ function parse_headers(conn::Connection, headers::Union{Nothing,Dict})
     copy(conn.headers)
   end
 end
+
+function Base.:(==)(src::Connection, other::Connection)
+  src.host[:protocol] == other.host[:protocol] &&
+    src.host[:host] == other.host[:host] &&
+    src.host[:port] == other.host[:port]
+end

src/elastic_transport/transport/sniffing.jl

@@ -0,0 +1,77 @@
+using HTTP
+
+const DEFAULT_SNIFFING_TIMEOUT = 1
+const SNIFFING_PROTOCOL = "http"
+
+struct SniffingTimetoutError <: Exception end
+
+function sniff_hosts(transport::Transport)
+  nodes = perform_sniff_request_with_timeout(transport).body
+  
+  map(collect(nodes["nodes"])) do (id, info)
+    if haskey(info, SNIFFING_PROTOCOL)
+      host, port = parse_publish_address(info[SNIFFING_PROTOCOL]["publish_address"])
+
+      Dict(
+        :id => id,
+        :name => get(info, "name", nothing),
+        :version => get(info, "version", nothing),
+        :host => host,
+        :port => parse(Int16, port),
+        :roles => get(info, "roles", nothing),
+        :attributes => get(info, "attributes", nothing)
+      )
+    else
+      missing
+    end
+  end |> skipmissing |> collect
+end
+
+function parse_publish_address(publish_address::String)
+  if !isnothing(match(r"^inet\[.*\]$", publish_address))
+    parse_address_port(publish_address[begin + 6:end - 1])
+  elseif !isnothing(match(r"/", publish_address))
+    parts = split(publish_address, "/") .|> String
+
+    [parts[begin], parse_address_port(parts[end])[end]]
+  else
+    parse_address_port(publish_address)
+  end
+
+end
+
+function parse_address_port(publish_address::String)
+  # If publish address is ipv6
+  if !isnothing(match(r"[\[\]]", publish_address))
+    parts = match(r"\A\[(.+)\](?::(\d+))?\z", publish_address)
+
+    [parts[1], parts[2]]
+  else
+    split(publish_address, ":")
+  end
+end
+
+function perform_sniff_request_with_timeout(transport::Transport)
+  task = @task(
+    perform_request(
+      transport,
+      "GET",
+      "/_nodes/$SNIFFING_PROTOCOL",
+      opts = Dict(:reload_on_failure => false)
+    )
+  )
+  schedule(task)
+  Timer(DEFAULT_SNIFFING_TIMEOUT) do _timer
+    istaskdone(task) || Base.throwto(task, SniffingTimetoutError())
+  end
+
+  try
+    fetch(task)
+  catch
+    throw(task.exception)
+  end
+end
+
+function sniffing_timeout(transport::Transport)
+  get(transport.options, :sniffing_timeout, DEFAULT_SNIFFING_TIMEOUT)
+end

src/elastic_transport/transport/transport.jl

@@ -5,8 +5,6 @@ using Retry
 using Mocking
 using JSON
 
-const DEFAULT_PORT = 9200
-const DEFAULT_PROTOCOL = "http"
 const DEFAULT_RELOAD_AFTER = 10_000 # Requests
 const DEFAULT_RESURRECT_AFTER = 60  # Seconds
 const DEFAULT_MAX_RETRIES = 3 # Requests
@@ -62,7 +60,7 @@ function get_connection(transport::Transport, options=Dict())
   end
 
   if transport.reload_connections && (transport.counter % transport.reload_after) == 0
-    reload_connections!
+    reload_connections!(transport)
   end
 
   Connections.get_connection(transport.connections)
@@ -72,9 +70,18 @@ function resurrect_dead_connections!(transport::Transport)
   foreach(Connections.resurrect!, Connections.dead(transport.connections))
 end
 
-@warn "Reload connections are not implemented"
+
 function reload_connections!(transport::Transport)
-  @warn "Reload connections are not implemented"
+  try
+    hosts = sniff_hosts(transport)
+    rebuild_connections!(transport, hosts = hosts)
+  catch e
+    if e isa SniffingTimetoutError
+      @error "[SnifferTimeoutError] Timeout when reloading connections."
+    else
+      throw(e)
+    end
+  end
 end
 
 function build_connections(hosts::Vector, options::Dict)
@@ -84,6 +91,23 @@ function build_connections(hosts::Vector, options::Dict)
   )
 end
 
+function rebuild_connections!(transport::Transport; hosts)
+  lock(transport.state_lock) do
+    transport.hosts = hosts
+
+    new_connections = build_connections(hosts, transport.options)
+    stale_connections = filter(transport.connections.connections) do conn
+      !any(new_conn -> new_conn == conn, new_connections)
+    end
+    new_connections = filter(new_connections) do conn
+      !any(new_conn -> new_conn == conn, transport.connections.connections)
+    end
+
+    foreach(conn -> Connections.remove!(transport.connections, conn), stale_connections)
+    push!(transport.connections, new_connections.connections...)
+  end
+end
+
 function connections_from_host(hosts::Vector, options::Dict)
   map(hosts) do host
     host[:protocol] = get(host, :scheme) do
@@ -169,8 +193,7 @@ function perform_request(
     if typeof(exception) in HOST_UNREACHABLE_EXCEPTIONS
       @error "[$(typeof(exception))] $(connection.host)"
 
-      # Disable dead connections, before reload_connections implementation
-      # Connections.dead!(connection)
+      Connections.dead!(connection)
     end
 
     @retry if reload_on_failure && tries < length(transport.connections) && in(typeof(exception), HOST_UNREACHABLE_EXCEPTIONS)