1 | //===----------------------------------------------------------------------===// |
---|---|
2 | // DuckDB |
3 | // |
4 | // duckdb/common/http_state.hpp |
5 | // |
6 | // |
7 | //===----------------------------------------------------------------------===// |
8 | |
9 | #pragma once |
10 | |
11 | #include "duckdb/common/file_opener.hpp" |
12 | #include "duckdb/main/client_context.hpp" |
13 | #include "duckdb/main/client_data.hpp" |
14 | #include "duckdb/common/atomic.hpp" |
15 | #include "duckdb/common/optional_ptr.hpp" |
16 | |
17 | namespace duckdb { |
18 | |
19 | struct CachedFile { |
20 | //! Cached Data |
21 | shared_ptr<char> data; |
22 | //! Data capacity |
23 | uint64_t capacity = 0; |
24 | //! If we finished downloading the file |
25 | bool finished = false; |
26 | }; |
27 | |
28 | class HTTPState { |
29 | public: |
30 | atomic<idx_t> head_count {0}; |
31 | atomic<idx_t> get_count {0}; |
32 | atomic<idx_t> put_count {0}; |
33 | atomic<idx_t> post_count {0}; |
34 | atomic<idx_t> total_bytes_received {0}; |
35 | atomic<idx_t> total_bytes_sent {0}; |
36 | //! Mutex to lock when getting the cached file(Parallel Only) |
37 | mutex cached_files_mutex; |
38 | //! In case of fully downloading the file, the cached files of this query |
39 | unordered_map<string, CachedFile> cached_files; |
40 | |
41 | void Reset() { |
42 | head_count = 0; |
43 | get_count = 0; |
44 | put_count = 0; |
45 | post_count = 0; |
46 | total_bytes_received = 0; |
47 | total_bytes_sent = 0; |
48 | cached_files.clear(); |
49 | } |
50 | |
51 | //! helper function to get the HTTP |
52 | static shared_ptr<HTTPState> TryGetState(FileOpener *opener) { |
53 | auto client_context = FileOpener::TryGetClientContext(opener); |
54 | if (client_context) { |
55 | return client_context->client_data->http_state; |
56 | } |
57 | return nullptr; |
58 | } |
59 | |
60 | bool IsEmpty() { |
61 | return head_count == 0 && get_count == 0 && put_count == 0 && post_count == 0 && total_bytes_received == 0 && |
62 | total_bytes_sent == 0; |
63 | } |
64 | }; |
65 | |
66 | } // namespace duckdb |
67 |