Deploying to gh-pages from @ GreptimeTeam/greptimedb@9a2ac7f9de 🚀

This commit is contained in:
github-merge-queue[bot]
2026-10-20 14:37:30 +00:00
commit a61d66e6f0
26072 changed files with 3030916 additions and 0 deletions
+464
View File
@@ -0,0 +1,464 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/admin.rs`."><title>admin.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>admin.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">mod </span>event;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">mod </span>layer;
<a href=#17 id=17 data-nosnippet>17</a>
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>std::sync::Arc;
<a href=#19 id=19 data-nosnippet>19</a>
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>common_function::function::FunctionContext;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>common_function::function_registry::{FUNCTION_REGISTRY, get_admin_function};
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>common_function::state::FunctionState;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>common_query::Output;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>common_recordbatch::{RecordBatch, RecordBatches};
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>common_sql::convert::sql_value_to_value;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>common_telemetry::tracing;
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>common_time::Timezone;
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>datafusion_expr::TypeSignature;
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span>datatypes::arrow::datatypes::DataType <span class="kw">as </span>ArrowDataType;
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span>datatypes::data_type::DataType;
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span>datatypes::prelude::ConcreteDataType;
<a href=#32 id=32 data-nosnippet>32</a><span class="kw">use </span>datatypes::schema::{ColumnSchema, Schema};
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">use </span>datatypes::value::Value;
<a href=#34 id=34 data-nosnippet>34</a><span class="kw">use </span>datatypes::vectors::VectorRef;
<a href=#35 id=35 data-nosnippet>35</a><span class="kw">pub use </span>layer::{
<a href=#36 id=36 data-nosnippet>36</a> AdminEventRecorderHandle, AdminFunctionLayer, AdminFunctionLayerRef,
<a href=#37 id=37 data-nosnippet>37</a> AdminFunctionRecordingLayer,
<a href=#38 id=38 data-nosnippet>38</a>};
<a href=#39 id=39 data-nosnippet>39</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#40 id=40 data-nosnippet>40</a><span class="kw">use </span>snafu::{OptionExt, ResultExt, ensure};
<a href=#41 id=41 data-nosnippet>41</a><span class="kw">use </span>sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value <span class="kw">as </span>SqlValue};
<a href=#42 id=42 data-nosnippet>42</a><span class="kw">use </span>sql::statements::admin::Admin;
<a href=#43 id=43 data-nosnippet>43</a>
<a href=#44 id=44 data-nosnippet>44</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, CastSnafu, ExecuteAdminFunctionSnafu, <span class="prelude-ty">Result</span>};
<a href=#45 id=45 data-nosnippet>45</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#46 id=46 data-nosnippet>46</a>
<a href=#47 id=47 data-nosnippet>47</a><span class="kw">const </span>DUMMY_COLUMN: <span class="kw-2">&amp;</span>str = <span class="string">"&lt;dummy&gt;"</span>;
<a href=#48 id=48 data-nosnippet>48</a>
<a href=#49 id=49 data-nosnippet>49</a><span class="doccomment">/// A request to execute one ADMIN function statement.
<a href=#50 id=50 data-nosnippet>50</a></span><span class="attr">#[derive(Clone)]
<a href=#51 id=51 data-nosnippet>51</a></span><span class="kw">pub struct </span>AdminFunctionRequest {
<a href=#52 id=52 data-nosnippet>52</a> <span class="doccomment">/// The parsed ADMIN statement to execute.
<a href=#53 id=53 data-nosnippet>53</a> </span><span class="kw">pub </span>statement: Admin,
<a href=#54 id=54 data-nosnippet>54</a> <span class="doccomment">/// The query context of the request.
<a href=#55 id=55 data-nosnippet>55</a> </span><span class="kw">pub </span>query_ctx: QueryContextRef,
<a href=#56 id=56 data-nosnippet>56</a>}
<a href=#57 id=57 data-nosnippet>57</a>
<a href=#58 id=58 data-nosnippet>58</a><span class="doccomment">/// The client output and immediate typed result of one ADMIN function execution.
<a href=#59 id=59 data-nosnippet>59</a></span><span class="kw">pub struct </span>AdminFunctionResponse {
<a href=#60 id=60 data-nosnippet>60</a> <span class="doccomment">/// The output returned to the client.
<a href=#61 id=61 data-nosnippet>61</a> </span><span class="kw">pub </span>output: Output,
<a href=#62 id=62 data-nosnippet>62</a> <span class="doccomment">/// The typed immediate result exposed to outer layers.
<a href=#63 id=63 data-nosnippet>63</a> </span><span class="kw">pub </span>immediate_result: <span class="prelude-ty">Option</span>&lt;Value&gt;,
<a href=#64 id=64 data-nosnippet>64</a>}
<a href=#65 id=65 data-nosnippet>65</a>
<a href=#66 id=66 data-nosnippet>66</a><span class="doccomment">/// Executes an ADMIN function request.
<a href=#67 id=67 data-nosnippet>67</a></span><span class="attr">#[async_trait::async_trait]
<a href=#68 id=68 data-nosnippet>68</a></span><span class="kw">pub trait </span>AdminFunctionService: Send + Sync {
<a href=#69 id=69 data-nosnippet>69</a> <span class="doccomment">/// Executes an ADMIN function request.
<a href=#70 id=70 data-nosnippet>70</a> </span><span class="kw">async fn </span>call(<span class="kw-2">&amp;</span><span class="self">self</span>, request: AdminFunctionRequest) -&gt; <span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt;;
<a href=#71 id=71 data-nosnippet>71</a>}
<a href=#72 id=72 data-nosnippet>72</a>
<a href=#73 id=73 data-nosnippet>73</a><span class="doccomment">/// A shared ADMIN function service.
<a href=#74 id=74 data-nosnippet>74</a></span><span class="kw">pub type </span>AdminFunctionServiceRef = Arc&lt;<span class="kw">dyn </span>AdminFunctionService&gt;;
<a href=#75 id=75 data-nosnippet>75</a>
<a href=#76 id=76 data-nosnippet>76</a><span class="attr">#[derive(Clone)]
<a href=#77 id=77 data-nosnippet>77</a></span><span class="kw">struct </span>CoreAdminFunctionService {
<a href=#78 id=78 data-nosnippet>78</a> query_engine: query::QueryEngineRef,
<a href=#79 id=79 data-nosnippet>79</a>}
<a href=#80 id=80 data-nosnippet>80</a>
<a href=#81 id=81 data-nosnippet>81</a><span class="doccomment">/// Parts of an `ADMIN` call needed both for execution and schema derivation.
<a href=#82 id=82 data-nosnippet>82</a></span><span class="kw">struct </span>ResolvedAdminFunction {
<a href=#83 id=83 data-nosnippet>83</a> admin_udf: datafusion_expr::ScalarUDF,
<a href=#84 id=84 data-nosnippet>84</a> fn_name: String,
<a href=#85 id=85 data-nosnippet>85</a> args: Vec&lt;VectorRef&gt;,
<a href=#86 id=86 data-nosnippet>86</a> arg_types: Vec&lt;ArrowDataType&gt;,
<a href=#87 id=87 data-nosnippet>87</a> ret_type: ArrowDataType,
<a href=#88 id=88 data-nosnippet>88</a>}
<a href=#89 id=89 data-nosnippet>89</a>
<a href=#90 id=90 data-nosnippet>90</a><span class="doccomment">/// Resolves the function, parses its literal arguments and derives its
<a href=#91 id=91 data-nosnippet>91</a>/// return type, without executing it.
<a href=#92 id=92 data-nosnippet>92</a></span><span class="kw">fn </span>resolve_admin_function(
<a href=#93 id=93 data-nosnippet>93</a> stmt: <span class="kw-2">&amp;</span>Admin,
<a href=#94 id=94 data-nosnippet>94</a> query_ctx: <span class="kw-2">&amp;</span>QueryContextRef,
<a href=#95 id=95 data-nosnippet>95</a> state: Arc&lt;FunctionState&gt;,
<a href=#96 id=96 data-nosnippet>96</a>) -&gt; <span class="prelude-ty">Result</span>&lt;ResolvedAdminFunction&gt; {
<a href=#97 id=97 data-nosnippet>97</a> <span class="kw">let </span>Admin::Func(func) = stmt;
<a href=#98 id=98 data-nosnippet>98</a> <span class="comment">// the function name should be in lower case.
<a href=#99 id=99 data-nosnippet>99</a> </span><span class="kw">let </span>func_name = func.name.to_string().to_lowercase();
<a href=#100 id=100 data-nosnippet>100</a> <span class="kw">let </span>factory = get_admin_function(<span class="kw-2">&amp;</span>func_name)
<a href=#101 id=101 data-nosnippet>101</a> .or_else(|| FUNCTION_REGISTRY.get_function(<span class="kw-2">&amp;</span>func_name))
<a href=#102 id=102 data-nosnippet>102</a> .context(error::AdminFunctionNotFoundSnafu {
<a href=#103 id=103 data-nosnippet>103</a> name: func_name.clone(),
<a href=#104 id=104 data-nosnippet>104</a> })<span class="question-mark">?</span>;
<a href=#105 id=105 data-nosnippet>105</a>
<a href=#106 id=106 data-nosnippet>106</a> <span class="kw">let </span>func_ctx = FunctionContext {
<a href=#107 id=107 data-nosnippet>107</a> query_ctx: query_ctx.clone(),
<a href=#108 id=108 data-nosnippet>108</a> state,
<a href=#109 id=109 data-nosnippet>109</a> };
<a href=#110 id=110 data-nosnippet>110</a>
<a href=#111 id=111 data-nosnippet>111</a> <span class="kw">let </span>admin_udf = factory.provide(func_ctx);
<a href=#112 id=112 data-nosnippet>112</a> admin_udf
<a href=#113 id=113 data-nosnippet>113</a> .as_async()
<a href=#114 id=114 data-nosnippet>114</a> .context(error::AdminFunctionNotFoundSnafu { name: func_name })<span class="question-mark">?</span>;
<a href=#115 id=115 data-nosnippet>115</a>
<a href=#116 id=116 data-nosnippet>116</a> <span class="kw">let </span>fn_name = admin_udf.name().to_string();
<a href=#117 id=117 data-nosnippet>117</a> <span class="kw">let </span>signature = admin_udf.signature();
<a href=#118 id=118 data-nosnippet>118</a>
<a href=#119 id=119 data-nosnippet>119</a> <span class="comment">// Parse function arguments
<a href=#120 id=120 data-nosnippet>120</a> </span><span class="kw">let </span>FunctionArguments::List(args) = <span class="kw-2">&amp;</span>func.args <span class="kw">else </span>{
<a href=#121 id=121 data-nosnippet>121</a> <span class="kw">return </span>error::BuildAdminFunctionArgsSnafu {
<a href=#122 id=122 data-nosnippet>122</a> msg: <span class="macro">format!</span>(<span class="string">"unsupported function args {} for {}"</span>, func.args, fn_name),
<a href=#123 id=123 data-nosnippet>123</a> }
<a href=#124 id=124 data-nosnippet>124</a> .fail();
<a href=#125 id=125 data-nosnippet>125</a> };
<a href=#126 id=126 data-nosnippet>126</a> <span class="kw">let </span>arg_values = args
<a href=#127 id=127 data-nosnippet>127</a> .args
<a href=#128 id=128 data-nosnippet>128</a> .iter()
<a href=#129 id=129 data-nosnippet>129</a> .map(|arg| {
<a href=#130 id=130 data-nosnippet>130</a> <span class="kw">let </span>FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) = arg <span class="kw">else </span>{
<a href=#131 id=131 data-nosnippet>131</a> <span class="kw">return </span>error::BuildAdminFunctionArgsSnafu {
<a href=#132 id=132 data-nosnippet>132</a> msg: <span class="macro">format!</span>(<span class="string">"unsupported function arg {arg} for {}"</span>, fn_name),
<a href=#133 id=133 data-nosnippet>133</a> }
<a href=#134 id=134 data-nosnippet>134</a> .fail();
<a href=#135 id=135 data-nosnippet>135</a> };
<a href=#136 id=136 data-nosnippet>136</a> <span class="prelude-val">Ok</span>(<span class="kw-2">&amp;</span>value.value)
<a href=#137 id=137 data-nosnippet>137</a> })
<a href=#138 id=138 data-nosnippet>138</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#139 id=139 data-nosnippet>139</a>
<a href=#140 id=140 data-nosnippet>140</a> <span class="kw">let </span>args = args_to_vector(<span class="kw-2">&amp;</span>signature.type_signature, <span class="kw-2">&amp;</span>arg_values, query_ctx)<span class="question-mark">?</span>;
<a href=#141 id=141 data-nosnippet>141</a> <span class="kw">let </span>arg_types = args
<a href=#142 id=142 data-nosnippet>142</a> .iter()
<a href=#143 id=143 data-nosnippet>143</a> .map(|arg| arg.data_type().as_arrow_type())
<a href=#144 id=144 data-nosnippet>144</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#145 id=145 data-nosnippet>145</a> <span class="kw">let </span>ret_type =
<a href=#146 id=146 data-nosnippet>146</a> admin_udf
<a href=#147 id=147 data-nosnippet>147</a> .return_type(<span class="kw-2">&amp;</span>arg_types)
<a href=#148 id=148 data-nosnippet>148</a> .map_err(|e| error::Error::BuildAdminFunctionArgs {
<a href=#149 id=149 data-nosnippet>149</a> msg: <span class="macro">format!</span>(
<a href=#150 id=150 data-nosnippet>150</a> <span class="string">"Failed to get return type of admin function {}: {}"</span>,
<a href=#151 id=151 data-nosnippet>151</a> fn_name, e
<a href=#152 id=152 data-nosnippet>152</a> ),
<a href=#153 id=153 data-nosnippet>153</a> })<span class="question-mark">?</span>;
<a href=#154 id=154 data-nosnippet>154</a>
<a href=#155 id=155 data-nosnippet>155</a> <span class="prelude-val">Ok</span>(ResolvedAdminFunction {
<a href=#156 id=156 data-nosnippet>156</a> admin_udf,
<a href=#157 id=157 data-nosnippet>157</a> fn_name,
<a href=#158 id=158 data-nosnippet>158</a> args,
<a href=#159 id=159 data-nosnippet>159</a> arg_types,
<a href=#160 id=160 data-nosnippet>160</a> ret_type,
<a href=#161 id=161 data-nosnippet>161</a> })
<a href=#162 id=162 data-nosnippet>162</a>}
<a href=#163 id=163 data-nosnippet>163</a>
<a href=#164 id=164 data-nosnippet>164</a><span class="doccomment">/// Output schema of an `ADMIN` statement, mirroring what
<a href=#165 id=165 data-nosnippet>165</a>/// [`CoreAdminFunctionService::execute`] produces. `None` if the function
<a href=#166 id=166 data-nosnippet>166</a>/// or arguments are unresolvable; execution will then surface the error.
<a href=#167 id=167 data-nosnippet>167</a></span><span class="kw">pub fn </span>admin_output_schema(stmt: <span class="kw-2">&amp;</span>Admin, query_ctx: <span class="kw-2">&amp;</span>QueryContextRef) -&gt; <span class="prelude-ty">Option</span>&lt;Schema&gt; {
<a href=#168 id=168 data-nosnippet>168</a> <span class="kw">let </span>resolved =
<a href=#169 id=169 data-nosnippet>169</a> resolve_admin_function(stmt, query_ctx, Arc::new(FunctionState::default())).ok()<span class="question-mark">?</span>;
<a href=#170 id=170 data-nosnippet>170</a> <span class="prelude-val">Some</span>(Schema::new(<span class="macro">vec!</span>[ColumnSchema::new(
<a href=#171 id=171 data-nosnippet>171</a> <span class="comment">// Use statement as the result column name
<a href=#172 id=172 data-nosnippet>172</a> </span>stmt.to_string(),
<a href=#173 id=173 data-nosnippet>173</a> ConcreteDataType::from_arrow_type(<span class="kw-2">&amp;</span>resolved.ret_type),
<a href=#174 id=174 data-nosnippet>174</a> <span class="bool-val">false</span>,
<a href=#175 id=175 data-nosnippet>175</a> )]))
<a href=#176 id=176 data-nosnippet>176</a>}
<a href=#177 id=177 data-nosnippet>177</a>
<a href=#178 id=178 data-nosnippet>178</a><span class="kw">impl </span>CoreAdminFunctionService {
<a href=#179 id=179 data-nosnippet>179</a> <span class="kw">fn </span>new(query_engine: query::QueryEngineRef) -&gt; <span class="self">Self </span>{
<a href=#180 id=180 data-nosnippet>180</a> <span class="self">Self </span>{ query_engine }
<a href=#181 id=181 data-nosnippet>181</a> }
<a href=#182 id=182 data-nosnippet>182</a>
<a href=#183 id=183 data-nosnippet>183</a> <span class="kw">async fn </span>execute(<span class="kw-2">&amp;</span><span class="self">self</span>, request: AdminFunctionRequest) -&gt; <span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt; {
<a href=#184 id=184 data-nosnippet>184</a> <span class="kw">let </span>AdminFunctionRequest {
<a href=#185 id=185 data-nosnippet>185</a> statement: stmt,
<a href=#186 id=186 data-nosnippet>186</a> query_ctx,
<a href=#187 id=187 data-nosnippet>187</a> } = request;
<a href=#188 id=188 data-nosnippet>188</a>
<a href=#189 id=189 data-nosnippet>189</a> <span class="kw">let </span>resolved = resolve_admin_function(
<a href=#190 id=190 data-nosnippet>190</a> <span class="kw-2">&amp;</span>stmt,
<a href=#191 id=191 data-nosnippet>191</a> <span class="kw-2">&amp;</span>query_ctx,
<a href=#192 id=192 data-nosnippet>192</a> <span class="self">self</span>.query_engine.engine_state().function_state(),
<a href=#193 id=193 data-nosnippet>193</a> )<span class="question-mark">?</span>;
<a href=#194 id=194 data-nosnippet>194</a> <span class="kw">let </span>ResolvedAdminFunction {
<a href=#195 id=195 data-nosnippet>195</a> admin_udf,
<a href=#196 id=196 data-nosnippet>196</a> fn_name,
<a href=#197 id=197 data-nosnippet>197</a> args,
<a href=#198 id=198 data-nosnippet>198</a> arg_types,
<a href=#199 id=199 data-nosnippet>199</a> ret_type,
<a href=#200 id=200 data-nosnippet>200</a> } = resolved;
<a href=#201 id=201 data-nosnippet>201</a> <span class="kw">let </span>admin_async_fn = admin_udf
<a href=#202 id=202 data-nosnippet>202</a> .as_async()
<a href=#203 id=203 data-nosnippet>203</a> .context(error::AdminFunctionNotFoundSnafu {
<a href=#204 id=204 data-nosnippet>204</a> name: fn_name.clone(),
<a href=#205 id=205 data-nosnippet>205</a> })<span class="question-mark">?</span>;
<a href=#206 id=206 data-nosnippet>206</a>
<a href=#207 id=207 data-nosnippet>207</a> <span class="comment">// Convert arguments to DataFusion ColumnarValue format
<a href=#208 id=208 data-nosnippet>208</a> </span><span class="kw">let </span>columnar_args: Vec&lt;datafusion_expr::ColumnarValue&gt; = args
<a href=#209 id=209 data-nosnippet>209</a> .iter()
<a href=#210 id=210 data-nosnippet>210</a> .map(|vector| datafusion_expr::ColumnarValue::Array(vector.to_arrow_array()))
<a href=#211 id=211 data-nosnippet>211</a> .collect();
<a href=#212 id=212 data-nosnippet>212</a>
<a href=#213 id=213 data-nosnippet>213</a> <span class="comment">// Create ScalarFunctionArgs following the same pattern as udf.rs
<a href=#214 id=214 data-nosnippet>214</a> </span><span class="kw">let </span>func_args = datafusion::logical_expr::ScalarFunctionArgs {
<a href=#215 id=215 data-nosnippet>215</a> args: columnar_args,
<a href=#216 id=216 data-nosnippet>216</a> arg_fields: args
<a href=#217 id=217 data-nosnippet>217</a> .iter()
<a href=#218 id=218 data-nosnippet>218</a> .enumerate()
<a href=#219 id=219 data-nosnippet>219</a> .map(|(i, vector)| {
<a href=#220 id=220 data-nosnippet>220</a> Arc::new(arrow::datatypes::Field::new(
<a href=#221 id=221 data-nosnippet>221</a> <span class="macro">format!</span>(<span class="string">"arg_{}"</span>, i),
<a href=#222 id=222 data-nosnippet>222</a> arg_types[i].clone(),
<a href=#223 id=223 data-nosnippet>223</a> vector.null_count() &gt; <span class="number">0</span>,
<a href=#224 id=224 data-nosnippet>224</a> ))
<a href=#225 id=225 data-nosnippet>225</a> })
<a href=#226 id=226 data-nosnippet>226</a> .collect(),
<a href=#227 id=227 data-nosnippet>227</a> return_field: Arc::new(arrow::datatypes::Field::new(<span class="string">"result"</span>, ret_type, <span class="bool-val">true</span>)),
<a href=#228 id=228 data-nosnippet>228</a> number_rows: <span class="kw">if </span>args.is_empty() { <span class="number">1 </span>} <span class="kw">else </span>{ args[<span class="number">0</span>].len() },
<a href=#229 id=229 data-nosnippet>229</a> config_options: Arc::new(query_ctx.create_config_options()),
<a href=#230 id=230 data-nosnippet>230</a> };
<a href=#231 id=231 data-nosnippet>231</a>
<a href=#232 id=232 data-nosnippet>232</a> <span class="comment">// Execute the async UDF
<a href=#233 id=233 data-nosnippet>233</a> </span><span class="kw">let </span>result_columnar = admin_async_fn
<a href=#234 id=234 data-nosnippet>234</a> .invoke_async_with_args(func_args)
<a href=#235 id=235 data-nosnippet>235</a> .<span class="kw">await
<a href=#236 id=236 data-nosnippet>236</a> </span>.with_context(|<span class="kw">_</span>| ExecuteAdminFunctionSnafu { msg: fn_name })<span class="question-mark">?</span>;
<a href=#237 id=237 data-nosnippet>237</a>
<a href=#238 id=238 data-nosnippet>238</a> <span class="comment">// Convert result back to VectorRef
<a href=#239 id=239 data-nosnippet>239</a> </span><span class="kw">let </span>result_columnar: common_query::prelude::ColumnarValue =
<a href=#240 id=240 data-nosnippet>240</a> (<span class="kw-2">&amp;</span>result_columnar).try_into().context(CastSnafu)<span class="question-mark">?</span>;
<a href=#241 id=241 data-nosnippet>241</a>
<a href=#242 id=242 data-nosnippet>242</a> <span class="kw">let </span>result_vector: VectorRef = result_columnar.try_into_vector(<span class="number">1</span>).context(CastSnafu)<span class="question-mark">?</span>;
<a href=#243 id=243 data-nosnippet>243</a> <span class="kw">let </span>immediate_result = immediate_result(<span class="kw-2">&amp;</span>result_vector);
<a href=#244 id=244 data-nosnippet>244</a>
<a href=#245 id=245 data-nosnippet>245</a> <span class="kw">let </span>column_schemas = <span class="macro">vec!</span>[ColumnSchema::new(
<a href=#246 id=246 data-nosnippet>246</a> <span class="comment">// Use statement as the result column name
<a href=#247 id=247 data-nosnippet>247</a> </span>stmt.to_string(),
<a href=#248 id=248 data-nosnippet>248</a> result_vector.data_type(),
<a href=#249 id=249 data-nosnippet>249</a> <span class="bool-val">false</span>,
<a href=#250 id=250 data-nosnippet>250</a> )];
<a href=#251 id=251 data-nosnippet>251</a> <span class="kw">let </span>schema = Arc::new(Schema::new(column_schemas));
<a href=#252 id=252 data-nosnippet>252</a> <span class="kw">let </span>batch = RecordBatch::new(schema.clone(), <span class="macro">vec!</span>[result_vector])
<a href=#253 id=253 data-nosnippet>253</a> .context(error::BuildRecordBatchSnafu)<span class="question-mark">?</span>;
<a href=#254 id=254 data-nosnippet>254</a> <span class="kw">let </span>batches =
<a href=#255 id=255 data-nosnippet>255</a> RecordBatches::try_new(schema, <span class="macro">vec!</span>[batch]).context(error::BuildRecordBatchSnafu)<span class="question-mark">?</span>;
<a href=#256 id=256 data-nosnippet>256</a>
<a href=#257 id=257 data-nosnippet>257</a> <span class="prelude-val">Ok</span>(AdminFunctionResponse {
<a href=#258 id=258 data-nosnippet>258</a> output: Output::new_with_record_batches(batches),
<a href=#259 id=259 data-nosnippet>259</a> immediate_result,
<a href=#260 id=260 data-nosnippet>260</a> })
<a href=#261 id=261 data-nosnippet>261</a> }
<a href=#262 id=262 data-nosnippet>262</a>}
<a href=#263 id=263 data-nosnippet>263</a>
<a href=#264 id=264 data-nosnippet>264</a><span class="kw">fn </span>immediate_result(result_vector: <span class="kw-2">&amp;</span>VectorRef) -&gt; <span class="prelude-ty">Option</span>&lt;Value&gt; {
<a href=#265 id=265 data-nosnippet>265</a> (!result_vector.is_empty()).then(|| result_vector.get(<span class="number">0</span>))
<a href=#266 id=266 data-nosnippet>266</a>}
<a href=#267 id=267 data-nosnippet>267</a>
<a href=#268 id=268 data-nosnippet>268</a><span class="attr">#[async_trait::async_trait]
<a href=#269 id=269 data-nosnippet>269</a></span><span class="kw">impl </span>AdminFunctionService <span class="kw">for </span>CoreAdminFunctionService {
<a href=#270 id=270 data-nosnippet>270</a> <span class="kw">async fn </span>call(<span class="kw-2">&amp;</span><span class="self">self</span>, request: AdminFunctionRequest) -&gt; <span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt; {
<a href=#271 id=271 data-nosnippet>271</a> <span class="self">self</span>.execute(request).<span class="kw">await
<a href=#272 id=272 data-nosnippet>272</a> </span>}
<a href=#273 id=273 data-nosnippet>273</a>}
<a href=#274 id=274 data-nosnippet>274</a>
<a href=#275 id=275 data-nosnippet>275</a><span class="doccomment">/// Creates the core ADMIN function service.
<a href=#276 id=276 data-nosnippet>276</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>new_admin_function_service(
<a href=#277 id=277 data-nosnippet>277</a> query_engine: query::QueryEngineRef,
<a href=#278 id=278 data-nosnippet>278</a>) -&gt; AdminFunctionServiceRef {
<a href=#279 id=279 data-nosnippet>279</a> Arc::new(CoreAdminFunctionService::new(query_engine))
<a href=#280 id=280 data-nosnippet>280</a>}
<a href=#281 id=281 data-nosnippet>281</a>
<a href=#282 id=282 data-nosnippet>282</a><span class="kw">impl </span>StatementExecutor {
<a href=#283 id=283 data-nosnippet>283</a> <span class="doccomment">/// Executes the [`Admin`] statement and returns the output.
<a href=#284 id=284 data-nosnippet>284</a> </span><span class="attr">#[tracing::instrument(skip_all)]
<a href=#285 id=285 data-nosnippet>285</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>execute_admin_command(
<a href=#286 id=286 data-nosnippet>286</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#287 id=287 data-nosnippet>287</a> stmt: Admin,
<a href=#288 id=288 data-nosnippet>288</a> query_ctx: QueryContextRef,
<a href=#289 id=289 data-nosnippet>289</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#290 id=290 data-nosnippet>290</a> <span class="self">self</span>.admin_function_service
<a href=#291 id=291 data-nosnippet>291</a> .call(AdminFunctionRequest {
<a href=#292 id=292 data-nosnippet>292</a> statement: stmt,
<a href=#293 id=293 data-nosnippet>293</a> query_ctx,
<a href=#294 id=294 data-nosnippet>294</a> })
<a href=#295 id=295 data-nosnippet>295</a> .<span class="kw">await
<a href=#296 id=296 data-nosnippet>296</a> </span>.map(|response| response.output)
<a href=#297 id=297 data-nosnippet>297</a> }
<a href=#298 id=298 data-nosnippet>298</a>}
<a href=#299 id=299 data-nosnippet>299</a>
<a href=#300 id=300 data-nosnippet>300</a><span class="doccomment">/// Try to cast the arguments to vectors by function's signature.
<a href=#301 id=301 data-nosnippet>301</a></span><span class="kw">fn </span>args_to_vector(
<a href=#302 id=302 data-nosnippet>302</a> type_signature: <span class="kw-2">&amp;</span>TypeSignature,
<a href=#303 id=303 data-nosnippet>303</a> args: <span class="kw-2">&amp;</span>Vec&lt;<span class="kw-2">&amp;</span>SqlValue&gt;,
<a href=#304 id=304 data-nosnippet>304</a> query_ctx: <span class="kw-2">&amp;</span>QueryContextRef,
<a href=#305 id=305 data-nosnippet>305</a>) -&gt; <span class="prelude-ty">Result</span>&lt;Vec&lt;VectorRef&gt;&gt; {
<a href=#306 id=306 data-nosnippet>306</a> <span class="kw">let </span>tz = query_ctx.timezone();
<a href=#307 id=307 data-nosnippet>307</a>
<a href=#308 id=308 data-nosnippet>308</a> <span class="kw">match </span>type_signature {
<a href=#309 id=309 data-nosnippet>309</a> TypeSignature::Variadic(valid_types) =&gt; {
<a href=#310 id=310 data-nosnippet>310</a> values_to_vectors_by_valid_types(valid_types, args, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>tz))
<a href=#311 id=311 data-nosnippet>311</a> }
<a href=#312 id=312 data-nosnippet>312</a>
<a href=#313 id=313 data-nosnippet>313</a> TypeSignature::Uniform(arity, valid_types) =&gt; {
<a href=#314 id=314 data-nosnippet>314</a> <span class="macro">ensure!</span>(
<a href=#315 id=315 data-nosnippet>315</a> <span class="kw-2">*</span>arity == args.len(),
<a href=#316 id=316 data-nosnippet>316</a> error::FunctionArityMismatchSnafu {
<a href=#317 id=317 data-nosnippet>317</a> actual: args.len(),
<a href=#318 id=318 data-nosnippet>318</a> expected: <span class="kw-2">*</span>arity,
<a href=#319 id=319 data-nosnippet>319</a> }
<a href=#320 id=320 data-nosnippet>320</a> );
<a href=#321 id=321 data-nosnippet>321</a>
<a href=#322 id=322 data-nosnippet>322</a> values_to_vectors_by_valid_types(valid_types, args, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>tz))
<a href=#323 id=323 data-nosnippet>323</a> }
<a href=#324 id=324 data-nosnippet>324</a>
<a href=#325 id=325 data-nosnippet>325</a> TypeSignature::Exact(data_types) =&gt; {
<a href=#326 id=326 data-nosnippet>326</a> values_to_vectors_by_exact_types(data_types, args, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>tz))
<a href=#327 id=327 data-nosnippet>327</a> }
<a href=#328 id=328 data-nosnippet>328</a>
<a href=#329 id=329 data-nosnippet>329</a> TypeSignature::VariadicAny =&gt; {
<a href=#330 id=330 data-nosnippet>330</a> <span class="kw">let </span>data_types = args
<a href=#331 id=331 data-nosnippet>331</a> .iter()
<a href=#332 id=332 data-nosnippet>332</a> .map(|value| try_get_data_type_for_sql_value(value))
<a href=#333 id=333 data-nosnippet>333</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#334 id=334 data-nosnippet>334</a>
<a href=#335 id=335 data-nosnippet>335</a> values_to_vectors_by_exact_types(<span class="kw-2">&amp;</span>data_types, args, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>tz))
<a href=#336 id=336 data-nosnippet>336</a> }
<a href=#337 id=337 data-nosnippet>337</a>
<a href=#338 id=338 data-nosnippet>338</a> TypeSignature::Any(arity) =&gt; {
<a href=#339 id=339 data-nosnippet>339</a> <span class="macro">ensure!</span>(
<a href=#340 id=340 data-nosnippet>340</a> <span class="kw-2">*</span>arity == args.len(),
<a href=#341 id=341 data-nosnippet>341</a> error::FunctionArityMismatchSnafu {
<a href=#342 id=342 data-nosnippet>342</a> actual: args.len(),
<a href=#343 id=343 data-nosnippet>343</a> expected: <span class="kw-2">*</span>arity,
<a href=#344 id=344 data-nosnippet>344</a> }
<a href=#345 id=345 data-nosnippet>345</a> );
<a href=#346 id=346 data-nosnippet>346</a>
<a href=#347 id=347 data-nosnippet>347</a> <span class="kw">let </span>data_types = args
<a href=#348 id=348 data-nosnippet>348</a> .iter()
<a href=#349 id=349 data-nosnippet>349</a> .map(|value| try_get_data_type_for_sql_value(value))
<a href=#350 id=350 data-nosnippet>350</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#351 id=351 data-nosnippet>351</a>
<a href=#352 id=352 data-nosnippet>352</a> values_to_vectors_by_exact_types(<span class="kw-2">&amp;</span>data_types, args, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>tz))
<a href=#353 id=353 data-nosnippet>353</a> }
<a href=#354 id=354 data-nosnippet>354</a>
<a href=#355 id=355 data-nosnippet>355</a> TypeSignature::OneOf(type_sigs) =&gt; {
<a href=#356 id=356 data-nosnippet>356</a> <span class="kw">for </span>type_sig <span class="kw">in </span>type_sigs {
<a href=#357 id=357 data-nosnippet>357</a> <span class="kw">if let </span><span class="prelude-val">Ok</span>(vectors) = args_to_vector(type_sig, args, query_ctx) {
<a href=#358 id=358 data-nosnippet>358</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(vectors);
<a href=#359 id=359 data-nosnippet>359</a> }
<a href=#360 id=360 data-nosnippet>360</a> }
<a href=#361 id=361 data-nosnippet>361</a>
<a href=#362 id=362 data-nosnippet>362</a> error::BuildAdminFunctionArgsSnafu {
<a href=#363 id=363 data-nosnippet>363</a> msg: <span class="string">"function signature not match"</span>,
<a href=#364 id=364 data-nosnippet>364</a> }
<a href=#365 id=365 data-nosnippet>365</a> .fail()
<a href=#366 id=366 data-nosnippet>366</a> }
<a href=#367 id=367 data-nosnippet>367</a>
<a href=#368 id=368 data-nosnippet>368</a> <span class="kw">_ </span>=&gt; error::BuildAdminFunctionArgsSnafu {
<a href=#369 id=369 data-nosnippet>369</a> msg: <span class="macro">format!</span>(<span class="string">"unknown function type signature: {type_signature:?}"</span>),
<a href=#370 id=370 data-nosnippet>370</a> }
<a href=#371 id=371 data-nosnippet>371</a> .fail(),
<a href=#372 id=372 data-nosnippet>372</a> }
<a href=#373 id=373 data-nosnippet>373</a>}
<a href=#374 id=374 data-nosnippet>374</a>
<a href=#375 id=375 data-nosnippet>375</a><span class="doccomment">/// Try to cast sql values to vectors by exact data types.
<a href=#376 id=376 data-nosnippet>376</a></span><span class="kw">fn </span>values_to_vectors_by_exact_types(
<a href=#377 id=377 data-nosnippet>377</a> exact_types: <span class="kw-2">&amp;</span>[ArrowDataType],
<a href=#378 id=378 data-nosnippet>378</a> args: <span class="kw-2">&amp;</span>[<span class="kw-2">&amp;</span>SqlValue],
<a href=#379 id=379 data-nosnippet>379</a> tz: <span class="prelude-ty">Option</span>&lt;<span class="kw-2">&amp;</span>Timezone&gt;,
<a href=#380 id=380 data-nosnippet>380</a>) -&gt; <span class="prelude-ty">Result</span>&lt;Vec&lt;VectorRef&gt;&gt; {
<a href=#381 id=381 data-nosnippet>381</a> args.iter()
<a href=#382 id=382 data-nosnippet>382</a> .zip(exact_types.iter())
<a href=#383 id=383 data-nosnippet>383</a> .map(|(value, data_type)| {
<a href=#384 id=384 data-nosnippet>384</a> <span class="kw">let </span>schema = ColumnSchema::new(
<a href=#385 id=385 data-nosnippet>385</a> DUMMY_COLUMN,
<a href=#386 id=386 data-nosnippet>386</a> ConcreteDataType::from_arrow_type(data_type),
<a href=#387 id=387 data-nosnippet>387</a> <span class="bool-val">true</span>,
<a href=#388 id=388 data-nosnippet>388</a> );
<a href=#389 id=389 data-nosnippet>389</a> <span class="kw">let </span>value = sql_value_to_value(<span class="kw-2">&amp;</span>schema, value, tz, <span class="prelude-val">None</span>, <span class="bool-val">false</span>)
<a href=#390 id=390 data-nosnippet>390</a> .context(error::SqlCommonSnafu)<span class="question-mark">?</span>;
<a href=#391 id=391 data-nosnippet>391</a>
<a href=#392 id=392 data-nosnippet>392</a> <span class="prelude-val">Ok</span>(value_to_vector(value))
<a href=#393 id=393 data-nosnippet>393</a> })
<a href=#394 id=394 data-nosnippet>394</a> .collect()
<a href=#395 id=395 data-nosnippet>395</a>}
<a href=#396 id=396 data-nosnippet>396</a>
<a href=#397 id=397 data-nosnippet>397</a><span class="doccomment">/// Try to cast sql values to vectors by valid data types.
<a href=#398 id=398 data-nosnippet>398</a></span><span class="kw">fn </span>values_to_vectors_by_valid_types(
<a href=#399 id=399 data-nosnippet>399</a> valid_types: <span class="kw-2">&amp;</span>[ArrowDataType],
<a href=#400 id=400 data-nosnippet>400</a> args: <span class="kw-2">&amp;</span>[<span class="kw-2">&amp;</span>SqlValue],
<a href=#401 id=401 data-nosnippet>401</a> tz: <span class="prelude-ty">Option</span>&lt;<span class="kw-2">&amp;</span>Timezone&gt;,
<a href=#402 id=402 data-nosnippet>402</a>) -&gt; <span class="prelude-ty">Result</span>&lt;Vec&lt;VectorRef&gt;&gt; {
<a href=#403 id=403 data-nosnippet>403</a> args.iter()
<a href=#404 id=404 data-nosnippet>404</a> .map(|value| {
<a href=#405 id=405 data-nosnippet>405</a> <span class="kw">for </span>data_type <span class="kw">in </span>valid_types {
<a href=#406 id=406 data-nosnippet>406</a> <span class="kw">let </span>schema = ColumnSchema::new(
<a href=#407 id=407 data-nosnippet>407</a> DUMMY_COLUMN,
<a href=#408 id=408 data-nosnippet>408</a> ConcreteDataType::from_arrow_type(data_type),
<a href=#409 id=409 data-nosnippet>409</a> <span class="bool-val">true</span>,
<a href=#410 id=410 data-nosnippet>410</a> );
<a href=#411 id=411 data-nosnippet>411</a> <span class="kw">if let </span><span class="prelude-val">Ok</span>(value) = sql_value_to_value(<span class="kw-2">&amp;</span>schema, value, tz, <span class="prelude-val">None</span>, <span class="bool-val">false</span>) {
<a href=#412 id=412 data-nosnippet>412</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(value_to_vector(value));
<a href=#413 id=413 data-nosnippet>413</a> }
<a href=#414 id=414 data-nosnippet>414</a> }
<a href=#415 id=415 data-nosnippet>415</a>
<a href=#416 id=416 data-nosnippet>416</a> error::BuildAdminFunctionArgsSnafu {
<a href=#417 id=417 data-nosnippet>417</a> msg: <span class="macro">format!</span>(<span class="string">"failed to cast {value}"</span>),
<a href=#418 id=418 data-nosnippet>418</a> }
<a href=#419 id=419 data-nosnippet>419</a> .fail()
<a href=#420 id=420 data-nosnippet>420</a> })
<a href=#421 id=421 data-nosnippet>421</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()
<a href=#422 id=422 data-nosnippet>422</a>}
<a href=#423 id=423 data-nosnippet>423</a>
<a href=#424 id=424 data-nosnippet>424</a><span class="doccomment">/// Build a [`VectorRef`] from [`Value`]
<a href=#425 id=425 data-nosnippet>425</a></span><span class="kw">fn </span>value_to_vector(value: Value) -&gt; VectorRef {
<a href=#426 id=426 data-nosnippet>426</a> <span class="kw">let </span>data_type = value.data_type();
<a href=#427 id=427 data-nosnippet>427</a> <span class="kw">let </span><span class="kw-2">mut </span>mutable_vector = data_type.create_mutable_vector(<span class="number">1</span>);
<a href=#428 id=428 data-nosnippet>428</a> mutable_vector.push_value_ref(<span class="kw-2">&amp;</span>value.as_value_ref());
<a href=#429 id=429 data-nosnippet>429</a>
<a href=#430 id=430 data-nosnippet>430</a> mutable_vector.to_vector()
<a href=#431 id=431 data-nosnippet>431</a>}
<a href=#432 id=432 data-nosnippet>432</a>
<a href=#433 id=433 data-nosnippet>433</a><span class="doccomment">/// Try to infer the data type from sql value.
<a href=#434 id=434 data-nosnippet>434</a></span><span class="kw">fn </span>try_get_data_type_for_sql_value(value: <span class="kw-2">&amp;</span>SqlValue) -&gt; <span class="prelude-ty">Result</span>&lt;ArrowDataType&gt; {
<a href=#435 id=435 data-nosnippet>435</a> <span class="kw">match </span>value {
<a href=#436 id=436 data-nosnippet>436</a> SqlValue::Number(<span class="kw">_</span>, <span class="kw">_</span>) =&gt; <span class="prelude-val">Ok</span>(ArrowDataType::Float64),
<a href=#437 id=437 data-nosnippet>437</a> SqlValue::Null =&gt; <span class="prelude-val">Ok</span>(ArrowDataType::Null),
<a href=#438 id=438 data-nosnippet>438</a> SqlValue::Boolean(<span class="kw">_</span>) =&gt; <span class="prelude-val">Ok</span>(ArrowDataType::Boolean),
<a href=#439 id=439 data-nosnippet>439</a> SqlValue::HexStringLiteral(<span class="kw">_</span>)
<a href=#440 id=440 data-nosnippet>440</a> | SqlValue::DoubleQuotedString(<span class="kw">_</span>)
<a href=#441 id=441 data-nosnippet>441</a> | SqlValue::SingleQuotedString(<span class="kw">_</span>) =&gt; <span class="prelude-val">Ok</span>(ArrowDataType::Utf8),
<a href=#442 id=442 data-nosnippet>442</a> <span class="kw">_ </span>=&gt; error::BuildAdminFunctionArgsSnafu {
<a href=#443 id=443 data-nosnippet>443</a> msg: <span class="macro">format!</span>(<span class="string">"unsupported sql value: {value}"</span>),
<a href=#444 id=444 data-nosnippet>444</a> }
<a href=#445 id=445 data-nosnippet>445</a> .fail(),
<a href=#446 id=446 data-nosnippet>446</a> }
<a href=#447 id=447 data-nosnippet>447</a>}
<a href=#448 id=448 data-nosnippet>448</a>
<a href=#449 id=449 data-nosnippet>449</a><span class="attr">#[cfg(test)]
<a href=#450 id=450 data-nosnippet>450</a></span><span class="kw">mod </span>tests {
<a href=#451 id=451 data-nosnippet>451</a> <span class="kw">use </span>std::sync::Arc;
<a href=#452 id=452 data-nosnippet>452</a>
<a href=#453 id=453 data-nosnippet>453</a> <span class="kw">use </span>datatypes::vectors::{Int32Vector, VectorRef};
<a href=#454 id=454 data-nosnippet>454</a>
<a href=#455 id=455 data-nosnippet>455</a> <span class="kw">use </span><span class="kw">crate</span>::statement::admin::immediate_result;
<a href=#456 id=456 data-nosnippet>456</a>
<a href=#457 id=457 data-nosnippet>457</a> <span class="attr">#[test]
<a href=#458 id=458 data-nosnippet>458</a> </span><span class="kw">fn </span>empty_admin_function_result_has_no_immediate_value() {
<a href=#459 id=459 data-nosnippet>459</a> <span class="kw">let </span>result_vector: VectorRef = Arc::new(Int32Vector::from(<span class="macro">vec!</span>[]));
<a href=#460 id=460 data-nosnippet>460</a>
<a href=#461 id=461 data-nosnippet>461</a> <span class="macro">assert_eq!</span>(immediate_result(<span class="kw-2">&amp;</span>result_vector), <span class="prelude-val">None</span>);
<a href=#462 id=462 data-nosnippet>462</a> }
<a href=#463 id=463 data-nosnippet>463</a>}
</code></pre></div></section></main></body></html>
+570
View File
@@ -0,0 +1,570 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/admin/event.rs`."><title>event.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../../" data-static-root-path="../../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../../static.files/src-script-813739b1.js"></script><script defer src="../../../../src-files.js"></script><script defer src="../../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/admin/</div>event.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::any::Any;
<a href=#16 id=16 data-nosnippet>16</a>
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>api::v1::value::ValueData;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>api::v1::{ColumnSchema, Row};
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>common_event_recorder::Event;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>common_event_recorder::error::Result <span class="kw">as </span>EventResult;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>common_event_recorder::event_table::{
<a href=#22 id=22 data-nosnippet>22</a> ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN,
<a href=#23 id=23 data-nosnippet>23</a> ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value,
<a href=#24 id=24 data-nosnippet>24</a>};
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>datatypes::value::Value;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>serde_json::{Value <span class="kw">as </span>JsonValue, json};
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>sql::ast::{Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Value <span class="kw">as </span>SqlValue};
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>sql::statements::admin::Admin;
<a href=#29 id=29 data-nosnippet>29</a>
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span><span class="kw">crate</span>::error::Error;
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span><span class="kw">crate</span>::statement::admin::AdminFunctionRequest;
<a href=#32 id=32 data-nosnippet>32</a>
<a href=#33 id=33 data-nosnippet>33</a><span class="doccomment">/// Event type emitted for an ADMIN function execution.
<a href=#34 id=34 data-nosnippet>34</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">const </span>ADMIN_FUNCTION_EVENT_TYPE: <span class="kw-2">&amp;</span>str = <span class="string">"admin_function"</span>;
<a href=#35 id=35 data-nosnippet>35</a>
<a href=#36 id=36 data-nosnippet>36</a><span class="kw">const </span>PAYLOAD_VERSION: u8 = <span class="number">1</span>;
<a href=#37 id=37 data-nosnippet>37</a><span class="kw">const </span>VERSION_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"version"</span>;
<a href=#38 id=38 data-nosnippet>38</a><span class="kw">const </span>ARGUMENTS_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"arguments"</span>;
<a href=#39 id=39 data-nosnippet>39</a><span class="kw">const </span>ARGUMENT_TYPE_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"type"</span>;
<a href=#40 id=40 data-nosnippet>40</a><span class="kw">const </span>ARGUMENT_VALUE_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"value"</span>;
<a href=#41 id=41 data-nosnippet>41</a><span class="kw">const </span>LIST_ARGUMENT_TYPE: <span class="kw-2">&amp;</span>str = <span class="string">"list"</span>;
<a href=#42 id=42 data-nosnippet>42</a><span class="kw">const </span>SUBQUERY_ARGUMENT_TYPE: <span class="kw-2">&amp;</span>str = <span class="string">"subquery"</span>;
<a href=#43 id=43 data-nosnippet>43</a><span class="kw">const </span>RESULT_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"result"</span>;
<a href=#44 id=44 data-nosnippet>44</a><span class="kw">const </span>ERROR_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"error"</span>;
<a href=#45 id=45 data-nosnippet>45</a><span class="kw">const </span>SUCCEEDED_STATUS: <span class="kw-2">&amp;</span>str = <span class="string">"Succeeded"</span>;
<a href=#46 id=46 data-nosnippet>46</a><span class="kw">const </span>FAILED_STATUS: <span class="kw-2">&amp;</span>str = <span class="string">"Failed"</span>;
<a href=#47 id=47 data-nosnippet>47</a><span class="kw">const </span>UNSUPPORTED: <span class="kw-2">&amp;</span>str = <span class="string">"&lt;unsupported&gt;"</span>;
<a href=#48 id=48 data-nosnippet>48</a>
<a href=#49 id=49 data-nosnippet>49</a><span class="doccomment">/// The ADMIN function input captured before execution.
<a href=#50 id=50 data-nosnippet>50</a></span><span class="attr">#[derive(Debug, Clone)]
<a href=#51 id=51 data-nosnippet>51</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>AdminFunctionEventInput {
<a href=#52 id=52 data-nosnippet>52</a> actor: String,
<a href=#53 id=53 data-nosnippet>53</a> function: String,
<a href=#54 id=54 data-nosnippet>54</a> arguments: <span class="prelude-ty">Option</span>&lt;JsonValue&gt;,
<a href=#55 id=55 data-nosnippet>55</a>}
<a href=#56 id=56 data-nosnippet>56</a>
<a href=#57 id=57 data-nosnippet>57</a><span class="kw">impl </span>AdminFunctionEventInput {
<a href=#58 id=58 data-nosnippet>58</a> <span class="doccomment">/// Creates an event input from an ADMIN function request.
<a href=#59 id=59 data-nosnippet>59</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>from_request(request: <span class="kw-2">&amp;</span>AdminFunctionRequest) -&gt; <span class="self">Self </span>{
<a href=#60 id=60 data-nosnippet>60</a> <span class="kw">let </span>Admin::Func(function) = <span class="kw-2">&amp;</span>request.statement;
<a href=#61 id=61 data-nosnippet>61</a> <span class="kw">let </span>function_name = function.name.to_string().to_lowercase();
<a href=#62 id=62 data-nosnippet>62</a>
<a href=#63 id=63 data-nosnippet>63</a> <span class="kw">let </span>arguments = <span class="kw">match </span><span class="kw-2">&amp;</span>function.args {
<a href=#64 id=64 data-nosnippet>64</a> FunctionArguments::List(arguments) =&gt; <span class="prelude-val">Some</span>(<span class="macro">json!</span>({
<a href=#65 id=65 data-nosnippet>65</a> (ARGUMENT_TYPE_KEY): LIST_ARGUMENT_TYPE,
<a href=#66 id=66 data-nosnippet>66</a> (ARGUMENT_VALUE_KEY): arguments.args.iter().map(argument_to_json).collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;(),
<a href=#67 id=67 data-nosnippet>67</a> })),
<a href=#68 id=68 data-nosnippet>68</a> FunctionArguments::Subquery(query) =&gt; <span class="prelude-val">Some</span>(<span class="macro">json!</span>({
<a href=#69 id=69 data-nosnippet>69</a> (ARGUMENT_TYPE_KEY): SUBQUERY_ARGUMENT_TYPE,
<a href=#70 id=70 data-nosnippet>70</a> (ARGUMENT_VALUE_KEY): query.to_string(),
<a href=#71 id=71 data-nosnippet>71</a> })),
<a href=#72 id=72 data-nosnippet>72</a> FunctionArguments::None =&gt; <span class="prelude-val">None</span>,
<a href=#73 id=73 data-nosnippet>73</a> };
<a href=#74 id=74 data-nosnippet>74</a>
<a href=#75 id=75 data-nosnippet>75</a> <span class="self">Self </span>{
<a href=#76 id=76 data-nosnippet>76</a> actor: request.query_ctx.current_user().username().to_string(),
<a href=#77 id=77 data-nosnippet>77</a> function: function_name,
<a href=#78 id=78 data-nosnippet>78</a> arguments,
<a href=#79 id=79 data-nosnippet>79</a> }
<a href=#80 id=80 data-nosnippet>80</a> }
<a href=#81 id=81 data-nosnippet>81</a>}
<a href=#82 id=82 data-nosnippet>82</a>
<a href=#83 id=83 data-nosnippet>83</a><span class="doccomment">/// An event describing the outcome of an ADMIN function execution.
<a href=#84 id=84 data-nosnippet>84</a></span><span class="attr">#[derive(Debug)]
<a href=#85 id=85 data-nosnippet>85</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>AdminFunctionEvent {
<a href=#86 id=86 data-nosnippet>86</a> actor: String,
<a href=#87 id=87 data-nosnippet>87</a> function_name: String,
<a href=#88 id=88 data-nosnippet>88</a> status: <span class="kw-2">&amp;</span><span class="lifetime">'static </span>str,
<a href=#89 id=89 data-nosnippet>89</a> output: JsonValue,
<a href=#90 id=90 data-nosnippet>90</a> payload: JsonValue,
<a href=#91 id=91 data-nosnippet>91</a>}
<a href=#92 id=92 data-nosnippet>92</a>
<a href=#93 id=93 data-nosnippet>93</a><span class="kw">impl </span>AdminFunctionEvent {
<a href=#94 id=94 data-nosnippet>94</a> <span class="doccomment">/// Creates a successful ADMIN function event.
<a href=#95 id=95 data-nosnippet>95</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>success(input: AdminFunctionEventInput, result: <span class="prelude-ty">Option</span>&lt;<span class="kw-2">&amp;</span>Value&gt;) -&gt; <span class="self">Self </span>{
<a href=#96 id=96 data-nosnippet>96</a> <span class="kw">let </span>AdminFunctionEventInput {
<a href=#97 id=97 data-nosnippet>97</a> actor,
<a href=#98 id=98 data-nosnippet>98</a> function,
<a href=#99 id=99 data-nosnippet>99</a> arguments,
<a href=#100 id=100 data-nosnippet>100</a> } = input;
<a href=#101 id=101 data-nosnippet>101</a> <span class="self">Self </span>{
<a href=#102 id=102 data-nosnippet>102</a> actor,
<a href=#103 id=103 data-nosnippet>103</a> function_name: function,
<a href=#104 id=104 data-nosnippet>104</a> status: SUCCEEDED_STATUS,
<a href=#105 id=105 data-nosnippet>105</a> output: result.map_or_else(
<a href=#106 id=106 data-nosnippet>106</a> || <span class="macro">json!</span>({}),
<a href=#107 id=107 data-nosnippet>107</a> |result| {
<a href=#108 id=108 data-nosnippet>108</a> <span class="macro">json!</span>({
<a href=#109 id=109 data-nosnippet>109</a> (RESULT_KEY): value_to_json(result),
<a href=#110 id=110 data-nosnippet>110</a> })
<a href=#111 id=111 data-nosnippet>111</a> },
<a href=#112 id=112 data-nosnippet>112</a> ),
<a href=#113 id=113 data-nosnippet>113</a> payload: input_payload(arguments),
<a href=#114 id=114 data-nosnippet>114</a> }
<a href=#115 id=115 data-nosnippet>115</a> }
<a href=#116 id=116 data-nosnippet>116</a>
<a href=#117 id=117 data-nosnippet>117</a> <span class="doccomment">/// Creates a failed ADMIN function event with the debug representation of the error.
<a href=#118 id=118 data-nosnippet>118</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>failure(input: AdminFunctionEventInput, error: <span class="kw-2">&amp;</span>Error) -&gt; <span class="self">Self </span>{
<a href=#119 id=119 data-nosnippet>119</a> <span class="kw">let </span>AdminFunctionEventInput {
<a href=#120 id=120 data-nosnippet>120</a> actor,
<a href=#121 id=121 data-nosnippet>121</a> function,
<a href=#122 id=122 data-nosnippet>122</a> arguments,
<a href=#123 id=123 data-nosnippet>123</a> } = input;
<a href=#124 id=124 data-nosnippet>124</a> <span class="self">Self </span>{
<a href=#125 id=125 data-nosnippet>125</a> actor,
<a href=#126 id=126 data-nosnippet>126</a> function_name: function,
<a href=#127 id=127 data-nosnippet>127</a> status: FAILED_STATUS,
<a href=#128 id=128 data-nosnippet>128</a> output: <span class="macro">json!</span>({
<a href=#129 id=129 data-nosnippet>129</a> (ERROR_KEY): <span class="macro">format!</span>(<span class="string">"{error:?}"</span>),
<a href=#130 id=130 data-nosnippet>130</a> }),
<a href=#131 id=131 data-nosnippet>131</a> payload: input_payload(arguments),
<a href=#132 id=132 data-nosnippet>132</a> }
<a href=#133 id=133 data-nosnippet>133</a> }
<a href=#134 id=134 data-nosnippet>134</a>
<a href=#135 id=135 data-nosnippet>135</a> <span class="doccomment">/// Creates a failed ADMIN function event for a cancelled execution.
<a href=#136 id=136 data-nosnippet>136</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>cancelled(input: AdminFunctionEventInput) -&gt; <span class="self">Self </span>{
<a href=#137 id=137 data-nosnippet>137</a> <span class="self">Self</span>::failure(input, <span class="kw-2">&amp;</span>Error::AdminFunctionCancelled)
<a href=#138 id=138 data-nosnippet>138</a> }
<a href=#139 id=139 data-nosnippet>139</a>}
<a href=#140 id=140 data-nosnippet>140</a>
<a href=#141 id=141 data-nosnippet>141</a><span class="kw">impl </span>Event <span class="kw">for </span>AdminFunctionEvent {
<a href=#142 id=142 data-nosnippet>142</a> <span class="kw">fn </span>event_type(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; <span class="kw-2">&amp;</span>str {
<a href=#143 id=143 data-nosnippet>143</a> ADMIN_FUNCTION_EVENT_TYPE
<a href=#144 id=144 data-nosnippet>144</a> }
<a href=#145 id=145 data-nosnippet>145</a>
<a href=#146 id=146 data-nosnippet>146</a> <span class="kw">fn </span>json_payload(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; EventResult&lt;JsonValue&gt; {
<a href=#147 id=147 data-nosnippet>147</a> <span class="prelude-val">Ok</span>(<span class="self">self</span>.payload.clone())
<a href=#148 id=148 data-nosnippet>148</a> }
<a href=#149 id=149 data-nosnippet>149</a>
<a href=#150 id=150 data-nosnippet>150</a> <span class="kw">fn </span>extra_schema(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; Vec&lt;ColumnSchema&gt; {
<a href=#151 id=151 data-nosnippet>151</a> column_schemas([
<a href=#152 id=152 data-nosnippet>152</a> <span class="kw-2">&amp;</span>ACTOR_COLUMN,
<a href=#153 id=153 data-nosnippet>153</a> <span class="kw-2">&amp;</span>ADMIN_FUNCTION_NAME_COLUMN,
<a href=#154 id=154 data-nosnippet>154</a> <span class="kw-2">&amp;</span>ADMIN_FUNCTION_STATUS_COLUMN,
<a href=#155 id=155 data-nosnippet>155</a> <span class="kw-2">&amp;</span>ADMIN_FUNCTION_OUTPUT_COLUMN,
<a href=#156 id=156 data-nosnippet>156</a> ])
<a href=#157 id=157 data-nosnippet>157</a> }
<a href=#158 id=158 data-nosnippet>158</a>
<a href=#159 id=159 data-nosnippet>159</a> <span class="kw">fn </span>extra_rows(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; EventResult&lt;Vec&lt;Row&gt;&gt; {
<a href=#160 id=160 data-nosnippet>160</a> <span class="prelude-val">Ok</span>(<span class="macro">vec!</span>[Row {
<a href=#161 id=161 data-nosnippet>161</a> values: <span class="macro">vec!</span>[
<a href=#162 id=162 data-nosnippet>162</a> ValueData::StringValue(<span class="self">self</span>.actor.clone()).into(),
<a href=#163 id=163 data-nosnippet>163</a> ValueData::StringValue(<span class="self">self</span>.function_name.clone()).into(),
<a href=#164 id=164 data-nosnippet>164</a> ValueData::StringValue(<span class="self">self</span>.status.to_string()).into(),
<a href=#165 id=165 data-nosnippet>165</a> jsonb_value(<span class="kw-2">&amp;</span><span class="self">self</span>.output),
<a href=#166 id=166 data-nosnippet>166</a> ],
<a href=#167 id=167 data-nosnippet>167</a> }])
<a href=#168 id=168 data-nosnippet>168</a> }
<a href=#169 id=169 data-nosnippet>169</a>
<a href=#170 id=170 data-nosnippet>170</a> <span class="kw">fn </span>as_any(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; <span class="kw-2">&amp;</span><span class="kw">dyn </span>Any {
<a href=#171 id=171 data-nosnippet>171</a> <span class="self">self
<a href=#172 id=172 data-nosnippet>172</a> </span>}
<a href=#173 id=173 data-nosnippet>173</a>}
<a href=#174 id=174 data-nosnippet>174</a>
<a href=#175 id=175 data-nosnippet>175</a><span class="kw">fn </span>input_payload(arguments: <span class="prelude-ty">Option</span>&lt;JsonValue&gt;) -&gt; JsonValue {
<a href=#176 id=176 data-nosnippet>176</a> <span class="kw">let </span><span class="kw-2">mut </span>payload = <span class="macro">json!</span>({ VERSION_KEY: PAYLOAD_VERSION });
<a href=#177 id=177 data-nosnippet>177</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(arguments) = arguments {
<a href=#178 id=178 data-nosnippet>178</a> payload[ARGUMENTS_KEY] = arguments;
<a href=#179 id=179 data-nosnippet>179</a> }
<a href=#180 id=180 data-nosnippet>180</a> payload
<a href=#181 id=181 data-nosnippet>181</a>}
<a href=#182 id=182 data-nosnippet>182</a>
<a href=#183 id=183 data-nosnippet>183</a><span class="kw">fn </span>argument_to_json(argument: <span class="kw-2">&amp;</span>FunctionArg) -&gt; JsonValue {
<a href=#184 id=184 data-nosnippet>184</a> <span class="kw">match </span>argument {
<a href=#185 id=185 data-nosnippet>185</a> FunctionArg::Unnamed(FunctionArgExpr::Expr(Expr::Value(value))) =&gt; {
<a href=#186 id=186 data-nosnippet>186</a> sql_value_to_json(<span class="kw-2">&amp;</span>value.value)
<a href=#187 id=187 data-nosnippet>187</a> }
<a href=#188 id=188 data-nosnippet>188</a> <span class="kw">_ </span>=&gt; JsonValue::String(argument.to_string()),
<a href=#189 id=189 data-nosnippet>189</a> }
<a href=#190 id=190 data-nosnippet>190</a>}
<a href=#191 id=191 data-nosnippet>191</a>
<a href=#192 id=192 data-nosnippet>192</a><span class="kw">fn </span>sql_value_to_json(value: <span class="kw-2">&amp;</span>SqlValue) -&gt; JsonValue {
<a href=#193 id=193 data-nosnippet>193</a> <span class="kw">match </span>value {
<a href=#194 id=194 data-nosnippet>194</a> SqlValue::Number(value, <span class="kw">_</span>) =&gt; {
<a href=#195 id=195 data-nosnippet>195</a> serde_json::from_str(value).unwrap_or_else(|<span class="kw">_</span>| JsonValue::String(value.clone()))
<a href=#196 id=196 data-nosnippet>196</a> }
<a href=#197 id=197 data-nosnippet>197</a> SqlValue::Boolean(value) =&gt; JsonValue::Bool(<span class="kw-2">*</span>value),
<a href=#198 id=198 data-nosnippet>198</a> SqlValue::Null =&gt; JsonValue::Null,
<a href=#199 id=199 data-nosnippet>199</a> SqlValue::SingleQuotedString(value) | SqlValue::DoubleQuotedString(value) =&gt; {
<a href=#200 id=200 data-nosnippet>200</a> JsonValue::String(value.clone())
<a href=#201 id=201 data-nosnippet>201</a> }
<a href=#202 id=202 data-nosnippet>202</a> SqlValue::Placeholder(value) =&gt; JsonValue::String(value.clone()),
<a href=#203 id=203 data-nosnippet>203</a> <span class="kw">_ </span>=&gt; value
<a href=#204 id=204 data-nosnippet>204</a> .clone()
<a href=#205 id=205 data-nosnippet>205</a> .into_string()
<a href=#206 id=206 data-nosnippet>206</a> .map(JsonValue::String)
<a href=#207 id=207 data-nosnippet>207</a> .unwrap_or_else(|| JsonValue::String(value.to_string())),
<a href=#208 id=208 data-nosnippet>208</a> }
<a href=#209 id=209 data-nosnippet>209</a>}
<a href=#210 id=210 data-nosnippet>210</a>
<a href=#211 id=211 data-nosnippet>211</a><span class="kw">fn </span>value_to_json(value: <span class="kw-2">&amp;</span>Value) -&gt; JsonValue {
<a href=#212 id=212 data-nosnippet>212</a> <span class="kw">match </span>value {
<a href=#213 id=213 data-nosnippet>213</a> Value::Float32(value) <span class="kw">if </span>!value.<span class="number">0</span>.is_finite() =&gt; JsonValue::String(value.to_string()),
<a href=#214 id=214 data-nosnippet>214</a> Value::Float64(value) <span class="kw">if </span>!value.<span class="number">0</span>.is_finite() =&gt; JsonValue::String(value.to_string()),
<a href=#215 id=215 data-nosnippet>215</a> <span class="kw">_ </span>=&gt; JsonValue::try_from(value.clone())
<a href=#216 id=216 data-nosnippet>216</a> .unwrap_or_else(|<span class="kw">_</span>| JsonValue::String(UNSUPPORTED.to_string())),
<a href=#217 id=217 data-nosnippet>217</a> }
<a href=#218 id=218 data-nosnippet>218</a>}
<a href=#219 id=219 data-nosnippet>219</a>
<a href=#220 id=220 data-nosnippet>220</a><span class="attr">#[cfg(test)]
<a href=#221 id=221 data-nosnippet>221</a></span><span class="kw">mod </span>tests {
<a href=#222 id=222 data-nosnippet>222</a> <span class="kw">use </span>api::v1::Row;
<a href=#223 id=223 data-nosnippet>223</a> <span class="kw">use </span>api::v1::value::ValueData;
<a href=#224 id=224 data-nosnippet>224</a> <span class="kw">use </span>common_event_recorder::Event;
<a href=#225 id=225 data-nosnippet>225</a> <span class="kw">use </span>common_event_recorder::event_table::{
<a href=#226 id=226 data-nosnippet>226</a> ACTOR_COLUMN, ADMIN_FUNCTION_NAME_COLUMN, ADMIN_FUNCTION_OUTPUT_COLUMN,
<a href=#227 id=227 data-nosnippet>227</a> ADMIN_FUNCTION_STATUS_COLUMN, column_schemas, jsonb_value,
<a href=#228 id=228 data-nosnippet>228</a> };
<a href=#229 id=229 data-nosnippet>229</a> <span class="kw">use </span>datatypes::value::Value;
<a href=#230 id=230 data-nosnippet>230</a> <span class="kw">use </span>serde_json::json;
<a href=#231 id=231 data-nosnippet>231</a> <span class="kw">use </span>session::context::QueryContext;
<a href=#232 id=232 data-nosnippet>232</a> <span class="kw">use </span>sql::ast::{
<a href=#233 id=233 data-nosnippet>233</a> Expr, FunctionArg, FunctionArgExpr, FunctionArguments, Ident, Value <span class="kw">as </span>SqlValue,
<a href=#234 id=234 data-nosnippet>234</a> };
<a href=#235 id=235 data-nosnippet>235</a> <span class="kw">use </span>sql::dialect::GreptimeDbDialect;
<a href=#236 id=236 data-nosnippet>236</a> <span class="kw">use </span>sql::parser::{ParseOptions, ParserContext};
<a href=#237 id=237 data-nosnippet>237</a> <span class="kw">use </span>sql::statements::admin::Admin;
<a href=#238 id=238 data-nosnippet>238</a> <span class="kw">use </span>sql::statements::statement::Statement;
<a href=#239 id=239 data-nosnippet>239</a> <span class="kw">use </span>sqlparser::ast::{DollarQuotedString, FunctionArgOperator, FunctionArgumentList};
<a href=#240 id=240 data-nosnippet>240</a>
<a href=#241 id=241 data-nosnippet>241</a> <span class="kw">use </span><span class="kw">crate</span>::error::Error;
<a href=#242 id=242 data-nosnippet>242</a> <span class="kw">use </span><span class="kw">crate</span>::statement::admin::AdminFunctionRequest;
<a href=#243 id=243 data-nosnippet>243</a> <span class="kw">use </span><span class="kw">crate</span>::statement::admin::event::{
<a href=#244 id=244 data-nosnippet>244</a> ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput, sql_value_to_json,
<a href=#245 id=245 data-nosnippet>245</a> value_to_json,
<a href=#246 id=246 data-nosnippet>246</a> };
<a href=#247 id=247 data-nosnippet>247</a>
<a href=#248 id=248 data-nosnippet>248</a> <span class="kw">fn </span>request(sql: <span class="kw-2">&amp;</span>str) -&gt; AdminFunctionRequest {
<a href=#249 id=249 data-nosnippet>249</a> <span class="kw">let </span>Statement::Admin(statement) =
<a href=#250 id=250 data-nosnippet>250</a> ParserContext::create_with_dialect(sql, <span class="kw-2">&amp;</span>GreptimeDbDialect {}, ParseOptions::default())
<a href=#251 id=251 data-nosnippet>251</a> .unwrap()
<a href=#252 id=252 data-nosnippet>252</a> .remove(<span class="number">0</span>)
<a href=#253 id=253 data-nosnippet>253</a> <span class="kw">else </span>{
<a href=#254 id=254 data-nosnippet>254</a> <span class="macro">panic!</span>(<span class="string">"expected ADMIN statement"</span>)
<a href=#255 id=255 data-nosnippet>255</a> };
<a href=#256 id=256 data-nosnippet>256</a> AdminFunctionRequest {
<a href=#257 id=257 data-nosnippet>257</a> statement,
<a href=#258 id=258 data-nosnippet>258</a> query_ctx: QueryContext::arc(),
<a href=#259 id=259 data-nosnippet>259</a> }
<a href=#260 id=260 data-nosnippet>260</a> }
<a href=#261 id=261 data-nosnippet>261</a>
<a href=#262 id=262 data-nosnippet>262</a> <span class="kw">fn </span>request_with_arguments(arguments: FunctionArguments) -&gt; AdminFunctionRequest {
<a href=#263 id=263 data-nosnippet>263</a> <span class="kw">let </span><span class="kw-2">mut </span>request = request(<span class="string">"ADMIN plugin_function()"</span>);
<a href=#264 id=264 data-nosnippet>264</a> <span class="kw">let </span>Admin::Func(function) = <span class="kw-2">&amp;mut </span>request.statement;
<a href=#265 id=265 data-nosnippet>265</a> function.args = arguments;
<a href=#266 id=266 data-nosnippet>266</a> request
<a href=#267 id=267 data-nosnippet>267</a> }
<a href=#268 id=268 data-nosnippet>268</a>
<a href=#269 id=269 data-nosnippet>269</a> <span class="kw">fn </span>assert_admin_columns(
<a href=#270 id=270 data-nosnippet>270</a> event: <span class="kw-2">&amp;</span>AdminFunctionEvent,
<a href=#271 id=271 data-nosnippet>271</a> function_name: <span class="kw-2">&amp;</span>str,
<a href=#272 id=272 data-nosnippet>272</a> status: <span class="kw-2">&amp;</span>str,
<a href=#273 id=273 data-nosnippet>273</a> output: serde_json::Value,
<a href=#274 id=274 data-nosnippet>274</a> ) {
<a href=#275 id=275 data-nosnippet>275</a> <span class="macro">assert_eq!</span>(event.event_type(), ADMIN_FUNCTION_EVENT_TYPE);
<a href=#276 id=276 data-nosnippet>276</a> <span class="macro">assert_eq!</span>(
<a href=#277 id=277 data-nosnippet>277</a> event.extra_schema(),
<a href=#278 id=278 data-nosnippet>278</a> column_schemas([
<a href=#279 id=279 data-nosnippet>279</a> <span class="kw-2">&amp;</span>ACTOR_COLUMN,
<a href=#280 id=280 data-nosnippet>280</a> <span class="kw-2">&amp;</span>ADMIN_FUNCTION_NAME_COLUMN,
<a href=#281 id=281 data-nosnippet>281</a> <span class="kw-2">&amp;</span>ADMIN_FUNCTION_STATUS_COLUMN,
<a href=#282 id=282 data-nosnippet>282</a> <span class="kw-2">&amp;</span>ADMIN_FUNCTION_OUTPUT_COLUMN,
<a href=#283 id=283 data-nosnippet>283</a> ])
<a href=#284 id=284 data-nosnippet>284</a> );
<a href=#285 id=285 data-nosnippet>285</a> <span class="macro">assert_eq!</span>(
<a href=#286 id=286 data-nosnippet>286</a> event.extra_rows().unwrap(),
<a href=#287 id=287 data-nosnippet>287</a> <span class="macro">vec!</span>[Row {
<a href=#288 id=288 data-nosnippet>288</a> values: <span class="macro">vec!</span>[
<a href=#289 id=289 data-nosnippet>289</a> ValueData::StringValue(<span class="string">"greptime"</span>.to_string()).into(),
<a href=#290 id=290 data-nosnippet>290</a> ValueData::StringValue(function_name.to_string()).into(),
<a href=#291 id=291 data-nosnippet>291</a> ValueData::StringValue(status.to_string()).into(),
<a href=#292 id=292 data-nosnippet>292</a> jsonb_value(<span class="kw-2">&amp;</span>output),
<a href=#293 id=293 data-nosnippet>293</a> ],
<a href=#294 id=294 data-nosnippet>294</a> }]
<a href=#295 id=295 data-nosnippet>295</a> );
<a href=#296 id=296 data-nosnippet>296</a> }
<a href=#297 id=297 data-nosnippet>297</a>
<a href=#298 id=298 data-nosnippet>298</a> <span class="attr">#[test]
<a href=#299 id=299 data-nosnippet>299</a> </span><span class="kw">fn </span>success_uses_separate_output_columns() {
<a href=#300 id=300 data-nosnippet>300</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request(
<a href=#301 id=301 data-nosnippet>301</a> <span class="string">"ADMIN flush_table('greptime.public.demo')"</span>,
<a href=#302 id=302 data-nosnippet>302</a> ));
<a href=#303 id=303 data-nosnippet>303</a> <span class="kw">let </span>event = AdminFunctionEvent::success(input, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>Value::UInt64(<span class="number">3</span>)));
<a href=#304 id=304 data-nosnippet>304</a>
<a href=#305 id=305 data-nosnippet>305</a> assert_admin_columns(<span class="kw-2">&amp;</span>event, <span class="string">"flush_table"</span>, <span class="string">"Succeeded"</span>, <span class="macro">json!</span>({<span class="string">"result"</span>: <span class="number">3</span>}));
<a href=#306 id=306 data-nosnippet>306</a> <span class="macro">assert_eq!</span>(
<a href=#307 id=307 data-nosnippet>307</a> event.json_payload().unwrap(),
<a href=#308 id=308 data-nosnippet>308</a> <span class="macro">json!</span>({
<a href=#309 id=309 data-nosnippet>309</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#310 id=310 data-nosnippet>310</a> <span class="string">"arguments"</span>: {<span class="string">"type"</span>: <span class="string">"list"</span>, <span class="string">"value"</span>: [<span class="string">"greptime.public.demo"</span>]},
<a href=#311 id=311 data-nosnippet>311</a> })
<a href=#312 id=312 data-nosnippet>312</a> );
<a href=#313 id=313 data-nosnippet>313</a> }
<a href=#314 id=314 data-nosnippet>314</a>
<a href=#315 id=315 data-nosnippet>315</a> <span class="attr">#[test]
<a href=#316 id=316 data-nosnippet>316</a> </span><span class="kw">fn </span>successful_null_is_explicit() {
<a href=#317 id=317 data-nosnippet>317</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request(
<a href=#318 id=318 data-nosnippet>318</a> <span class="string">"ADMIN migrate_region(NULL, NULL, NULL)"</span>,
<a href=#319 id=319 data-nosnippet>319</a> ));
<a href=#320 id=320 data-nosnippet>320</a> <span class="kw">let </span>event = AdminFunctionEvent::success(input, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>Value::Null));
<a href=#321 id=321 data-nosnippet>321</a>
<a href=#322 id=322 data-nosnippet>322</a> assert_admin_columns(
<a href=#323 id=323 data-nosnippet>323</a> <span class="kw-2">&amp;</span>event,
<a href=#324 id=324 data-nosnippet>324</a> <span class="string">"migrate_region"</span>,
<a href=#325 id=325 data-nosnippet>325</a> <span class="string">"Succeeded"</span>,
<a href=#326 id=326 data-nosnippet>326</a> <span class="macro">json!</span>({<span class="string">"result"</span>: null}),
<a href=#327 id=327 data-nosnippet>327</a> );
<a href=#328 id=328 data-nosnippet>328</a> <span class="macro">assert_eq!</span>(
<a href=#329 id=329 data-nosnippet>329</a> event.json_payload().unwrap(),
<a href=#330 id=330 data-nosnippet>330</a> <span class="macro">json!</span>({
<a href=#331 id=331 data-nosnippet>331</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#332 id=332 data-nosnippet>332</a> <span class="string">"arguments"</span>: {<span class="string">"type"</span>: <span class="string">"list"</span>, <span class="string">"value"</span>: [null, null, null]},
<a href=#333 id=333 data-nosnippet>333</a> })
<a href=#334 id=334 data-nosnippet>334</a> );
<a href=#335 id=335 data-nosnippet>335</a> }
<a href=#336 id=336 data-nosnippet>336</a>
<a href=#337 id=337 data-nosnippet>337</a> <span class="attr">#[test]
<a href=#338 id=338 data-nosnippet>338</a> </span><span class="kw">fn </span>unknown_function_values_are_recorded() {
<a href=#339 id=339 data-nosnippet>339</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request(
<a href=#340 id=340 data-nosnippet>340</a> <span class="string">"ADMIN plugin_function('plugin-value', NULL)"</span>,
<a href=#341 id=341 data-nosnippet>341</a> ));
<a href=#342 id=342 data-nosnippet>342</a> <span class="kw">let </span>event = AdminFunctionEvent::success(input, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>Value::from(<span class="string">"plugin-result"</span>)));
<a href=#343 id=343 data-nosnippet>343</a>
<a href=#344 id=344 data-nosnippet>344</a> assert_admin_columns(
<a href=#345 id=345 data-nosnippet>345</a> <span class="kw-2">&amp;</span>event,
<a href=#346 id=346 data-nosnippet>346</a> <span class="string">"plugin_function"</span>,
<a href=#347 id=347 data-nosnippet>347</a> <span class="string">"Succeeded"</span>,
<a href=#348 id=348 data-nosnippet>348</a> <span class="macro">json!</span>({<span class="string">"result"</span>: <span class="string">"plugin-result"</span>}),
<a href=#349 id=349 data-nosnippet>349</a> );
<a href=#350 id=350 data-nosnippet>350</a> <span class="macro">assert_eq!</span>(
<a href=#351 id=351 data-nosnippet>351</a> event.json_payload().unwrap(),
<a href=#352 id=352 data-nosnippet>352</a> <span class="macro">json!</span>({
<a href=#353 id=353 data-nosnippet>353</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#354 id=354 data-nosnippet>354</a> <span class="string">"arguments"</span>: {<span class="string">"type"</span>: <span class="string">"list"</span>, <span class="string">"value"</span>: [<span class="string">"plugin-value"</span>, null]},
<a href=#355 id=355 data-nosnippet>355</a> })
<a href=#356 id=356 data-nosnippet>356</a> );
<a href=#357 id=357 data-nosnippet>357</a> }
<a href=#358 id=358 data-nosnippet>358</a>
<a href=#359 id=359 data-nosnippet>359</a> <span class="attr">#[test]
<a href=#360 id=360 data-nosnippet>360</a> </span><span class="kw">fn </span>failure_payload_has_error_only() {
<a href=#361 id=361 data-nosnippet>361</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request(
<a href=#362 id=362 data-nosnippet>362</a> <span class="string">"ADMIN flush_table('greptime.public.demo')"</span>,
<a href=#363 id=363 data-nosnippet>363</a> ));
<a href=#364 id=364 data-nosnippet>364</a> <span class="kw">let </span>error = Error::BuildAdminFunctionArgs {
<a href=#365 id=365 data-nosnippet>365</a> msg: <span class="string">"invalid input"</span>.to_string(),
<a href=#366 id=366 data-nosnippet>366</a> };
<a href=#367 id=367 data-nosnippet>367</a> <span class="kw">let </span>expected_error = <span class="macro">format!</span>(<span class="string">"{error:?}"</span>);
<a href=#368 id=368 data-nosnippet>368</a> <span class="kw">let </span>event = AdminFunctionEvent::failure(input, <span class="kw-2">&amp;</span>error);
<a href=#369 id=369 data-nosnippet>369</a>
<a href=#370 id=370 data-nosnippet>370</a> assert_admin_columns(
<a href=#371 id=371 data-nosnippet>371</a> <span class="kw-2">&amp;</span>event,
<a href=#372 id=372 data-nosnippet>372</a> <span class="string">"flush_table"</span>,
<a href=#373 id=373 data-nosnippet>373</a> <span class="string">"Failed"</span>,
<a href=#374 id=374 data-nosnippet>374</a> <span class="macro">json!</span>({<span class="string">"error"</span>: expected_error}),
<a href=#375 id=375 data-nosnippet>375</a> );
<a href=#376 id=376 data-nosnippet>376</a> <span class="macro">assert_eq!</span>(
<a href=#377 id=377 data-nosnippet>377</a> event.json_payload().unwrap(),
<a href=#378 id=378 data-nosnippet>378</a> <span class="macro">json!</span>({
<a href=#379 id=379 data-nosnippet>379</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#380 id=380 data-nosnippet>380</a> <span class="string">"arguments"</span>: {<span class="string">"type"</span>: <span class="string">"list"</span>, <span class="string">"value"</span>: [<span class="string">"greptime.public.demo"</span>]},
<a href=#381 id=381 data-nosnippet>381</a> })
<a href=#382 id=382 data-nosnippet>382</a> );
<a href=#383 id=383 data-nosnippet>383</a> }
<a href=#384 id=384 data-nosnippet>384</a>
<a href=#385 id=385 data-nosnippet>385</a> <span class="attr">#[test]
<a href=#386 id=386 data-nosnippet>386</a> </span><span class="kw">fn </span>cancellation_uses_debug_error_format() {
<a href=#387 id=387 data-nosnippet>387</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request(<span class="string">"ADMIN flush_table('demo')"</span>));
<a href=#388 id=388 data-nosnippet>388</a> <span class="kw">let </span>event = AdminFunctionEvent::cancelled(input);
<a href=#389 id=389 data-nosnippet>389</a>
<a href=#390 id=390 data-nosnippet>390</a> assert_admin_columns(
<a href=#391 id=391 data-nosnippet>391</a> <span class="kw-2">&amp;</span>event,
<a href=#392 id=392 data-nosnippet>392</a> <span class="string">"flush_table"</span>,
<a href=#393 id=393 data-nosnippet>393</a> <span class="string">"Failed"</span>,
<a href=#394 id=394 data-nosnippet>394</a> <span class="macro">json!</span>({<span class="string">"error"</span>: <span class="macro">format!</span>(<span class="string">"{:?}"</span>, Error::AdminFunctionCancelled)}),
<a href=#395 id=395 data-nosnippet>395</a> );
<a href=#396 id=396 data-nosnippet>396</a> }
<a href=#397 id=397 data-nosnippet>397</a>
<a href=#398 id=398 data-nosnippet>398</a> <span class="attr">#[test]
<a href=#399 id=399 data-nosnippet>399</a> </span><span class="kw">fn </span>subquery_arguments_use_display() {
<a href=#400 id=400 data-nosnippet>400</a> <span class="kw">let </span>Statement::Query(query) = ParserContext::create_with_dialect(
<a href=#401 id=401 data-nosnippet>401</a> <span class="string">"SELECT 1"</span>,
<a href=#402 id=402 data-nosnippet>402</a> <span class="kw-2">&amp;</span>GreptimeDbDialect {},
<a href=#403 id=403 data-nosnippet>403</a> ParseOptions::default(),
<a href=#404 id=404 data-nosnippet>404</a> )
<a href=#405 id=405 data-nosnippet>405</a> .unwrap()
<a href=#406 id=406 data-nosnippet>406</a> .remove(<span class="number">0</span>) <span class="kw">else </span>{
<a href=#407 id=407 data-nosnippet>407</a> <span class="macro">panic!</span>(<span class="string">"expected query statement"</span>)
<a href=#408 id=408 data-nosnippet>408</a> };
<a href=#409 id=409 data-nosnippet>409</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request_with_arguments(
<a href=#410 id=410 data-nosnippet>410</a> FunctionArguments::Subquery(Box::new(query.inner)),
<a href=#411 id=411 data-nosnippet>411</a> ));
<a href=#412 id=412 data-nosnippet>412</a> <span class="kw">let </span>event = AdminFunctionEvent::failure(
<a href=#413 id=413 data-nosnippet>413</a> input,
<a href=#414 id=414 data-nosnippet>414</a> <span class="kw-2">&amp;</span>Error::BuildAdminFunctionArgs {
<a href=#415 id=415 data-nosnippet>415</a> msg: <span class="string">"subquery is not executable"</span>.to_string(),
<a href=#416 id=416 data-nosnippet>416</a> },
<a href=#417 id=417 data-nosnippet>417</a> );
<a href=#418 id=418 data-nosnippet>418</a>
<a href=#419 id=419 data-nosnippet>419</a> <span class="macro">assert_eq!</span>(
<a href=#420 id=420 data-nosnippet>420</a> event.json_payload().unwrap(),
<a href=#421 id=421 data-nosnippet>421</a> <span class="macro">json!</span>({
<a href=#422 id=422 data-nosnippet>422</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#423 id=423 data-nosnippet>423</a> <span class="string">"arguments"</span>: {<span class="string">"type"</span>: <span class="string">"subquery"</span>, <span class="string">"value"</span>: <span class="string">"SELECT 1"</span>},
<a href=#424 id=424 data-nosnippet>424</a> })
<a href=#425 id=425 data-nosnippet>425</a> );
<a href=#426 id=426 data-nosnippet>426</a> }
<a href=#427 id=427 data-nosnippet>427</a>
<a href=#428 id=428 data-nosnippet>428</a> <span class="attr">#[test]
<a href=#429 id=429 data-nosnippet>429</a> </span><span class="kw">fn </span>no_arguments_omits_arguments_field_and_empty_result() {
<a href=#430 id=430 data-nosnippet>430</a> <span class="kw">let </span>input =
<a href=#431 id=431 data-nosnippet>431</a> AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request_with_arguments(FunctionArguments::None));
<a href=#432 id=432 data-nosnippet>432</a> <span class="kw">let </span>event = AdminFunctionEvent::success(input, <span class="prelude-val">None</span>);
<a href=#433 id=433 data-nosnippet>433</a>
<a href=#434 id=434 data-nosnippet>434</a> assert_admin_columns(<span class="kw-2">&amp;</span>event, <span class="string">"plugin_function"</span>, <span class="string">"Succeeded"</span>, <span class="macro">json!</span>({}));
<a href=#435 id=435 data-nosnippet>435</a> <span class="macro">assert_eq!</span>(event.json_payload().unwrap(), <span class="macro">json!</span>({<span class="string">"version"</span>: <span class="number">1</span>}));
<a href=#436 id=436 data-nosnippet>436</a> }
<a href=#437 id=437 data-nosnippet>437</a>
<a href=#438 id=438 data-nosnippet>438</a> <span class="attr">#[test]
<a href=#439 id=439 data-nosnippet>439</a> </span><span class="kw">fn </span>empty_argument_list_is_recorded() {
<a href=#440 id=440 data-nosnippet>440</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request(<span class="string">"ADMIN plugin_function()"</span>));
<a href=#441 id=441 data-nosnippet>441</a> <span class="kw">let </span>event = AdminFunctionEvent::success(input, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>Value::UInt64(<span class="number">0</span>)));
<a href=#442 id=442 data-nosnippet>442</a>
<a href=#443 id=443 data-nosnippet>443</a> <span class="macro">assert_eq!</span>(
<a href=#444 id=444 data-nosnippet>444</a> event.json_payload().unwrap(),
<a href=#445 id=445 data-nosnippet>445</a> <span class="macro">json!</span>({
<a href=#446 id=446 data-nosnippet>446</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#447 id=447 data-nosnippet>447</a> <span class="string">"arguments"</span>: {<span class="string">"type"</span>: <span class="string">"list"</span>, <span class="string">"value"</span>: []},
<a href=#448 id=448 data-nosnippet>448</a> })
<a href=#449 id=449 data-nosnippet>449</a> );
<a href=#450 id=450 data-nosnippet>450</a> }
<a href=#451 id=451 data-nosnippet>451</a>
<a href=#452 id=452 data-nosnippet>452</a> <span class="attr">#[test]
<a href=#453 id=453 data-nosnippet>453</a> </span><span class="kw">fn </span>records_extended_sql_values_and_non_literal_arguments() {
<a href=#454 id=454 data-nosnippet>454</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request(
<a href=#455 id=455 data-nosnippet>455</a> <span class="string">"ADMIN plugin_function(1, true, NULL, X'48656c6c6f', table_name)"</span>,
<a href=#456 id=456 data-nosnippet>456</a> ));
<a href=#457 id=457 data-nosnippet>457</a> <span class="kw">let </span>event = AdminFunctionEvent::success(input, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>Value::UInt64(<span class="number">0</span>)));
<a href=#458 id=458 data-nosnippet>458</a>
<a href=#459 id=459 data-nosnippet>459</a> <span class="macro">assert_eq!</span>(
<a href=#460 id=460 data-nosnippet>460</a> event.json_payload().unwrap(),
<a href=#461 id=461 data-nosnippet>461</a> <span class="macro">json!</span>({
<a href=#462 id=462 data-nosnippet>462</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#463 id=463 data-nosnippet>463</a> <span class="string">"arguments"</span>: {
<a href=#464 id=464 data-nosnippet>464</a> <span class="string">"type"</span>: <span class="string">"list"</span>,
<a href=#465 id=465 data-nosnippet>465</a> <span class="string">"value"</span>: [<span class="number">1</span>, <span class="bool-val">true</span>, null, <span class="string">"48656c6c6f"</span>, <span class="string">"table_name"</span>],
<a href=#466 id=466 data-nosnippet>466</a> },
<a href=#467 id=467 data-nosnippet>467</a> })
<a href=#468 id=468 data-nosnippet>468</a> );
<a href=#469 id=469 data-nosnippet>469</a> }
<a href=#470 id=470 data-nosnippet>470</a>
<a href=#471 id=471 data-nosnippet>471</a> <span class="attr">#[test]
<a href=#472 id=472 data-nosnippet>472</a> </span><span class="kw">fn </span>string_literals_preserve_logical_content() {
<a href=#473 id=473 data-nosnippet>473</a> <span class="kw">let </span>values = [
<a href=#474 id=474 data-nosnippet>474</a> SqlValue::SingleQuotedString(<span class="string">"single"</span>.to_string()),
<a href=#475 id=475 data-nosnippet>475</a> SqlValue::DoubleQuotedString(<span class="string">"double"</span>.to_string()),
<a href=#476 id=476 data-nosnippet>476</a> SqlValue::DollarQuotedString(DollarQuotedString {
<a href=#477 id=477 data-nosnippet>477</a> value: <span class="string">"dollar"</span>.to_string(),
<a href=#478 id=478 data-nosnippet>478</a> tag: <span class="prelude-val">None</span>,
<a href=#479 id=479 data-nosnippet>479</a> }),
<a href=#480 id=480 data-nosnippet>480</a> SqlValue::SingleQuotedRawStringLiteral(<span class="string">"raw"</span>.to_string()),
<a href=#481 id=481 data-nosnippet>481</a> SqlValue::SingleQuotedByteStringLiteral(<span class="string">"bytes"</span>.to_string()),
<a href=#482 id=482 data-nosnippet>482</a> SqlValue::NationalStringLiteral(<span class="string">"national"</span>.to_string()),
<a href=#483 id=483 data-nosnippet>483</a> SqlValue::UnicodeStringLiteral(<span class="string">"unicode"</span>.to_string()),
<a href=#484 id=484 data-nosnippet>484</a> SqlValue::HexStringLiteral(<span class="string">"48656c6c6f"</span>.to_string()),
<a href=#485 id=485 data-nosnippet>485</a> ];
<a href=#486 id=486 data-nosnippet>486</a>
<a href=#487 id=487 data-nosnippet>487</a> <span class="kw">let </span>expected = [
<a href=#488 id=488 data-nosnippet>488</a> <span class="string">"single"</span>,
<a href=#489 id=489 data-nosnippet>489</a> <span class="string">"double"</span>,
<a href=#490 id=490 data-nosnippet>490</a> <span class="string">"dollar"</span>,
<a href=#491 id=491 data-nosnippet>491</a> <span class="string">"raw"</span>,
<a href=#492 id=492 data-nosnippet>492</a> <span class="string">"bytes"</span>,
<a href=#493 id=493 data-nosnippet>493</a> <span class="string">"national"</span>,
<a href=#494 id=494 data-nosnippet>494</a> <span class="string">"unicode"</span>,
<a href=#495 id=495 data-nosnippet>495</a> <span class="string">"48656c6c6f"</span>,
<a href=#496 id=496 data-nosnippet>496</a> ];
<a href=#497 id=497 data-nosnippet>497</a> <span class="kw">for </span>(value, expected) <span class="kw">in </span>values.iter().zip(expected) {
<a href=#498 id=498 data-nosnippet>498</a> <span class="macro">assert_eq!</span>(sql_value_to_json(value), <span class="macro">json!</span>(expected));
<a href=#499 id=499 data-nosnippet>499</a> }
<a href=#500 id=500 data-nosnippet>500</a> }
<a href=#501 id=501 data-nosnippet>501</a>
<a href=#502 id=502 data-nosnippet>502</a> <span class="attr">#[test]
<a href=#503 id=503 data-nosnippet>503</a> </span><span class="kw">fn </span>non_literal_arguments_use_sql_display() {
<a href=#504 id=504 data-nosnippet>504</a> <span class="kw">let </span>arguments = FunctionArgumentList {
<a href=#505 id=505 data-nosnippet>505</a> duplicate_treatment: <span class="prelude-val">None</span>,
<a href=#506 id=506 data-nosnippet>506</a> args: <span class="macro">vec!</span>[
<a href=#507 id=507 data-nosnippet>507</a> FunctionArg::Named {
<a href=#508 id=508 data-nosnippet>508</a> name: Ident::new(<span class="string">"named"</span>),
<a href=#509 id=509 data-nosnippet>509</a> arg: FunctionArgExpr::Expr(Expr::Identifier(Ident::new(<span class="string">"value"</span>))),
<a href=#510 id=510 data-nosnippet>510</a> operator: FunctionArgOperator::RightArrow,
<a href=#511 id=511 data-nosnippet>511</a> },
<a href=#512 id=512 data-nosnippet>512</a> FunctionArg::Unnamed(FunctionArgExpr::Wildcard),
<a href=#513 id=513 data-nosnippet>513</a> ],
<a href=#514 id=514 data-nosnippet>514</a> clauses: <span class="macro">vec!</span>[],
<a href=#515 id=515 data-nosnippet>515</a> };
<a href=#516 id=516 data-nosnippet>516</a> <span class="kw">let </span>input = AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request_with_arguments(
<a href=#517 id=517 data-nosnippet>517</a> FunctionArguments::List(arguments),
<a href=#518 id=518 data-nosnippet>518</a> ));
<a href=#519 id=519 data-nosnippet>519</a> <span class="kw">let </span>event = AdminFunctionEvent::success(input, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>Value::UInt64(<span class="number">0</span>)));
<a href=#520 id=520 data-nosnippet>520</a>
<a href=#521 id=521 data-nosnippet>521</a> <span class="macro">assert_eq!</span>(
<a href=#522 id=522 data-nosnippet>522</a> event.json_payload().unwrap(),
<a href=#523 id=523 data-nosnippet>523</a> <span class="macro">json!</span>({
<a href=#524 id=524 data-nosnippet>524</a> <span class="string">"version"</span>: <span class="number">1</span>,
<a href=#525 id=525 data-nosnippet>525</a> <span class="string">"arguments"</span>: {<span class="string">"type"</span>: <span class="string">"list"</span>, <span class="string">"value"</span>: [<span class="string">"named =&gt; value"</span>, <span class="string">"*"</span>]},
<a href=#526 id=526 data-nosnippet>526</a> })
<a href=#527 id=527 data-nosnippet>527</a> );
<a href=#528 id=528 data-nosnippet>528</a> <span class="macro">assert_eq!</span>(
<a href=#529 id=529 data-nosnippet>529</a> sql_value_to_json(<span class="kw-2">&amp;</span>SqlValue::Placeholder(<span class="string">"$1"</span>.to_string())),
<a href=#530 id=530 data-nosnippet>530</a> <span class="macro">json!</span>(<span class="string">"$1"</span>)
<a href=#531 id=531 data-nosnippet>531</a> );
<a href=#532 id=532 data-nosnippet>532</a> }
<a href=#533 id=533 data-nosnippet>533</a>
<a href=#534 id=534 data-nosnippet>534</a> <span class="attr">#[test]
<a href=#535 id=535 data-nosnippet>535</a> </span><span class="kw">fn </span>serializes_binary_result() {
<a href=#536 id=536 data-nosnippet>536</a> <span class="macro">assert_eq!</span>(
<a href=#537 id=537 data-nosnippet>537</a> value_to_json(<span class="kw-2">&amp;</span>Value::from(<span class="macro">vec!</span>[<span class="number">1_u8</span>, <span class="number">2</span>, <span class="number">3</span>])),
<a href=#538 id=538 data-nosnippet>538</a> <span class="macro">json!</span>([<span class="number">1</span>, <span class="number">2</span>, <span class="number">3</span>])
<a href=#539 id=539 data-nosnippet>539</a> );
<a href=#540 id=540 data-nosnippet>540</a> }
<a href=#541 id=541 data-nosnippet>541</a>
<a href=#542 id=542 data-nosnippet>542</a> <span class="attr">#[test]
<a href=#543 id=543 data-nosnippet>543</a> </span><span class="kw">fn </span>serializes_non_finite_float_results_as_strings() {
<a href=#544 id=544 data-nosnippet>544</a> <span class="macro">assert_eq!</span>(
<a href=#545 id=545 data-nosnippet>545</a> value_to_json(<span class="kw-2">&amp;</span>Value::Float32(f32::NAN.into())),
<a href=#546 id=546 data-nosnippet>546</a> <span class="macro">json!</span>(<span class="string">"NaN"</span>)
<a href=#547 id=547 data-nosnippet>547</a> );
<a href=#548 id=548 data-nosnippet>548</a> <span class="macro">assert_eq!</span>(
<a href=#549 id=549 data-nosnippet>549</a> value_to_json(<span class="kw-2">&amp;</span>Value::Float32(f32::INFINITY.into())),
<a href=#550 id=550 data-nosnippet>550</a> <span class="macro">json!</span>(<span class="string">"inf"</span>)
<a href=#551 id=551 data-nosnippet>551</a> );
<a href=#552 id=552 data-nosnippet>552</a> <span class="macro">assert_eq!</span>(
<a href=#553 id=553 data-nosnippet>553</a> value_to_json(<span class="kw-2">&amp;</span>Value::Float32(f32::NEG_INFINITY.into())),
<a href=#554 id=554 data-nosnippet>554</a> <span class="macro">json!</span>(<span class="string">"-inf"</span>)
<a href=#555 id=555 data-nosnippet>555</a> );
<a href=#556 id=556 data-nosnippet>556</a> <span class="macro">assert_eq!</span>(
<a href=#557 id=557 data-nosnippet>557</a> value_to_json(<span class="kw-2">&amp;</span>Value::Float64(f64::NAN.into())),
<a href=#558 id=558 data-nosnippet>558</a> <span class="macro">json!</span>(<span class="string">"NaN"</span>)
<a href=#559 id=559 data-nosnippet>559</a> );
<a href=#560 id=560 data-nosnippet>560</a> <span class="macro">assert_eq!</span>(
<a href=#561 id=561 data-nosnippet>561</a> value_to_json(<span class="kw-2">&amp;</span>Value::Float64(f64::INFINITY.into())),
<a href=#562 id=562 data-nosnippet>562</a> <span class="macro">json!</span>(<span class="string">"inf"</span>)
<a href=#563 id=563 data-nosnippet>563</a> );
<a href=#564 id=564 data-nosnippet>564</a> <span class="macro">assert_eq!</span>(
<a href=#565 id=565 data-nosnippet>565</a> value_to_json(<span class="kw-2">&amp;</span>Value::Float64(f64::NEG_INFINITY.into())),
<a href=#566 id=566 data-nosnippet>566</a> <span class="macro">json!</span>(<span class="string">"-inf"</span>)
<a href=#567 id=567 data-nosnippet>567</a> );
<a href=#568 id=568 data-nosnippet>568</a> }
<a href=#569 id=569 data-nosnippet>569</a>}
</code></pre></div></section></main></body></html>
+441
View File
@@ -0,0 +1,441 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/admin/layer.rs`."><title>layer.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../../" data-static-root-path="../../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../../static.files/src-script-813739b1.js"></script><script defer src="../../../../src-files.js"></script><script defer src="../../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/admin/</div>layer.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::sync::{Arc, RwLock, Weak};
<a href=#16 id=16 data-nosnippet>16</a>
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_event_recorder::{EventRecorder, EventRecorderRef};
<a href=#18 id=18 data-nosnippet>18</a>
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span><span class="kw">crate</span>::error::Result;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span><span class="kw">crate</span>::statement::admin::event::{
<a href=#21 id=21 data-nosnippet>21</a> ADMIN_FUNCTION_EVENT_TYPE, AdminFunctionEvent, AdminFunctionEventInput,
<a href=#22 id=22 data-nosnippet>22</a>};
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span><span class="kw">crate</span>::statement::admin::{
<a href=#24 id=24 data-nosnippet>24</a> AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService, AdminFunctionServiceRef,
<a href=#25 id=25 data-nosnippet>25</a>};
<a href=#26 id=26 data-nosnippet>26</a>
<a href=#27 id=27 data-nosnippet>27</a><span class="doccomment">/// A composable layer over ADMIN function execution.
<a href=#28 id=28 data-nosnippet>28</a></span><span class="kw">pub trait </span>AdminFunctionLayer: Send + Sync {
<a href=#29 id=29 data-nosnippet>29</a> <span class="doccomment">/// Wraps `inner` and returns the resulting ADMIN function service.
<a href=#30 id=30 data-nosnippet>30</a> </span><span class="kw">fn </span>layer(<span class="kw-2">&amp;</span><span class="self">self</span>, inner: AdminFunctionServiceRef) -&gt; AdminFunctionServiceRef;
<a href=#31 id=31 data-nosnippet>31</a>}
<a href=#32 id=32 data-nosnippet>32</a>
<a href=#33 id=33 data-nosnippet>33</a><span class="doccomment">/// A shared composable ADMIN function layer.
<a href=#34 id=34 data-nosnippet>34</a></span><span class="kw">pub type </span>AdminFunctionLayerRef = Arc&lt;<span class="kw">dyn </span>AdminFunctionLayer&gt;;
<a href=#35 id=35 data-nosnippet>35</a>
<a href=#36 id=36 data-nosnippet>36</a><span class="doccomment">/// A late-bound weak reference to the event recorder.
<a href=#37 id=37 data-nosnippet>37</a>///
<a href=#38 id=38 data-nosnippet>38</a>/// The weak reference avoids retaining the frontend event handler, which owns a
<a href=#39 id=39 data-nosnippet>39</a>/// statement executor itself.
<a href=#40 id=40 data-nosnippet>40</a></span><span class="attr">#[derive(Clone, Default)]
<a href=#41 id=41 data-nosnippet>41</a></span><span class="kw">pub struct </span>AdminEventRecorderHandle {
<a href=#42 id=42 data-nosnippet>42</a> recorder: Arc&lt;RwLock&lt;<span class="prelude-ty">Option</span>&lt;Weak&lt;<span class="kw">dyn </span>EventRecorder&gt;&gt;&gt;&gt;,
<a href=#43 id=43 data-nosnippet>43</a>}
<a href=#44 id=44 data-nosnippet>44</a>
<a href=#45 id=45 data-nosnippet>45</a><span class="kw">impl </span>AdminEventRecorderHandle {
<a href=#46 id=46 data-nosnippet>46</a> <span class="doccomment">/// Installs the event recorder as a weak reference.
<a href=#47 id=47 data-nosnippet>47</a> </span><span class="kw">pub fn </span>install(<span class="kw-2">&amp;</span><span class="self">self</span>, recorder: <span class="kw-2">&amp;</span>EventRecorderRef) {
<a href=#48 id=48 data-nosnippet>48</a> <span class="kw-2">*</span><span class="self">self</span>.recorder.write().unwrap_or_else(|e| e.into_inner()) = <span class="prelude-val">Some</span>(Arc::downgrade(recorder));
<a href=#49 id=49 data-nosnippet>49</a> }
<a href=#50 id=50 data-nosnippet>50</a>
<a href=#51 id=51 data-nosnippet>51</a> <span class="kw">fn </span>upgrade(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; <span class="prelude-ty">Option</span>&lt;EventRecorderRef&gt; {
<a href=#52 id=52 data-nosnippet>52</a> <span class="self">self</span>.recorder
<a href=#53 id=53 data-nosnippet>53</a> .read()
<a href=#54 id=54 data-nosnippet>54</a> .unwrap_or_else(|e| e.into_inner())
<a href=#55 id=55 data-nosnippet>55</a> .as_ref()
<a href=#56 id=56 data-nosnippet>56</a> .and_then(Weak::upgrade)
<a href=#57 id=57 data-nosnippet>57</a> }
<a href=#58 id=58 data-nosnippet>58</a>}
<a href=#59 id=59 data-nosnippet>59</a>
<a href=#60 id=60 data-nosnippet>60</a><span class="doccomment">/// A layer that records ADMIN function inputs and outcomes as events.
<a href=#61 id=61 data-nosnippet>61</a></span><span class="attr">#[derive(Clone)]
<a href=#62 id=62 data-nosnippet>62</a></span><span class="kw">pub struct </span>AdminFunctionRecordingLayer {
<a href=#63 id=63 data-nosnippet>63</a> recorder: AdminEventRecorderHandle,
<a href=#64 id=64 data-nosnippet>64</a>}
<a href=#65 id=65 data-nosnippet>65</a>
<a href=#66 id=66 data-nosnippet>66</a><span class="kw">impl </span>AdminFunctionRecordingLayer {
<a href=#67 id=67 data-nosnippet>67</a> <span class="doccomment">/// Creates a recording layer backed by the late-bound recorder handle.
<a href=#68 id=68 data-nosnippet>68</a> </span><span class="kw">pub fn </span>new(recorder: AdminEventRecorderHandle) -&gt; <span class="self">Self </span>{
<a href=#69 id=69 data-nosnippet>69</a> <span class="self">Self </span>{ recorder }
<a href=#70 id=70 data-nosnippet>70</a> }
<a href=#71 id=71 data-nosnippet>71</a>}
<a href=#72 id=72 data-nosnippet>72</a>
<a href=#73 id=73 data-nosnippet>73</a><span class="kw">impl </span>AdminFunctionLayer <span class="kw">for </span>AdminFunctionRecordingLayer {
<a href=#74 id=74 data-nosnippet>74</a> <span class="kw">fn </span>layer(<span class="kw-2">&amp;</span><span class="self">self</span>, inner: AdminFunctionServiceRef) -&gt; AdminFunctionServiceRef {
<a href=#75 id=75 data-nosnippet>75</a> Arc::new(AdminFunctionRecordingService {
<a href=#76 id=76 data-nosnippet>76</a> inner,
<a href=#77 id=77 data-nosnippet>77</a> recorder: <span class="self">self</span>.recorder.clone(),
<a href=#78 id=78 data-nosnippet>78</a> })
<a href=#79 id=79 data-nosnippet>79</a> }
<a href=#80 id=80 data-nosnippet>80</a>}
<a href=#81 id=81 data-nosnippet>81</a>
<a href=#82 id=82 data-nosnippet>82</a><span class="attr">#[derive(Clone)]
<a href=#83 id=83 data-nosnippet>83</a></span><span class="kw">struct </span>AdminFunctionRecordingService {
<a href=#84 id=84 data-nosnippet>84</a> inner: AdminFunctionServiceRef,
<a href=#85 id=85 data-nosnippet>85</a> recorder: AdminEventRecorderHandle,
<a href=#86 id=86 data-nosnippet>86</a>}
<a href=#87 id=87 data-nosnippet>87</a>
<a href=#88 id=88 data-nosnippet>88</a><span class="attr">#[async_trait::async_trait]
<a href=#89 id=89 data-nosnippet>89</a></span><span class="kw">impl </span>AdminFunctionService <span class="kw">for </span>AdminFunctionRecordingService {
<a href=#90 id=90 data-nosnippet>90</a> <span class="kw">async fn </span>call(<span class="kw-2">&amp;</span><span class="self">self</span>, request: AdminFunctionRequest) -&gt; <span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt; {
<a href=#91 id=91 data-nosnippet>91</a> <span class="comment">// Avoid building the event input when this event type is disabled. Do
<a href=#92 id=92 data-nosnippet>92</a> // not retain the recorder across ADMIN execution because its handler can
<a href=#93 id=93 data-nosnippet>93</a> // own the statement executor that contains this layer.
<a href=#94 id=94 data-nosnippet>94</a> </span><span class="kw">let </span>enabled = <span class="self">self</span>.recorder.upgrade().is_some_and(|recorder| {
<a href=#95 id=95 data-nosnippet>95</a> recorder
<a href=#96 id=96 data-nosnippet>96</a> .event_type_filter()
<a href=#97 id=97 data-nosnippet>97</a> .allows(ADMIN_FUNCTION_EVENT_TYPE)
<a href=#98 id=98 data-nosnippet>98</a> });
<a href=#99 id=99 data-nosnippet>99</a> <span class="kw">let </span><span class="kw-2">mut </span>recording = AdminFunctionRecordingGuard::new(
<a href=#100 id=100 data-nosnippet>100</a> enabled.then(|| AdminFunctionEventInput::from_request(<span class="kw-2">&amp;</span>request)),
<a href=#101 id=101 data-nosnippet>101</a> <span class="self">self</span>.recorder.clone(),
<a href=#102 id=102 data-nosnippet>102</a> );
<a href=#103 id=103 data-nosnippet>103</a> <span class="kw">let </span>result = <span class="self">self</span>.inner.call(request).<span class="kw">await</span>;
<a href=#104 id=104 data-nosnippet>104</a> recording.record_result(<span class="kw-2">&amp;</span>result);
<a href=#105 id=105 data-nosnippet>105</a> result
<a href=#106 id=106 data-nosnippet>106</a> }
<a href=#107 id=107 data-nosnippet>107</a>}
<a href=#108 id=108 data-nosnippet>108</a>
<a href=#109 id=109 data-nosnippet>109</a><span class="doccomment">/// Records an ADMIN event after completion or when an in-flight call is dropped.
<a href=#110 id=110 data-nosnippet>110</a></span><span class="kw">struct </span>AdminFunctionRecordingGuard {
<a href=#111 id=111 data-nosnippet>111</a> input: <span class="prelude-ty">Option</span>&lt;AdminFunctionEventInput&gt;,
<a href=#112 id=112 data-nosnippet>112</a> recorder: AdminEventRecorderHandle,
<a href=#113 id=113 data-nosnippet>113</a>}
<a href=#114 id=114 data-nosnippet>114</a>
<a href=#115 id=115 data-nosnippet>115</a><span class="kw">impl </span>AdminFunctionRecordingGuard {
<a href=#116 id=116 data-nosnippet>116</a> <span class="kw">fn </span>new(input: <span class="prelude-ty">Option</span>&lt;AdminFunctionEventInput&gt;, recorder: AdminEventRecorderHandle) -&gt; <span class="self">Self </span>{
<a href=#117 id=117 data-nosnippet>117</a> <span class="self">Self </span>{ input, recorder }
<a href=#118 id=118 data-nosnippet>118</a> }
<a href=#119 id=119 data-nosnippet>119</a>
<a href=#120 id=120 data-nosnippet>120</a> <span class="kw">fn </span>record_result(<span class="kw-2">&amp;mut </span><span class="self">self</span>, result: <span class="kw-2">&amp;</span><span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt;) {
<a href=#121 id=121 data-nosnippet>121</a> <span class="kw">let </span><span class="prelude-val">Some</span>(input) = <span class="self">self</span>.input.take() <span class="kw">else </span>{
<a href=#122 id=122 data-nosnippet>122</a> <span class="kw">return</span>;
<a href=#123 id=123 data-nosnippet>123</a> };
<a href=#124 id=124 data-nosnippet>124</a> <span class="kw">let </span><span class="prelude-val">Some</span>(recorder) = <span class="self">self</span>.recorder.upgrade() <span class="kw">else </span>{
<a href=#125 id=125 data-nosnippet>125</a> <span class="kw">return</span>;
<a href=#126 id=126 data-nosnippet>126</a> };
<a href=#127 id=127 data-nosnippet>127</a> <span class="kw">let </span>event = <span class="kw">match </span>result {
<a href=#128 id=128 data-nosnippet>128</a> <span class="prelude-val">Ok</span>(response) =&gt; AdminFunctionEvent::success(input, response.immediate_result.as_ref()),
<a href=#129 id=129 data-nosnippet>129</a> <span class="prelude-val">Err</span>(error) =&gt; AdminFunctionEvent::failure(input, error),
<a href=#130 id=130 data-nosnippet>130</a> };
<a href=#131 id=131 data-nosnippet>131</a> recorder.record(Box::new(event));
<a href=#132 id=132 data-nosnippet>132</a> }
<a href=#133 id=133 data-nosnippet>133</a>}
<a href=#134 id=134 data-nosnippet>134</a>
<a href=#135 id=135 data-nosnippet>135</a><span class="kw">impl </span>Drop <span class="kw">for </span>AdminFunctionRecordingGuard {
<a href=#136 id=136 data-nosnippet>136</a> <span class="kw">fn </span>drop(<span class="kw-2">&amp;mut </span><span class="self">self</span>) {
<a href=#137 id=137 data-nosnippet>137</a> <span class="kw">let </span><span class="prelude-val">Some</span>(input) = <span class="self">self</span>.input.take() <span class="kw">else </span>{
<a href=#138 id=138 data-nosnippet>138</a> <span class="kw">return</span>;
<a href=#139 id=139 data-nosnippet>139</a> };
<a href=#140 id=140 data-nosnippet>140</a> <span class="kw">let </span><span class="prelude-val">Some</span>(recorder) = <span class="self">self</span>.recorder.upgrade() <span class="kw">else </span>{
<a href=#141 id=141 data-nosnippet>141</a> <span class="kw">return</span>;
<a href=#142 id=142 data-nosnippet>142</a> };
<a href=#143 id=143 data-nosnippet>143</a> recorder.record(Box::new(AdminFunctionEvent::cancelled(input)));
<a href=#144 id=144 data-nosnippet>144</a> }
<a href=#145 id=145 data-nosnippet>145</a>}
<a href=#146 id=146 data-nosnippet>146</a>
<a href=#147 id=147 data-nosnippet>147</a><span class="attr">#[cfg(test)]
<a href=#148 id=148 data-nosnippet>148</a></span><span class="kw">mod </span>tests {
<a href=#149 id=149 data-nosnippet>149</a> <span class="kw">use </span>std::collections::HashSet;
<a href=#150 id=150 data-nosnippet>150</a> <span class="kw">use </span>std::fmt::{Debug, Formatter};
<a href=#151 id=151 data-nosnippet>151</a> <span class="kw">use </span>std::sync::{Arc, Mutex};
<a href=#152 id=152 data-nosnippet>152</a> <span class="kw">use </span>std::time::Duration;
<a href=#153 id=153 data-nosnippet>153</a>
<a href=#154 id=154 data-nosnippet>154</a> <span class="kw">use </span>common_event_recorder::event_table::jsonb_value;
<a href=#155 id=155 data-nosnippet>155</a> <span class="kw">use </span>common_event_recorder::{
<a href=#156 id=156 data-nosnippet>156</a> Event, EventRecorder, EventRecorderRef, EventTypeFilter, EventTypeFilterRef,
<a href=#157 id=157 data-nosnippet>157</a> };
<a href=#158 id=158 data-nosnippet>158</a> <span class="kw">use </span>common_query::{Output, OutputData};
<a href=#159 id=159 data-nosnippet>159</a> <span class="kw">use </span>datatypes::value::Value;
<a href=#160 id=160 data-nosnippet>160</a> <span class="kw">use </span>serde_json::json;
<a href=#161 id=161 data-nosnippet>161</a> <span class="kw">use </span>session::context::QueryContext;
<a href=#162 id=162 data-nosnippet>162</a> <span class="kw">use </span>sql::dialect::GreptimeDbDialect;
<a href=#163 id=163 data-nosnippet>163</a> <span class="kw">use </span>sql::parser::{ParseOptions, ParserContext};
<a href=#164 id=164 data-nosnippet>164</a> <span class="kw">use </span>sql::statements::statement::Statement;
<a href=#165 id=165 data-nosnippet>165</a> <span class="kw">use </span>tokio::sync::Notify;
<a href=#166 id=166 data-nosnippet>166</a>
<a href=#167 id=167 data-nosnippet>167</a> <span class="kw">use </span><span class="kw">crate</span>::error::Result;
<a href=#168 id=168 data-nosnippet>168</a> <span class="kw">use </span><span class="kw">crate</span>::statement::admin::layer::{AdminEventRecorderHandle, AdminFunctionRecordingLayer};
<a href=#169 id=169 data-nosnippet>169</a> <span class="kw">use </span><span class="kw">crate</span>::statement::admin::{
<a href=#170 id=170 data-nosnippet>170</a> AdminFunctionLayer, AdminFunctionRequest, AdminFunctionResponse, AdminFunctionService,
<a href=#171 id=171 data-nosnippet>171</a> AdminFunctionServiceRef,
<a href=#172 id=172 data-nosnippet>172</a> };
<a href=#173 id=173 data-nosnippet>173</a>
<a href=#174 id=174 data-nosnippet>174</a> <span class="attr">#[derive(Clone)]
<a href=#175 id=175 data-nosnippet>175</a> </span><span class="kw">struct </span>TraceLayer {
<a href=#176 id=176 data-nosnippet>176</a> trace: Arc&lt;Mutex&lt;Vec&lt;<span class="kw-2">&amp;</span><span class="lifetime">'static </span>str&gt;&gt;&gt;,
<a href=#177 id=177 data-nosnippet>177</a> before: <span class="kw-2">&amp;</span><span class="lifetime">'static </span>str,
<a href=#178 id=178 data-nosnippet>178</a> after: <span class="kw-2">&amp;</span><span class="lifetime">'static </span>str,
<a href=#179 id=179 data-nosnippet>179</a> }
<a href=#180 id=180 data-nosnippet>180</a>
<a href=#181 id=181 data-nosnippet>181</a> <span class="kw">impl </span>AdminFunctionLayer <span class="kw">for </span>TraceLayer {
<a href=#182 id=182 data-nosnippet>182</a> <span class="kw">fn </span>layer(<span class="kw-2">&amp;</span><span class="self">self</span>, inner: AdminFunctionServiceRef) -&gt; AdminFunctionServiceRef {
<a href=#183 id=183 data-nosnippet>183</a> Arc::new(TraceService {
<a href=#184 id=184 data-nosnippet>184</a> inner,
<a href=#185 id=185 data-nosnippet>185</a> trace: <span class="self">self</span>.trace.clone(),
<a href=#186 id=186 data-nosnippet>186</a> before: <span class="self">self</span>.before,
<a href=#187 id=187 data-nosnippet>187</a> after: <span class="self">self</span>.after,
<a href=#188 id=188 data-nosnippet>188</a> })
<a href=#189 id=189 data-nosnippet>189</a> }
<a href=#190 id=190 data-nosnippet>190</a> }
<a href=#191 id=191 data-nosnippet>191</a>
<a href=#192 id=192 data-nosnippet>192</a> <span class="kw">struct </span>TraceService {
<a href=#193 id=193 data-nosnippet>193</a> inner: AdminFunctionServiceRef,
<a href=#194 id=194 data-nosnippet>194</a> trace: Arc&lt;Mutex&lt;Vec&lt;<span class="kw-2">&amp;</span><span class="lifetime">'static </span>str&gt;&gt;&gt;,
<a href=#195 id=195 data-nosnippet>195</a> before: <span class="kw-2">&amp;</span><span class="lifetime">'static </span>str,
<a href=#196 id=196 data-nosnippet>196</a> after: <span class="kw-2">&amp;</span><span class="lifetime">'static </span>str,
<a href=#197 id=197 data-nosnippet>197</a> }
<a href=#198 id=198 data-nosnippet>198</a>
<a href=#199 id=199 data-nosnippet>199</a> <span class="attr">#[async_trait::async_trait]
<a href=#200 id=200 data-nosnippet>200</a> </span><span class="kw">impl </span>AdminFunctionService <span class="kw">for </span>TraceService {
<a href=#201 id=201 data-nosnippet>201</a> <span class="kw">async fn </span>call(<span class="kw-2">&amp;</span><span class="self">self</span>, request: AdminFunctionRequest) -&gt; <span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt; {
<a href=#202 id=202 data-nosnippet>202</a> <span class="self">self</span>.trace.lock().unwrap().push(<span class="self">self</span>.before);
<a href=#203 id=203 data-nosnippet>203</a> <span class="kw">let </span>result = <span class="self">self</span>.inner.call(request).<span class="kw">await</span>;
<a href=#204 id=204 data-nosnippet>204</a> <span class="self">self</span>.trace.lock().unwrap().push(<span class="self">self</span>.after);
<a href=#205 id=205 data-nosnippet>205</a> result
<a href=#206 id=206 data-nosnippet>206</a> }
<a href=#207 id=207 data-nosnippet>207</a> }
<a href=#208 id=208 data-nosnippet>208</a>
<a href=#209 id=209 data-nosnippet>209</a> <span class="kw">struct </span>TestService {
<a href=#210 id=210 data-nosnippet>210</a> trace: <span class="prelude-ty">Option</span>&lt;Arc&lt;Mutex&lt;Vec&lt;<span class="kw-2">&amp;</span><span class="lifetime">'static </span>str&gt;&gt;&gt;&gt;,
<a href=#211 id=211 data-nosnippet>211</a> immediate_result: <span class="prelude-ty">Option</span>&lt;Value&gt;,
<a href=#212 id=212 data-nosnippet>212</a> }
<a href=#213 id=213 data-nosnippet>213</a>
<a href=#214 id=214 data-nosnippet>214</a> <span class="attr">#[async_trait::async_trait]
<a href=#215 id=215 data-nosnippet>215</a> </span><span class="kw">impl </span>AdminFunctionService <span class="kw">for </span>TestService {
<a href=#216 id=216 data-nosnippet>216</a> <span class="kw">async fn </span>call(<span class="kw-2">&amp;</span><span class="self">self</span>, _request: AdminFunctionRequest) -&gt; <span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt; {
<a href=#217 id=217 data-nosnippet>217</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(trace) = <span class="kw-2">&amp;</span><span class="self">self</span>.trace {
<a href=#218 id=218 data-nosnippet>218</a> trace.lock().unwrap().push(<span class="string">"core"</span>);
<a href=#219 id=219 data-nosnippet>219</a> }
<a href=#220 id=220 data-nosnippet>220</a> <span class="prelude-val">Ok</span>(response(<span class="self">self</span>.immediate_result.clone()))
<a href=#221 id=221 data-nosnippet>221</a> }
<a href=#222 id=222 data-nosnippet>222</a> }
<a href=#223 id=223 data-nosnippet>223</a>
<a href=#224 id=224 data-nosnippet>224</a> <span class="kw">struct </span>PendingService {
<a href=#225 id=225 data-nosnippet>225</a> started: Arc&lt;Notify&gt;,
<a href=#226 id=226 data-nosnippet>226</a> }
<a href=#227 id=227 data-nosnippet>227</a>
<a href=#228 id=228 data-nosnippet>228</a> <span class="attr">#[async_trait::async_trait]
<a href=#229 id=229 data-nosnippet>229</a> </span><span class="kw">impl </span>AdminFunctionService <span class="kw">for </span>PendingService {
<a href=#230 id=230 data-nosnippet>230</a> <span class="kw">async fn </span>call(<span class="kw-2">&amp;</span><span class="self">self</span>, _request: AdminFunctionRequest) -&gt; <span class="prelude-ty">Result</span>&lt;AdminFunctionResponse&gt; {
<a href=#231 id=231 data-nosnippet>231</a> <span class="self">self</span>.started.notify_one();
<a href=#232 id=232 data-nosnippet>232</a> std::future::pending().<span class="kw">await
<a href=#233 id=233 data-nosnippet>233</a> </span>}
<a href=#234 id=234 data-nosnippet>234</a> }
<a href=#235 id=235 data-nosnippet>235</a>
<a href=#236 id=236 data-nosnippet>236</a> <span class="kw">struct </span>TestRecorder {
<a href=#237 id=237 data-nosnippet>237</a> events: Mutex&lt;Vec&lt;Box&lt;<span class="kw">dyn </span>Event&gt;&gt;&gt;,
<a href=#238 id=238 data-nosnippet>238</a> filter: EventTypeFilterRef,
<a href=#239 id=239 data-nosnippet>239</a> }
<a href=#240 id=240 data-nosnippet>240</a>
<a href=#241 id=241 data-nosnippet>241</a> <span class="kw">impl </span>TestRecorder {
<a href=#242 id=242 data-nosnippet>242</a> <span class="kw">fn </span>new(filter: EventTypeFilter) -&gt; <span class="self">Self </span>{
<a href=#243 id=243 data-nosnippet>243</a> <span class="self">Self </span>{
<a href=#244 id=244 data-nosnippet>244</a> events: Mutex::new(Vec::new()),
<a href=#245 id=245 data-nosnippet>245</a> filter: Arc::new(filter),
<a href=#246 id=246 data-nosnippet>246</a> }
<a href=#247 id=247 data-nosnippet>247</a> }
<a href=#248 id=248 data-nosnippet>248</a> }
<a href=#249 id=249 data-nosnippet>249</a>
<a href=#250 id=250 data-nosnippet>250</a> <span class="kw">impl </span>Debug <span class="kw">for </span>TestRecorder {
<a href=#251 id=251 data-nosnippet>251</a> <span class="kw">fn </span>fmt(<span class="kw-2">&amp;</span><span class="self">self</span>, f: <span class="kw-2">&amp;mut </span>Formatter&lt;<span class="lifetime">'_</span>&gt;) -&gt; std::fmt::Result {
<a href=#252 id=252 data-nosnippet>252</a> f.debug_struct(<span class="string">"TestRecorder"</span>).finish()
<a href=#253 id=253 data-nosnippet>253</a> }
<a href=#254 id=254 data-nosnippet>254</a> }
<a href=#255 id=255 data-nosnippet>255</a>
<a href=#256 id=256 data-nosnippet>256</a> <span class="kw">impl </span>EventRecorder <span class="kw">for </span>TestRecorder {
<a href=#257 id=257 data-nosnippet>257</a> <span class="kw">fn </span>record(<span class="kw-2">&amp;</span><span class="self">self</span>, event: Box&lt;<span class="kw">dyn </span>Event&gt;) {
<a href=#258 id=258 data-nosnippet>258</a> <span class="self">self</span>.events.lock().unwrap().push(event);
<a href=#259 id=259 data-nosnippet>259</a> }
<a href=#260 id=260 data-nosnippet>260</a>
<a href=#261 id=261 data-nosnippet>261</a> <span class="kw">fn </span>event_type_filter(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; EventTypeFilterRef {
<a href=#262 id=262 data-nosnippet>262</a> <span class="self">self</span>.filter.clone()
<a href=#263 id=263 data-nosnippet>263</a> }
<a href=#264 id=264 data-nosnippet>264</a>
<a href=#265 id=265 data-nosnippet>265</a> <span class="kw">fn </span>close(<span class="kw-2">&amp;</span><span class="self">self</span>) {}
<a href=#266 id=266 data-nosnippet>266</a> }
<a href=#267 id=267 data-nosnippet>267</a>
<a href=#268 id=268 data-nosnippet>268</a> <span class="attr">#[tokio::test]
<a href=#269 id=269 data-nosnippet>269</a> </span><span class="kw">async fn </span>last_added_layer_runs_outermost() {
<a href=#270 id=270 data-nosnippet>270</a> <span class="kw">let </span>trace = Arc::new(Mutex::new(Vec::new()));
<a href=#271 id=271 data-nosnippet>271</a> <span class="kw">let </span><span class="kw-2">mut </span>service: AdminFunctionServiceRef = Arc::new(TestService {
<a href=#272 id=272 data-nosnippet>272</a> trace: <span class="prelude-val">Some</span>(trace.clone()),
<a href=#273 id=273 data-nosnippet>273</a> immediate_result: <span class="prelude-val">Some</span>(Value::UInt64(<span class="number">0</span>)),
<a href=#274 id=274 data-nosnippet>274</a> });
<a href=#275 id=275 data-nosnippet>275</a> <span class="kw">for </span>(before, after) <span class="kw">in </span>[(<span class="string">"a_before"</span>, <span class="string">"a_after"</span>), (<span class="string">"b_before"</span>, <span class="string">"b_after"</span>)] {
<a href=#276 id=276 data-nosnippet>276</a> service = TraceLayer {
<a href=#277 id=277 data-nosnippet>277</a> trace: trace.clone(),
<a href=#278 id=278 data-nosnippet>278</a> before,
<a href=#279 id=279 data-nosnippet>279</a> after,
<a href=#280 id=280 data-nosnippet>280</a> }
<a href=#281 id=281 data-nosnippet>281</a> .layer(service);
<a href=#282 id=282 data-nosnippet>282</a> }
<a href=#283 id=283 data-nosnippet>283</a>
<a href=#284 id=284 data-nosnippet>284</a> service.call(request()).<span class="kw">await</span>.unwrap();
<a href=#285 id=285 data-nosnippet>285</a> <span class="macro">assert_eq!</span>(
<a href=#286 id=286 data-nosnippet>286</a> <span class="kw-2">*</span>trace.lock().unwrap(),
<a href=#287 id=287 data-nosnippet>287</a> [<span class="string">"b_before"</span>, <span class="string">"a_before"</span>, <span class="string">"core"</span>, <span class="string">"a_after"</span>, <span class="string">"b_after"</span>]
<a href=#288 id=288 data-nosnippet>288</a> );
<a href=#289 id=289 data-nosnippet>289</a> }
<a href=#290 id=290 data-nosnippet>290</a>
<a href=#291 id=291 data-nosnippet>291</a> <span class="attr">#[tokio::test]
<a href=#292 id=292 data-nosnippet>292</a> </span><span class="kw">async fn </span>recording_respects_filter_and_weak_lifetime() {
<a href=#293 id=293 data-nosnippet>293</a> <span class="kw">let </span>enabled = Arc::new(TestRecorder::new(EventTypeFilter::All));
<a href=#294 id=294 data-nosnippet>294</a> <span class="kw">let </span>enabled_ref: EventRecorderRef = enabled.clone();
<a href=#295 id=295 data-nosnippet>295</a> <span class="kw">let </span>handle = AdminEventRecorderHandle::default();
<a href=#296 id=296 data-nosnippet>296</a> handle.install(<span class="kw-2">&amp;</span>enabled_ref);
<a href=#297 id=297 data-nosnippet>297</a> recording_service(handle.clone())
<a href=#298 id=298 data-nosnippet>298</a> .call(request())
<a href=#299 id=299 data-nosnippet>299</a> .<span class="kw">await
<a href=#300 id=300 data-nosnippet>300</a> </span>.unwrap();
<a href=#301 id=301 data-nosnippet>301</a> <span class="macro">assert_eq!</span>(enabled.events.lock().unwrap().len(), <span class="number">1</span>);
<a href=#302 id=302 data-nosnippet>302</a>
<a href=#303 id=303 data-nosnippet>303</a> <span class="kw">let </span>disabled = Arc::new(TestRecorder::new(EventTypeFilter::Only(HashSet::new())));
<a href=#304 id=304 data-nosnippet>304</a> <span class="kw">let </span>disabled_ref: EventRecorderRef = disabled.clone();
<a href=#305 id=305 data-nosnippet>305</a> handle.install(<span class="kw-2">&amp;</span>disabled_ref);
<a href=#306 id=306 data-nosnippet>306</a> recording_service(handle.clone())
<a href=#307 id=307 data-nosnippet>307</a> .call(request())
<a href=#308 id=308 data-nosnippet>308</a> .<span class="kw">await
<a href=#309 id=309 data-nosnippet>309</a> </span>.unwrap();
<a href=#310 id=310 data-nosnippet>310</a> <span class="macro">assert!</span>(disabled.events.lock().unwrap().is_empty());
<a href=#311 id=311 data-nosnippet>311</a>
<a href=#312 id=312 data-nosnippet>312</a> drop(disabled_ref);
<a href=#313 id=313 data-nosnippet>313</a> drop(disabled);
<a href=#314 id=314 data-nosnippet>314</a> <span class="macro">assert!</span>(handle.upgrade().is_none());
<a href=#315 id=315 data-nosnippet>315</a> recording_service(handle).call(request()).<span class="kw">await</span>.unwrap();
<a href=#316 id=316 data-nosnippet>316</a>
<a href=#317 id=317 data-nosnippet>317</a> <span class="kw">let </span>recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
<a href=#318 id=318 data-nosnippet>318</a> <span class="kw">let </span>recorder_ref: EventRecorderRef = recorder.clone();
<a href=#319 id=319 data-nosnippet>319</a> <span class="kw">let </span>handle = AdminEventRecorderHandle::default();
<a href=#320 id=320 data-nosnippet>320</a> handle.install(<span class="kw-2">&amp;</span>recorder_ref);
<a href=#321 id=321 data-nosnippet>321</a> <span class="kw">let </span>service = recording_service(handle.clone());
<a href=#322 id=322 data-nosnippet>322</a> <span class="kw">let </span>future = service.call(request());
<a href=#323 id=323 data-nosnippet>323</a> drop(recorder_ref);
<a href=#324 id=324 data-nosnippet>324</a> drop(recorder);
<a href=#325 id=325 data-nosnippet>325</a> <span class="macro">assert!</span>(handle.upgrade().is_none());
<a href=#326 id=326 data-nosnippet>326</a> future.<span class="kw">await</span>.unwrap();
<a href=#327 id=327 data-nosnippet>327</a> }
<a href=#328 id=328 data-nosnippet>328</a>
<a href=#329 id=329 data-nosnippet>329</a> <span class="attr">#[tokio::test]
<a href=#330 id=330 data-nosnippet>330</a> </span><span class="kw">async fn </span>recording_does_not_retain_recorder_while_pending() {
<a href=#331 id=331 data-nosnippet>331</a> <span class="kw">let </span>recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
<a href=#332 id=332 data-nosnippet>332</a> <span class="kw">let </span>recorder_ref: EventRecorderRef = recorder.clone();
<a href=#333 id=333 data-nosnippet>333</a> <span class="kw">let </span>handle = AdminEventRecorderHandle::default();
<a href=#334 id=334 data-nosnippet>334</a> handle.install(<span class="kw-2">&amp;</span>recorder_ref);
<a href=#335 id=335 data-nosnippet>335</a> <span class="kw">let </span>started = Arc::new(Notify::new());
<a href=#336 id=336 data-nosnippet>336</a> <span class="kw">let </span>service =
<a href=#337 id=337 data-nosnippet>337</a> AdminFunctionRecordingLayer::new(handle.clone()).layer(Arc::new(PendingService {
<a href=#338 id=338 data-nosnippet>338</a> started: started.clone(),
<a href=#339 id=339 data-nosnippet>339</a> }));
<a href=#340 id=340 data-nosnippet>340</a>
<a href=#341 id=341 data-nosnippet>341</a> <span class="kw">let </span>task = tokio::spawn(<span class="kw">async move </span>{ service.call(request()).<span class="kw">await </span>});
<a href=#342 id=342 data-nosnippet>342</a> started.notified().<span class="kw">await</span>;
<a href=#343 id=343 data-nosnippet>343</a> drop(recorder_ref);
<a href=#344 id=344 data-nosnippet>344</a> drop(recorder);
<a href=#345 id=345 data-nosnippet>345</a>
<a href=#346 id=346 data-nosnippet>346</a> <span class="macro">assert!</span>(handle.upgrade().is_none());
<a href=#347 id=347 data-nosnippet>347</a> task.abort();
<a href=#348 id=348 data-nosnippet>348</a> <span class="macro">assert!</span>(<span class="macro">matches!</span>(task.<span class="kw">await</span>, <span class="prelude-val">Err</span>(error) <span class="kw">if </span>error.is_cancelled()));
<a href=#349 id=349 data-nosnippet>349</a> }
<a href=#350 id=350 data-nosnippet>350</a>
<a href=#351 id=351 data-nosnippet>351</a> <span class="attr">#[tokio::test]
<a href=#352 id=352 data-nosnippet>352</a> </span><span class="kw">async fn </span>dropping_pending_call_records_failure() {
<a href=#353 id=353 data-nosnippet>353</a> <span class="kw">let </span>recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
<a href=#354 id=354 data-nosnippet>354</a> <span class="kw">let </span>recorder_ref: EventRecorderRef = recorder.clone();
<a href=#355 id=355 data-nosnippet>355</a> <span class="kw">let </span>handle = AdminEventRecorderHandle::default();
<a href=#356 id=356 data-nosnippet>356</a> handle.install(<span class="kw-2">&amp;</span>recorder_ref);
<a href=#357 id=357 data-nosnippet>357</a> <span class="kw">let </span>started = Arc::new(Notify::new());
<a href=#358 id=358 data-nosnippet>358</a> <span class="kw">let </span>service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService {
<a href=#359 id=359 data-nosnippet>359</a> started: started.clone(),
<a href=#360 id=360 data-nosnippet>360</a> }));
<a href=#361 id=361 data-nosnippet>361</a>
<a href=#362 id=362 data-nosnippet>362</a> <span class="kw">let </span>task = tokio::spawn(<span class="kw">async move </span>{ service.call(request()).<span class="kw">await </span>});
<a href=#363 id=363 data-nosnippet>363</a> started.notified().<span class="kw">await</span>;
<a href=#364 id=364 data-nosnippet>364</a> task.abort();
<a href=#365 id=365 data-nosnippet>365</a> <span class="macro">assert!</span>(<span class="macro">matches!</span>(task.<span class="kw">await</span>, <span class="prelude-val">Err</span>(error) <span class="kw">if </span>error.is_cancelled()));
<a href=#366 id=366 data-nosnippet>366</a> <span class="macro">assert_eq!</span>(recorder.events.lock().unwrap().len(), <span class="number">1</span>);
<a href=#367 id=367 data-nosnippet>367</a> }
<a href=#368 id=368 data-nosnippet>368</a>
<a href=#369 id=369 data-nosnippet>369</a> <span class="attr">#[tokio::test]
<a href=#370 id=370 data-nosnippet>370</a> </span><span class="kw">async fn </span>timing_out_pending_call_records_failure() {
<a href=#371 id=371 data-nosnippet>371</a> <span class="kw">let </span>recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
<a href=#372 id=372 data-nosnippet>372</a> <span class="kw">let </span>recorder_ref: EventRecorderRef = recorder.clone();
<a href=#373 id=373 data-nosnippet>373</a> <span class="kw">let </span>handle = AdminEventRecorderHandle::default();
<a href=#374 id=374 data-nosnippet>374</a> handle.install(<span class="kw-2">&amp;</span>recorder_ref);
<a href=#375 id=375 data-nosnippet>375</a> <span class="kw">let </span>service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(PendingService {
<a href=#376 id=376 data-nosnippet>376</a> started: Arc::new(Notify::new()),
<a href=#377 id=377 data-nosnippet>377</a> }));
<a href=#378 id=378 data-nosnippet>378</a>
<a href=#379 id=379 data-nosnippet>379</a> <span class="macro">assert!</span>(
<a href=#380 id=380 data-nosnippet>380</a> tokio::time::timeout(Duration::from_millis(<span class="number">10</span>), service.call(request()))
<a href=#381 id=381 data-nosnippet>381</a> .<span class="kw">await
<a href=#382 id=382 data-nosnippet>382</a> </span>.is_err()
<a href=#383 id=383 data-nosnippet>383</a> );
<a href=#384 id=384 data-nosnippet>384</a> <span class="macro">assert_eq!</span>(recorder.events.lock().unwrap().len(), <span class="number">1</span>);
<a href=#385 id=385 data-nosnippet>385</a> }
<a href=#386 id=386 data-nosnippet>386</a>
<a href=#387 id=387 data-nosnippet>387</a> <span class="attr">#[tokio::test]
<a href=#388 id=388 data-nosnippet>388</a> </span><span class="kw">async fn </span>empty_immediate_result_preserves_success_output_and_records_no_result() {
<a href=#389 id=389 data-nosnippet>389</a> <span class="kw">let </span>recorder = Arc::new(TestRecorder::new(EventTypeFilter::All));
<a href=#390 id=390 data-nosnippet>390</a> <span class="kw">let </span>recorder_ref: EventRecorderRef = recorder.clone();
<a href=#391 id=391 data-nosnippet>391</a> <span class="kw">let </span>handle = AdminEventRecorderHandle::default();
<a href=#392 id=392 data-nosnippet>392</a> handle.install(<span class="kw-2">&amp;</span>recorder_ref);
<a href=#393 id=393 data-nosnippet>393</a> <span class="kw">let </span>service = AdminFunctionRecordingLayer::new(handle).layer(Arc::new(TestService {
<a href=#394 id=394 data-nosnippet>394</a> trace: <span class="prelude-val">None</span>,
<a href=#395 id=395 data-nosnippet>395</a> immediate_result: <span class="prelude-val">None</span>,
<a href=#396 id=396 data-nosnippet>396</a> }));
<a href=#397 id=397 data-nosnippet>397</a>
<a href=#398 id=398 data-nosnippet>398</a> <span class="kw">let </span>response = service.call(request()).<span class="kw">await</span>.unwrap();
<a href=#399 id=399 data-nosnippet>399</a> <span class="macro">assert!</span>(<span class="macro">matches!</span>(response.output.data, OutputData::AffectedRows(<span class="number">0</span>)));
<a href=#400 id=400 data-nosnippet>400</a> <span class="macro">assert_eq!</span>(response.immediate_result, <span class="prelude-val">None</span>);
<a href=#401 id=401 data-nosnippet>401</a>
<a href=#402 id=402 data-nosnippet>402</a> <span class="kw">let </span>events = recorder.events.lock().unwrap();
<a href=#403 id=403 data-nosnippet>403</a> <span class="macro">assert_eq!</span>(events.len(), <span class="number">1</span>);
<a href=#404 id=404 data-nosnippet>404</a> <span class="macro">assert_eq!</span>(
<a href=#405 id=405 data-nosnippet>405</a> events[<span class="number">0</span>].extra_rows().unwrap()[<span class="number">0</span>].values[<span class="number">3</span>],
<a href=#406 id=406 data-nosnippet>406</a> jsonb_value(<span class="kw-2">&amp;</span><span class="macro">json!</span>({}))
<a href=#407 id=407 data-nosnippet>407</a> );
<a href=#408 id=408 data-nosnippet>408</a> }
<a href=#409 id=409 data-nosnippet>409</a>
<a href=#410 id=410 data-nosnippet>410</a> <span class="kw">fn </span>recording_service(handle: AdminEventRecorderHandle) -&gt; AdminFunctionServiceRef {
<a href=#411 id=411 data-nosnippet>411</a> <span class="kw">let </span>core = Arc::new(TestService {
<a href=#412 id=412 data-nosnippet>412</a> trace: <span class="prelude-val">None</span>,
<a href=#413 id=413 data-nosnippet>413</a> immediate_result: <span class="prelude-val">Some</span>(Value::UInt64(<span class="number">0</span>)),
<a href=#414 id=414 data-nosnippet>414</a> });
<a href=#415 id=415 data-nosnippet>415</a> AdminFunctionRecordingLayer::new(handle).layer(core)
<a href=#416 id=416 data-nosnippet>416</a> }
<a href=#417 id=417 data-nosnippet>417</a>
<a href=#418 id=418 data-nosnippet>418</a> <span class="kw">fn </span>response(immediate_result: <span class="prelude-ty">Option</span>&lt;Value&gt;) -&gt; AdminFunctionResponse {
<a href=#419 id=419 data-nosnippet>419</a> AdminFunctionResponse {
<a href=#420 id=420 data-nosnippet>420</a> output: Output::new_with_affected_rows(<span class="number">0</span>),
<a href=#421 id=421 data-nosnippet>421</a> immediate_result,
<a href=#422 id=422 data-nosnippet>422</a> }
<a href=#423 id=423 data-nosnippet>423</a> }
<a href=#424 id=424 data-nosnippet>424</a>
<a href=#425 id=425 data-nosnippet>425</a> <span class="kw">fn </span>request() -&gt; AdminFunctionRequest {
<a href=#426 id=426 data-nosnippet>426</a> <span class="kw">let </span>Statement::Admin(statement) = ParserContext::create_with_dialect(
<a href=#427 id=427 data-nosnippet>427</a> <span class="string">"ADMIN flush_table('demo')"</span>,
<a href=#428 id=428 data-nosnippet>428</a> <span class="kw-2">&amp;</span>GreptimeDbDialect {},
<a href=#429 id=429 data-nosnippet>429</a> ParseOptions::default(),
<a href=#430 id=430 data-nosnippet>430</a> )
<a href=#431 id=431 data-nosnippet>431</a> .unwrap()
<a href=#432 id=432 data-nosnippet>432</a> .remove(<span class="number">0</span>) <span class="kw">else </span>{
<a href=#433 id=433 data-nosnippet>433</a> <span class="macro">panic!</span>(<span class="string">"expected ADMIN statement"</span>)
<a href=#434 id=434 data-nosnippet>434</a> };
<a href=#435 id=435 data-nosnippet>435</a> AdminFunctionRequest {
<a href=#436 id=436 data-nosnippet>436</a> statement,
<a href=#437 id=437 data-nosnippet>437</a> query_ctx: QueryContext::arc(),
<a href=#438 id=438 data-nosnippet>438</a> }
<a href=#439 id=439 data-nosnippet>439</a> }
<a href=#440 id=440 data-nosnippet>440</a>}
</code></pre></div></section></main></body></html>
+208
View File
@@ -0,0 +1,208 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/comment.rs`."><title>comment.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>comment.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>api::v1::CommentOnExpr;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>common_error::ext::BoxedError;
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_meta::cache_invalidator::Context;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_meta::rpc::ddl::{
<a href=#19 id=19 data-nosnippet>19</a> CommentObjectType, CommentOnTask, DdlTask, SubmitDdlTaskRequest, TriggerReason,
<a href=#20 id=20 data-nosnippet>20</a>};
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>common_query::Output;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>session::table_name::table_idents_to_full_name;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>snafu::ResultExt;
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>sql::ast::ObjectNamePartExt;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>sql::statements::comment::{Comment, CommentObject};
<a href=#27 id=27 data-nosnippet>27</a>
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span><span class="kw">crate</span>::error::{
<a href=#29 id=29 data-nosnippet>29</a> <span class="self">self</span>, ExecuteDdlSnafu, ExternalSnafu, InvalidSqlSnafu, <span class="prelude-ty">Result</span>, TableMetadataManagerSnafu,
<a href=#30 id=30 data-nosnippet>30</a>};
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#32 id=32 data-nosnippet>32</a><span class="kw">use </span><span class="kw">crate</span>::utils::to_executor_context;
<a href=#33 id=33 data-nosnippet>33</a>
<a href=#34 id=34 data-nosnippet>34</a><span class="kw">impl </span>StatementExecutor {
<a href=#35 id=35 data-nosnippet>35</a> <span class="doccomment">/// Adds a comment to a database object (table, column, or flow).
<a href=#36 id=36 data-nosnippet>36</a> ///
<a href=#37 id=37 data-nosnippet>37</a> /// # Arguments
<a href=#38 id=38 data-nosnippet>38</a> ///
<a href=#39 id=39 data-nosnippet>39</a> /// * `stmt`: A `Comment` struct containing the object to comment on and the comment text.
<a href=#40 id=40 data-nosnippet>40</a> /// * `query_ctx`: A `QueryContextRef` providing contextual information for the query.
<a href=#41 id=41 data-nosnippet>41</a> ///
<a href=#42 id=42 data-nosnippet>42</a> /// # Returns
<a href=#43 id=43 data-nosnippet>43</a> ///
<a href=#44 id=44 data-nosnippet>44</a> /// A `Result` containing the `Output` of the operation, or an error if the operation fails.
<a href=#45 id=45 data-nosnippet>45</a> </span><span class="kw">pub async fn </span>comment(<span class="kw-2">&amp;</span><span class="self">self</span>, stmt: Comment, query_ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#46 id=46 data-nosnippet>46</a> <span class="kw">let </span><span class="kw-2">mut </span>comment_on_task = <span class="self">self</span>.create_comment_on_task_from_stmt(stmt, <span class="kw-2">&amp;</span>query_ctx)<span class="question-mark">?</span>;
<a href=#47 id=47 data-nosnippet>47</a> comment_on_task
<a href=#48 id=48 data-nosnippet>48</a> .enrich_object_id(
<a href=#49 id=49 data-nosnippet>49</a> <span class="self">self</span>.table_metadata_manager.table_name_manager(),
<a href=#50 id=50 data-nosnippet>50</a> <span class="self">self</span>.flow_metadata_manager.flow_name_manager(),
<a href=#51 id=51 data-nosnippet>51</a> )
<a href=#52 id=52 data-nosnippet>52</a> .<span class="kw">await
<a href=#53 id=53 data-nosnippet>53</a> </span>.context(TableMetadataManagerSnafu)<span class="question-mark">?</span>;
<a href=#54 id=54 data-nosnippet>54</a> <span class="kw">let </span>cache_idents = comment_on_task.cache_idents();
<a href=#55 id=55 data-nosnippet>55</a>
<a href=#56 id=56 data-nosnippet>56</a> <span class="kw">let </span>executor_context = to_executor_context(query_ctx, TriggerReason::Manual);
<a href=#57 id=57 data-nosnippet>57</a> <span class="kw">let </span>request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task));
<a href=#58 id=58 data-nosnippet>58</a>
<a href=#59 id=59 data-nosnippet>59</a> <span class="self">self</span>.procedure_executor
<a href=#60 id=60 data-nosnippet>60</a> .submit_ddl_task(executor_context, request)
<a href=#61 id=61 data-nosnippet>61</a> .<span class="kw">await
<a href=#62 id=62 data-nosnippet>62</a> </span>.context(ExecuteDdlSnafu)<span class="question-mark">?</span>;
<a href=#63 id=63 data-nosnippet>63</a>
<a href=#64 id=64 data-nosnippet>64</a> <span class="comment">// Invalidates local cache ASAP.
<a href=#65 id=65 data-nosnippet>65</a> </span><span class="self">self</span>.cache_invalidator
<a href=#66 id=66 data-nosnippet>66</a> .invalidate(<span class="kw-2">&amp;</span>Context::default(), <span class="kw-2">&amp;</span>cache_idents)
<a href=#67 id=67 data-nosnippet>67</a> .<span class="kw">await
<a href=#68 id=68 data-nosnippet>68</a> </span>.context(error::InvalidateTableCacheSnafu)<span class="question-mark">?</span>;
<a href=#69 id=69 data-nosnippet>69</a>
<a href=#70 id=70 data-nosnippet>70</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="number">0</span>))
<a href=#71 id=71 data-nosnippet>71</a> }
<a href=#72 id=72 data-nosnippet>72</a>
<a href=#73 id=73 data-nosnippet>73</a> <span class="kw">pub async fn </span>comment_by_expr(
<a href=#74 id=74 data-nosnippet>74</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#75 id=75 data-nosnippet>75</a> expr: CommentOnExpr,
<a href=#76 id=76 data-nosnippet>76</a> query_ctx: QueryContextRef,
<a href=#77 id=77 data-nosnippet>77</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#78 id=78 data-nosnippet>78</a> <span class="kw">let </span><span class="kw-2">mut </span>comment_on_task = <span class="self">self</span>.create_comment_on_task_from_expr(expr)<span class="question-mark">?</span>;
<a href=#79 id=79 data-nosnippet>79</a> comment_on_task
<a href=#80 id=80 data-nosnippet>80</a> .enrich_object_id(
<a href=#81 id=81 data-nosnippet>81</a> <span class="self">self</span>.table_metadata_manager.table_name_manager(),
<a href=#82 id=82 data-nosnippet>82</a> <span class="self">self</span>.flow_metadata_manager.flow_name_manager(),
<a href=#83 id=83 data-nosnippet>83</a> )
<a href=#84 id=84 data-nosnippet>84</a> .<span class="kw">await
<a href=#85 id=85 data-nosnippet>85</a> </span>.context(TableMetadataManagerSnafu)<span class="question-mark">?</span>;
<a href=#86 id=86 data-nosnippet>86</a> <span class="kw">let </span>cache_idents = comment_on_task.cache_idents();
<a href=#87 id=87 data-nosnippet>87</a>
<a href=#88 id=88 data-nosnippet>88</a> <span class="kw">let </span>executor_context = to_executor_context(query_ctx, TriggerReason::Manual);
<a href=#89 id=89 data-nosnippet>89</a> <span class="kw">let </span>request = SubmitDdlTaskRequest::new(DdlTask::new_comment_on(comment_on_task));
<a href=#90 id=90 data-nosnippet>90</a>
<a href=#91 id=91 data-nosnippet>91</a> <span class="self">self</span>.procedure_executor
<a href=#92 id=92 data-nosnippet>92</a> .submit_ddl_task(executor_context, request)
<a href=#93 id=93 data-nosnippet>93</a> .<span class="kw">await
<a href=#94 id=94 data-nosnippet>94</a> </span>.context(ExecuteDdlSnafu)<span class="question-mark">?</span>;
<a href=#95 id=95 data-nosnippet>95</a>
<a href=#96 id=96 data-nosnippet>96</a> <span class="comment">// Invalidates local cache ASAP.
<a href=#97 id=97 data-nosnippet>97</a> </span><span class="self">self</span>.cache_invalidator
<a href=#98 id=98 data-nosnippet>98</a> .invalidate(<span class="kw-2">&amp;</span>Context::default(), <span class="kw-2">&amp;</span>cache_idents)
<a href=#99 id=99 data-nosnippet>99</a> .<span class="kw">await
<a href=#100 id=100 data-nosnippet>100</a> </span>.context(error::InvalidateTableCacheSnafu)<span class="question-mark">?</span>;
<a href=#101 id=101 data-nosnippet>101</a>
<a href=#102 id=102 data-nosnippet>102</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="number">0</span>))
<a href=#103 id=103 data-nosnippet>103</a> }
<a href=#104 id=104 data-nosnippet>104</a>
<a href=#105 id=105 data-nosnippet>105</a> <span class="kw">fn </span>create_comment_on_task_from_expr(<span class="kw-2">&amp;</span><span class="self">self</span>, expr: CommentOnExpr) -&gt; <span class="prelude-ty">Result</span>&lt;CommentOnTask&gt; {
<a href=#106 id=106 data-nosnippet>106</a> <span class="kw">let </span>object_type = <span class="kw">match </span>expr.object_type {
<a href=#107 id=107 data-nosnippet>107</a> <span class="number">0 </span>=&gt; CommentObjectType::Table,
<a href=#108 id=108 data-nosnippet>108</a> <span class="number">1 </span>=&gt; CommentObjectType::Column,
<a href=#109 id=109 data-nosnippet>109</a> <span class="number">2 </span>=&gt; CommentObjectType::Flow,
<a href=#110 id=110 data-nosnippet>110</a> <span class="kw">_ </span>=&gt; {
<a href=#111 id=111 data-nosnippet>111</a> <span class="kw">return </span>InvalidSqlSnafu {
<a href=#112 id=112 data-nosnippet>112</a> err_msg: <span class="macro">format!</span>(
<a href=#113 id=113 data-nosnippet>113</a> <span class="string">"Invalid CommentObjectType value: {}. Valid values are: 0 (Table), 1 (Column), 2 (Flow)"</span>,
<a href=#114 id=114 data-nosnippet>114</a> expr.object_type
<a href=#115 id=115 data-nosnippet>115</a> ),
<a href=#116 id=116 data-nosnippet>116</a> }
<a href=#117 id=117 data-nosnippet>117</a> .fail();
<a href=#118 id=118 data-nosnippet>118</a> }
<a href=#119 id=119 data-nosnippet>119</a> };
<a href=#120 id=120 data-nosnippet>120</a>
<a href=#121 id=121 data-nosnippet>121</a> <span class="prelude-val">Ok</span>(CommentOnTask {
<a href=#122 id=122 data-nosnippet>122</a> catalog_name: expr.catalog_name,
<a href=#123 id=123 data-nosnippet>123</a> schema_name: expr.schema_name,
<a href=#124 id=124 data-nosnippet>124</a> object_type,
<a href=#125 id=125 data-nosnippet>125</a> object_name: expr.object_name,
<a href=#126 id=126 data-nosnippet>126</a> column_name: <span class="kw">if </span>expr.column_name.is_empty() {
<a href=#127 id=127 data-nosnippet>127</a> <span class="prelude-val">None
<a href=#128 id=128 data-nosnippet>128</a> </span>} <span class="kw">else </span>{
<a href=#129 id=129 data-nosnippet>129</a> <span class="prelude-val">Some</span>(expr.column_name)
<a href=#130 id=130 data-nosnippet>130</a> },
<a href=#131 id=131 data-nosnippet>131</a> object_id: <span class="prelude-val">None</span>,
<a href=#132 id=132 data-nosnippet>132</a> comment: <span class="kw">if </span>expr.comment.is_empty() {
<a href=#133 id=133 data-nosnippet>133</a> <span class="prelude-val">None
<a href=#134 id=134 data-nosnippet>134</a> </span>} <span class="kw">else </span>{
<a href=#135 id=135 data-nosnippet>135</a> <span class="prelude-val">Some</span>(expr.comment)
<a href=#136 id=136 data-nosnippet>136</a> },
<a href=#137 id=137 data-nosnippet>137</a> })
<a href=#138 id=138 data-nosnippet>138</a> }
<a href=#139 id=139 data-nosnippet>139</a>
<a href=#140 id=140 data-nosnippet>140</a> <span class="kw">fn </span>create_comment_on_task_from_stmt(
<a href=#141 id=141 data-nosnippet>141</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#142 id=142 data-nosnippet>142</a> stmt: Comment,
<a href=#143 id=143 data-nosnippet>143</a> query_ctx: <span class="kw-2">&amp;</span>QueryContextRef,
<a href=#144 id=144 data-nosnippet>144</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;CommentOnTask&gt; {
<a href=#145 id=145 data-nosnippet>145</a> <span class="kw">match </span>stmt.object {
<a href=#146 id=146 data-nosnippet>146</a> CommentObject::Table(table) =&gt; {
<a href=#147 id=147 data-nosnippet>147</a> <span class="kw">let </span>(catalog_name, schema_name, table_name) =
<a href=#148 id=148 data-nosnippet>148</a> table_idents_to_full_name(<span class="kw-2">&amp;</span>table, query_ctx)
<a href=#149 id=149 data-nosnippet>149</a> .map_err(BoxedError::new)
<a href=#150 id=150 data-nosnippet>150</a> .context(ExternalSnafu)<span class="question-mark">?</span>;
<a href=#151 id=151 data-nosnippet>151</a>
<a href=#152 id=152 data-nosnippet>152</a> <span class="prelude-val">Ok</span>(CommentOnTask {
<a href=#153 id=153 data-nosnippet>153</a> catalog_name,
<a href=#154 id=154 data-nosnippet>154</a> schema_name,
<a href=#155 id=155 data-nosnippet>155</a> object_type: CommentObjectType::Table,
<a href=#156 id=156 data-nosnippet>156</a> object_name: table_name,
<a href=#157 id=157 data-nosnippet>157</a> column_name: <span class="prelude-val">None</span>,
<a href=#158 id=158 data-nosnippet>158</a> object_id: <span class="prelude-val">None</span>,
<a href=#159 id=159 data-nosnippet>159</a> comment: stmt.comment,
<a href=#160 id=160 data-nosnippet>160</a> })
<a href=#161 id=161 data-nosnippet>161</a> }
<a href=#162 id=162 data-nosnippet>162</a> CommentObject::Column { table, column } =&gt; {
<a href=#163 id=163 data-nosnippet>163</a> <span class="kw">let </span>(catalog_name, schema_name, table_name) =
<a href=#164 id=164 data-nosnippet>164</a> table_idents_to_full_name(<span class="kw-2">&amp;</span>table, query_ctx)
<a href=#165 id=165 data-nosnippet>165</a> .map_err(BoxedError::new)
<a href=#166 id=166 data-nosnippet>166</a> .context(ExternalSnafu)<span class="question-mark">?</span>;
<a href=#167 id=167 data-nosnippet>167</a>
<a href=#168 id=168 data-nosnippet>168</a> <span class="prelude-val">Ok</span>(CommentOnTask {
<a href=#169 id=169 data-nosnippet>169</a> catalog_name,
<a href=#170 id=170 data-nosnippet>170</a> schema_name,
<a href=#171 id=171 data-nosnippet>171</a> object_type: CommentObjectType::Column,
<a href=#172 id=172 data-nosnippet>172</a> object_name: table_name,
<a href=#173 id=173 data-nosnippet>173</a> column_name: <span class="prelude-val">Some</span>(column.value),
<a href=#174 id=174 data-nosnippet>174</a> object_id: <span class="prelude-val">None</span>,
<a href=#175 id=175 data-nosnippet>175</a> comment: stmt.comment,
<a href=#176 id=176 data-nosnippet>176</a> })
<a href=#177 id=177 data-nosnippet>177</a> }
<a href=#178 id=178 data-nosnippet>178</a> CommentObject::Flow(flow_name) =&gt; {
<a href=#179 id=179 data-nosnippet>179</a> <span class="kw">let </span>(catalog_name, flow_name_str) = <span class="kw">match </span><span class="kw-2">&amp;</span>flow_name.<span class="number">0</span>[..] {
<a href=#180 id=180 data-nosnippet>180</a> [flow] =&gt; (
<a href=#181 id=181 data-nosnippet>181</a> query_ctx.current_catalog().to_string(),
<a href=#182 id=182 data-nosnippet>182</a> flow.to_string_unquoted(),
<a href=#183 id=183 data-nosnippet>183</a> ),
<a href=#184 id=184 data-nosnippet>184</a> [catalog, flow] =&gt; (catalog.to_string_unquoted(), flow.to_string_unquoted()),
<a href=#185 id=185 data-nosnippet>185</a> <span class="kw">_ </span>=&gt; {
<a href=#186 id=186 data-nosnippet>186</a> <span class="kw">return </span>InvalidSqlSnafu {
<a href=#187 id=187 data-nosnippet>187</a> err_msg: <span class="macro">format!</span>(
<a href=#188 id=188 data-nosnippet>188</a> <span class="string">"expect flow name to be &lt;catalog&gt;.&lt;flow_name&gt; or &lt;flow_name&gt;, actual: {flow_name}"
<a href=#189 id=189 data-nosnippet>189</a> </span>),
<a href=#190 id=190 data-nosnippet>190</a> }
<a href=#191 id=191 data-nosnippet>191</a> .fail();
<a href=#192 id=192 data-nosnippet>192</a> }
<a href=#193 id=193 data-nosnippet>193</a> };
<a href=#194 id=194 data-nosnippet>194</a>
<a href=#195 id=195 data-nosnippet>195</a> <span class="prelude-val">Ok</span>(CommentOnTask {
<a href=#196 id=196 data-nosnippet>196</a> catalog_name,
<a href=#197 id=197 data-nosnippet>197</a> schema_name: String::new(), <span class="comment">// Flow doesn't use schema
<a href=#198 id=198 data-nosnippet>198</a> </span>object_type: CommentObjectType::Flow,
<a href=#199 id=199 data-nosnippet>199</a> object_name: flow_name_str,
<a href=#200 id=200 data-nosnippet>200</a> column_name: <span class="prelude-val">None</span>,
<a href=#201 id=201 data-nosnippet>201</a> object_id: <span class="prelude-val">None</span>,
<a href=#202 id=202 data-nosnippet>202</a> comment: stmt.comment,
<a href=#203 id=203 data-nosnippet>203</a> })
<a href=#204 id=204 data-nosnippet>204</a> }
<a href=#205 id=205 data-nosnippet>205</a> }
<a href=#206 id=206 data-nosnippet>206</a> }
<a href=#207 id=207 data-nosnippet>207</a>}
</code></pre></div></section></main></body></html>
@@ -0,0 +1,289 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/copy_database.rs`."><title>copy_database.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>copy_database.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::str::FromStr;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>std::sync::Arc;
<a href=#17 id=17 data-nosnippet>17</a>
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>client::{Output, OutputData, OutputMeta};
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>common_catalog::format_full_table_name;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>common_datasource::file_format::Format;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>common_datasource::lister::{Lister, Source};
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>common_datasource::object_store::{LocalFileAccess, build_backend, build_backend_for_write};
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>common_telemetry::{debug, error, info, tracing};
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>futures::future::try_join_all;
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>object_store::Entry;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>regex::Regex;
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>snafu::ResultExt;
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span>table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest};
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span>tokio::sync::Semaphore;
<a href=#31 id=31 data-nosnippet>31</a>
<a href=#32 id=32 data-nosnippet>32</a><span class="kw">use </span><span class="kw">crate</span>::error;
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#34 id=34 data-nosnippet>34</a><span class="kw">use </span><span class="kw">crate</span>::statement::database_copy::{
<a href=#35 id=35 data-nosnippet>35</a> DatabaseExportFile, database_import_source, parse_parallelism_from_option_map,
<a href=#36 id=36 data-nosnippet>36</a> validate_database_directory, validate_database_export_layout,
<a href=#37 id=37 data-nosnippet>37</a>};
<a href=#38 id=38 data-nosnippet>38</a>
<a href=#39 id=39 data-nosnippet>39</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">const </span>COPY_DATABASE_TIME_START_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"start_time"</span>;
<a href=#40 id=40 data-nosnippet>40</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">const </span>COPY_DATABASE_TIME_END_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"end_time"</span>;
<a href=#41 id=41 data-nosnippet>41</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">const </span>CONTINUE_ON_ERROR_KEY: <span class="kw-2">&amp;</span>str = <span class="string">"continue_on_error"</span>;
<a href=#42 id=42 data-nosnippet>42</a>
<a href=#43 id=43 data-nosnippet>43</a><span class="kw">impl </span>StatementExecutor {
<a href=#44 id=44 data-nosnippet>44</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#45 id=45 data-nosnippet>45</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>copy_database_to(
<a href=#46 id=46 data-nosnippet>46</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#47 id=47 data-nosnippet>47</a> req: CopyDatabaseRequest,
<a href=#48 id=48 data-nosnippet>48</a> ctx: QueryContextRef,
<a href=#49 id=49 data-nosnippet>49</a> ) -&gt; error::Result&lt;Output&gt; {
<a href=#50 id=50 data-nosnippet>50</a> validate_database_export_layout(<span class="kw-2">&amp;</span>req.with)<span class="question-mark">?</span>;
<a href=#51 id=51 data-nosnippet>51</a> validate_database_directory(<span class="kw-2">&amp;</span>req.location)<span class="question-mark">?</span>;
<a href=#52 id=52 data-nosnippet>52</a> build_backend_for_write(<span class="kw-2">&amp;</span>req.location, <span class="kw-2">&amp;</span>req.connection, <span class="kw-2">&amp;</span><span class="self">self</span>.local_file_access)
<a href=#53 id=53 data-nosnippet>53</a> .<span class="kw">await
<a href=#54 id=54 data-nosnippet>54</a> </span>.context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#55 id=55 data-nosnippet>55</a>
<a href=#56 id=56 data-nosnippet>56</a> <span class="kw">let </span>parallelism = parse_parallelism_from_option_map(<span class="kw-2">&amp;</span>req.with);
<a href=#57 id=57 data-nosnippet>57</a> <span class="macro">info!</span>(
<a href=#58 id=58 data-nosnippet>58</a> <span class="string">"Copy database {}.{} to dir: {}, time: {:?}, parallelism: {}"</span>,
<a href=#59 id=59 data-nosnippet>59</a> req.catalog_name, req.schema_name, req.location, req.time_range, parallelism
<a href=#60 id=60 data-nosnippet>60</a> );
<a href=#61 id=61 data-nosnippet>61</a> <span class="kw">let </span>tables = <span class="self">self
<a href=#62 id=62 data-nosnippet>62</a> </span>.capture_database_export_tables(<span class="kw-2">&amp;</span>req, <span class="prelude-val">None</span>, <span class="kw-2">&amp;</span>ctx)
<a href=#63 id=63 data-nosnippet>63</a> .<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#64 id=64 data-nosnippet>64</a> <span class="kw">let </span>num_tables = tables.len();
<a href=#65 id=65 data-nosnippet>65</a>
<a href=#66 id=66 data-nosnippet>66</a> <span class="kw">let </span>suffix = Format::try_from(<span class="kw-2">&amp;</span>req.with)
<a href=#67 id=67 data-nosnippet>67</a> .context(error::ParseFileFormatSnafu)<span class="question-mark">?
<a href=#68 id=68 data-nosnippet>68</a> </span>.suffix();
<a href=#69 id=69 data-nosnippet>69</a>
<a href=#70 id=70 data-nosnippet>70</a> <span class="kw">let </span><span class="kw-2">mut </span>tasks = Vec::with_capacity(num_tables);
<a href=#71 id=71 data-nosnippet>71</a> <span class="kw">let </span>semaphore = Arc::new(Semaphore::new(parallelism));
<a href=#72 id=72 data-nosnippet>72</a>
<a href=#73 id=73 data-nosnippet>73</a> <span class="kw">for </span>(i, table) <span class="kw">in </span>tables.into_iter().enumerate() {
<a href=#74 id=74 data-nosnippet>74</a> <span class="kw">let </span>table_name = table.table_info().name.clone();
<a href=#75 id=75 data-nosnippet>75</a> <span class="kw">let </span>semaphore_moved = semaphore.clone();
<a href=#76 id=76 data-nosnippet>76</a> <span class="kw">let </span>table_file = DatabaseExportFile::new(<span class="kw-2">&amp;</span>req.location, <span class="kw-2">&amp;</span>table_name, suffix)<span class="question-mark">?</span>.location;
<a href=#77 id=77 data-nosnippet>77</a> <span class="kw">let </span>table_no = i + <span class="number">1</span>;
<a href=#78 id=78 data-nosnippet>78</a> <span class="kw">let </span>moved_ctx = ctx.clone();
<a href=#79 id=79 data-nosnippet>79</a> <span class="kw">let </span>full_table_name =
<a href=#80 id=80 data-nosnippet>80</a> format_full_table_name(<span class="kw-2">&amp;</span>req.catalog_name, <span class="kw-2">&amp;</span>req.schema_name, <span class="kw-2">&amp;</span>table_name);
<a href=#81 id=81 data-nosnippet>81</a> <span class="kw">let </span>copy_table_req = CopyTableRequest {
<a href=#82 id=82 data-nosnippet>82</a> catalog_name: req.catalog_name.clone(),
<a href=#83 id=83 data-nosnippet>83</a> schema_name: req.schema_name.clone(),
<a href=#84 id=84 data-nosnippet>84</a> table_name,
<a href=#85 id=85 data-nosnippet>85</a> location: table_file.clone(),
<a href=#86 id=86 data-nosnippet>86</a> with: req.with.clone(),
<a href=#87 id=87 data-nosnippet>87</a> connection: req.connection.clone(),
<a href=#88 id=88 data-nosnippet>88</a> pattern: <span class="prelude-val">None</span>,
<a href=#89 id=89 data-nosnippet>89</a> direction: CopyDirection::Export,
<a href=#90 id=90 data-nosnippet>90</a> timestamp_range: req.time_range,
<a href=#91 id=91 data-nosnippet>91</a> limit: <span class="prelude-val">None</span>,
<a href=#92 id=92 data-nosnippet>92</a> };
<a href=#93 id=93 data-nosnippet>93</a>
<a href=#94 id=94 data-nosnippet>94</a> tasks.push(<span class="kw">async move </span>{
<a href=#95 id=95 data-nosnippet>95</a> <span class="kw">let </span>_permit = semaphore_moved.acquire().<span class="kw">await</span>.unwrap();
<a href=#96 id=96 data-nosnippet>96</a> <span class="macro">info!</span>(
<a href=#97 id=97 data-nosnippet>97</a> <span class="string">"Copy table({}/{}): {} to {}"</span>,
<a href=#98 id=98 data-nosnippet>98</a> table_no, num_tables, full_table_name, table_file
<a href=#99 id=99 data-nosnippet>99</a> );
<a href=#100 id=100 data-nosnippet>100</a> <span class="self">self</span>.copy_captured_table_to(table, copy_table_req, moved_ctx)
<a href=#101 id=101 data-nosnippet>101</a> .<span class="kw">await
<a href=#102 id=102 data-nosnippet>102</a> </span>});
<a href=#103 id=103 data-nosnippet>103</a> }
<a href=#104 id=104 data-nosnippet>104</a>
<a href=#105 id=105 data-nosnippet>105</a> <span class="kw">let </span>results = try_join_all(tasks).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#106 id=106 data-nosnippet>106</a> <span class="kw">let </span>exported_rows = results.into_iter().sum();
<a href=#107 id=107 data-nosnippet>107</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(exported_rows))
<a href=#108 id=108 data-nosnippet>108</a> }
<a href=#109 id=109 data-nosnippet>109</a>
<a href=#110 id=110 data-nosnippet>110</a> <span class="doccomment">/// Imports data to database from a given location and returns total rows imported.
<a href=#111 id=111 data-nosnippet>111</a> </span><span class="attr">#[tracing::instrument(skip_all)]
<a href=#112 id=112 data-nosnippet>112</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>copy_database_from(
<a href=#113 id=113 data-nosnippet>113</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#114 id=114 data-nosnippet>114</a> req: CopyDatabaseRequest,
<a href=#115 id=115 data-nosnippet>115</a> ctx: QueryContextRef,
<a href=#116 id=116 data-nosnippet>116</a> ) -&gt; error::Result&lt;Output&gt; {
<a href=#117 id=117 data-nosnippet>117</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(layout) = req.with.get(<span class="string">"metric_data_layout"</span>) {
<a href=#118 id=118 data-nosnippet>118</a> <span class="kw">return </span>error::InvalidCopyParameterSnafu {
<a href=#119 id=119 data-nosnippet>119</a> key: <span class="string">"metric_data_layout"</span>,
<a href=#120 id=120 data-nosnippet>120</a> value: layout,
<a href=#121 id=121 data-nosnippet>121</a> }
<a href=#122 id=122 data-nosnippet>122</a> .fail();
<a href=#123 id=123 data-nosnippet>123</a> }
<a href=#124 id=124 data-nosnippet>124</a> validate_database_directory(<span class="kw-2">&amp;</span>req.location)<span class="question-mark">?</span>;
<a href=#125 id=125 data-nosnippet>125</a>
<a href=#126 id=126 data-nosnippet>126</a> <span class="kw">let </span>parallelism = parse_parallelism_from_option_map(<span class="kw-2">&amp;</span>req.with);
<a href=#127 id=127 data-nosnippet>127</a> <span class="macro">info!</span>(
<a href=#128 id=128 data-nosnippet>128</a> <span class="string">"Copy database {}.{} from dir: {}, time: {:?}, parallelism: {}"</span>,
<a href=#129 id=129 data-nosnippet>129</a> req.catalog_name, req.schema_name, req.location, req.time_range, parallelism
<a href=#130 id=130 data-nosnippet>130</a> );
<a href=#131 id=131 data-nosnippet>131</a> <span class="kw">let </span>suffix = Format::try_from(<span class="kw-2">&amp;</span>req.with)
<a href=#132 id=132 data-nosnippet>132</a> .context(error::ParseFileFormatSnafu)<span class="question-mark">?
<a href=#133 id=133 data-nosnippet>133</a> </span>.suffix();
<a href=#134 id=134 data-nosnippet>134</a>
<a href=#135 id=135 data-nosnippet>135</a> <span class="kw">let </span>entries = list_files_to_copy(<span class="kw-2">&amp;</span>req, suffix, <span class="kw-2">&amp;</span><span class="self">self</span>.local_file_access).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#136 id=136 data-nosnippet>136</a>
<a href=#137 id=137 data-nosnippet>137</a> <span class="kw">let </span>continue_on_error = req
<a href=#138 id=138 data-nosnippet>138</a> .with
<a href=#139 id=139 data-nosnippet>139</a> .get(CONTINUE_ON_ERROR_KEY)
<a href=#140 id=140 data-nosnippet>140</a> .and_then(|v| bool::from_str(v).ok())
<a href=#141 id=141 data-nosnippet>141</a> .unwrap_or(<span class="bool-val">false</span>);
<a href=#142 id=142 data-nosnippet>142</a>
<a href=#143 id=143 data-nosnippet>143</a> <span class="kw">let </span><span class="kw-2">mut </span>tasks = Vec::with_capacity(entries.len());
<a href=#144 id=144 data-nosnippet>144</a> <span class="kw">let </span>semaphore = Arc::new(Semaphore::new(parallelism));
<a href=#145 id=145 data-nosnippet>145</a>
<a href=#146 id=146 data-nosnippet>146</a> <span class="kw">for </span>e <span class="kw">in </span>entries {
<a href=#147 id=147 data-nosnippet>147</a> <span class="kw">let </span>(table_name, location) = <span class="kw">match </span>database_import_source(<span class="kw-2">&amp;</span>req.location, e.path()) {
<a href=#148 id=148 data-nosnippet>148</a> <span class="prelude-val">Ok</span>(source) =&gt; source,
<a href=#149 id=149 data-nosnippet>149</a> <span class="prelude-val">Err</span>(err) =&gt; {
<a href=#150 id=150 data-nosnippet>150</a> <span class="kw">if </span>continue_on_error {
<a href=#151 id=151 data-nosnippet>151</a> <span class="macro">error!</span>(err; <span class="string">"Failed to import table from file: {:?}"</span>, e);
<a href=#152 id=152 data-nosnippet>152</a> <span class="kw">continue</span>;
<a href=#153 id=153 data-nosnippet>153</a> } <span class="kw">else </span>{
<a href=#154 id=154 data-nosnippet>154</a> <span class="kw">return </span><span class="prelude-val">Err</span>(err);
<a href=#155 id=155 data-nosnippet>155</a> }
<a href=#156 id=156 data-nosnippet>156</a> }
<a href=#157 id=157 data-nosnippet>157</a> };
<a href=#158 id=158 data-nosnippet>158</a>
<a href=#159 id=159 data-nosnippet>159</a> <span class="kw">let </span>req = CopyTableRequest {
<a href=#160 id=160 data-nosnippet>160</a> catalog_name: req.catalog_name.clone(),
<a href=#161 id=161 data-nosnippet>161</a> schema_name: req.schema_name.clone(),
<a href=#162 id=162 data-nosnippet>162</a> table_name: table_name.clone(),
<a href=#163 id=163 data-nosnippet>163</a> location,
<a href=#164 id=164 data-nosnippet>164</a> with: req.with.clone(),
<a href=#165 id=165 data-nosnippet>165</a> connection: req.connection.clone(),
<a href=#166 id=166 data-nosnippet>166</a> pattern: <span class="prelude-val">None</span>,
<a href=#167 id=167 data-nosnippet>167</a> direction: CopyDirection::Import,
<a href=#168 id=168 data-nosnippet>168</a> timestamp_range: <span class="prelude-val">None</span>,
<a href=#169 id=169 data-nosnippet>169</a> limit: <span class="prelude-val">None</span>,
<a href=#170 id=170 data-nosnippet>170</a> };
<a href=#171 id=171 data-nosnippet>171</a> <span class="kw">let </span>moved_ctx = ctx.clone();
<a href=#172 id=172 data-nosnippet>172</a> <span class="kw">let </span>moved_table_name = table_name.clone();
<a href=#173 id=173 data-nosnippet>173</a> <span class="kw">let </span>moved_semaphore = semaphore.clone();
<a href=#174 id=174 data-nosnippet>174</a> tasks.push(<span class="kw">async move </span>{
<a href=#175 id=175 data-nosnippet>175</a> <span class="kw">let </span>_permit = moved_semaphore.acquire().<span class="kw">await</span>.unwrap();
<a href=#176 id=176 data-nosnippet>176</a> <span class="macro">debug!</span>(<span class="string">"Copy table, arg: {:?}"</span>, req);
<a href=#177 id=177 data-nosnippet>177</a> <span class="kw">match </span><span class="self">self</span>.copy_table_from(req, moved_ctx).<span class="kw">await </span>{
<a href=#178 id=178 data-nosnippet>178</a> <span class="prelude-val">Ok</span>(o) =&gt; {
<a href=#179 id=179 data-nosnippet>179</a> <span class="kw">let </span>(rows, cost) = o.extract_rows_and_cost();
<a href=#180 id=180 data-nosnippet>180</a> <span class="prelude-val">Ok</span>((rows, cost))
<a href=#181 id=181 data-nosnippet>181</a> }
<a href=#182 id=182 data-nosnippet>182</a> <span class="prelude-val">Err</span>(err) =&gt; {
<a href=#183 id=183 data-nosnippet>183</a> <span class="kw">if </span>continue_on_error {
<a href=#184 id=184 data-nosnippet>184</a> <span class="macro">error!</span>(err; <span class="string">"Failed to import file to table: {}"</span>, moved_table_name);
<a href=#185 id=185 data-nosnippet>185</a> <span class="prelude-val">Ok</span>((<span class="number">0</span>, <span class="number">0</span>))
<a href=#186 id=186 data-nosnippet>186</a> } <span class="kw">else </span>{
<a href=#187 id=187 data-nosnippet>187</a> <span class="prelude-val">Err</span>(err)
<a href=#188 id=188 data-nosnippet>188</a> }
<a href=#189 id=189 data-nosnippet>189</a> }
<a href=#190 id=190 data-nosnippet>190</a> }
<a href=#191 id=191 data-nosnippet>191</a> });
<a href=#192 id=192 data-nosnippet>192</a> }
<a href=#193 id=193 data-nosnippet>193</a>
<a href=#194 id=194 data-nosnippet>194</a> <span class="kw">let </span>results = try_join_all(tasks).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#195 id=195 data-nosnippet>195</a> <span class="kw">let </span>(rows_inserted, insert_cost) = results
<a href=#196 id=196 data-nosnippet>196</a> .into_iter()
<a href=#197 id=197 data-nosnippet>197</a> .fold((<span class="number">0</span>, <span class="number">0</span>), |(acc_rows, acc_cost), (rows, cost)| {
<a href=#198 id=198 data-nosnippet>198</a> (acc_rows + rows, acc_cost + cost)
<a href=#199 id=199 data-nosnippet>199</a> });
<a href=#200 id=200 data-nosnippet>200</a>
<a href=#201 id=201 data-nosnippet>201</a> <span class="prelude-val">Ok</span>(Output::new(
<a href=#202 id=202 data-nosnippet>202</a> OutputData::AffectedRows(rows_inserted),
<a href=#203 id=203 data-nosnippet>203</a> OutputMeta::new_with_cost(insert_cost),
<a href=#204 id=204 data-nosnippet>204</a> ))
<a href=#205 id=205 data-nosnippet>205</a> }
<a href=#206 id=206 data-nosnippet>206</a>}
<a href=#207 id=207 data-nosnippet>207</a>
<a href=#208 id=208 data-nosnippet>208</a><span class="doccomment">/// Lists all files with expected suffix that can be imported to database.
<a href=#209 id=209 data-nosnippet>209</a></span><span class="kw">async fn </span>list_files_to_copy(
<a href=#210 id=210 data-nosnippet>210</a> req: <span class="kw-2">&amp;</span>CopyDatabaseRequest,
<a href=#211 id=211 data-nosnippet>211</a> suffix: <span class="kw-2">&amp;</span>str,
<a href=#212 id=212 data-nosnippet>212</a> local_file_access: <span class="kw-2">&amp;</span>LocalFileAccess,
<a href=#213 id=213 data-nosnippet>213</a>) -&gt; error::Result&lt;Vec&lt;Entry&gt;&gt; {
<a href=#214 id=214 data-nosnippet>214</a> <span class="kw">let </span>object_store = build_backend(<span class="kw-2">&amp;</span>req.location, <span class="kw-2">&amp;</span>req.connection, local_file_access)
<a href=#215 id=215 data-nosnippet>215</a> .<span class="kw">await
<a href=#216 id=216 data-nosnippet>216</a> </span>.context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#217 id=217 data-nosnippet>217</a>
<a href=#218 id=218 data-nosnippet>218</a> <span class="kw">let </span>pattern = Regex::try_from(<span class="macro">format!</span>(<span class="string">".*{}"</span>, suffix)).context(error::BuildRegexSnafu)<span class="question-mark">?</span>;
<a href=#219 id=219 data-nosnippet>219</a> <span class="kw">let </span>lister = Lister::new(
<a href=#220 id=220 data-nosnippet>220</a> object_store.clone(),
<a href=#221 id=221 data-nosnippet>221</a> Source::Dir,
<a href=#222 id=222 data-nosnippet>222</a> <span class="string">"/"</span>.to_string(),
<a href=#223 id=223 data-nosnippet>223</a> <span class="prelude-val">Some</span>(pattern),
<a href=#224 id=224 data-nosnippet>224</a> );
<a href=#225 id=225 data-nosnippet>225</a> lister.list().<span class="kw">await</span>.context(error::ListObjectsSnafu)
<a href=#226 id=226 data-nosnippet>226</a>}
<a href=#227 id=227 data-nosnippet>227</a>
<a href=#228 id=228 data-nosnippet>228</a><span class="attr">#[cfg(test)]
<a href=#229 id=229 data-nosnippet>229</a></span><span class="kw">mod </span>tests {
<a href=#230 id=230 data-nosnippet>230</a> <span class="kw">use </span>std::collections::HashSet;
<a href=#231 id=231 data-nosnippet>231</a>
<a href=#232 id=232 data-nosnippet>232</a> <span class="kw">use </span>common_datasource::object_store::LocalFileAccess;
<a href=#233 id=233 data-nosnippet>233</a> <span class="kw">use </span>object_store::ObjectStore;
<a href=#234 id=234 data-nosnippet>234</a> <span class="kw">use </span>object_store::services::Fs;
<a href=#235 id=235 data-nosnippet>235</a> <span class="kw">use </span>object_store::util::normalize_dir;
<a href=#236 id=236 data-nosnippet>236</a> <span class="attr">#[cfg(not(windows))]
<a href=#237 id=237 data-nosnippet>237</a> </span><span class="kw">use </span>path_slash::PathExt;
<a href=#238 id=238 data-nosnippet>238</a> <span class="kw">use </span>table::requests::CopyDatabaseRequest;
<a href=#239 id=239 data-nosnippet>239</a>
<a href=#240 id=240 data-nosnippet>240</a> <span class="kw">use </span><span class="kw">crate</span>::statement::copy_database::list_files_to_copy;
<a href=#241 id=241 data-nosnippet>241</a> <span class="kw">use </span><span class="kw">crate</span>::statement::database_copy::database_import_source;
<a href=#242 id=242 data-nosnippet>242</a>
<a href=#243 id=243 data-nosnippet>243</a> <span class="attr">#[tokio::test]
<a href=#244 id=244 data-nosnippet>244</a> </span><span class="kw">async fn </span>test_list_files_and_parse_table_name() {
<a href=#245 id=245 data-nosnippet>245</a> <span class="kw">let </span>dir = common_test_util::temp_dir::create_temp_dir(<span class="string">"test_list_files_to_copy"</span>);
<a href=#246 id=246 data-nosnippet>246</a> <span class="kw">let </span>store_dir = normalize_dir(dir.path().to_str().unwrap());
<a href=#247 id=247 data-nosnippet>247</a> <span class="kw">let </span>builder = Fs::default().root(<span class="kw-2">&amp;</span>store_dir);
<a href=#248 id=248 data-nosnippet>248</a> <span class="kw">let </span>object_store = ObjectStore::new(builder).unwrap();
<a href=#249 id=249 data-nosnippet>249</a> object_store.write(<span class="string">"a.parquet"</span>, <span class="string">""</span>).<span class="kw">await</span>.unwrap();
<a href=#250 id=250 data-nosnippet>250</a> object_store.write(<span class="string">"b.parquet"</span>, <span class="string">""</span>).<span class="kw">await</span>.unwrap();
<a href=#251 id=251 data-nosnippet>251</a> object_store.write(<span class="string">"c.csv"</span>, <span class="string">""</span>).<span class="kw">await</span>.unwrap();
<a href=#252 id=252 data-nosnippet>252</a> object_store.write(<span class="string">"d"</span>, <span class="string">""</span>).<span class="kw">await</span>.unwrap();
<a href=#253 id=253 data-nosnippet>253</a> object_store.write(<span class="string">"e.f.parquet"</span>, <span class="string">""</span>).<span class="kw">await</span>.unwrap();
<a href=#254 id=254 data-nosnippet>254</a>
<a href=#255 id=255 data-nosnippet>255</a> <span class="attr">#[cfg(not(windows))]
<a href=#256 id=256 data-nosnippet>256</a> </span><span class="kw">let </span>location = normalize_dir(<span class="kw-2">&amp;</span>dir.path().to_slash().unwrap());
<a href=#257 id=257 data-nosnippet>257</a> <span class="attr">#[cfg(windows)]
<a href=#258 id=258 data-nosnippet>258</a> </span><span class="kw">let </span>location = <span class="macro">format!</span>(<span class="string">"{}\\"</span>, dir.path().display());
<a href=#259 id=259 data-nosnippet>259</a> <span class="kw">let </span>request = CopyDatabaseRequest {
<a href=#260 id=260 data-nosnippet>260</a> catalog_name: <span class="string">"catalog_0"</span>.to_string(),
<a href=#261 id=261 data-nosnippet>261</a> schema_name: <span class="string">"schema_0"</span>.to_string(),
<a href=#262 id=262 data-nosnippet>262</a> location,
<a href=#263 id=263 data-nosnippet>263</a> with: [(<span class="string">"FORMAT"</span>.to_string(), <span class="string">"parquet"</span>.to_string())]
<a href=#264 id=264 data-nosnippet>264</a> .into_iter()
<a href=#265 id=265 data-nosnippet>265</a> .collect(),
<a href=#266 id=266 data-nosnippet>266</a> connection: Default::default(),
<a href=#267 id=267 data-nosnippet>267</a> time_range: <span class="prelude-val">None</span>,
<a href=#268 id=268 data-nosnippet>268</a> };
<a href=#269 id=269 data-nosnippet>269</a> <span class="kw">let </span>local_file_access = LocalFileAccess::sandboxed(dir.path()).unwrap();
<a href=#270 id=270 data-nosnippet>270</a> <span class="kw">let </span>listed = list_files_to_copy(<span class="kw-2">&amp;</span>request, <span class="string">".parquet"</span>, <span class="kw-2">&amp;</span>local_file_access)
<a href=#271 id=271 data-nosnippet>271</a> .<span class="kw">await
<a href=#272 id=272 data-nosnippet>272</a> </span>.unwrap()
<a href=#273 id=273 data-nosnippet>273</a> .into_iter()
<a href=#274 id=274 data-nosnippet>274</a> .map(|e| {
<a href=#275 id=275 data-nosnippet>275</a> database_import_source(<span class="kw-2">&amp;</span>request.location, e.path())
<a href=#276 id=276 data-nosnippet>276</a> .unwrap()
<a href=#277 id=277 data-nosnippet>277</a> .<span class="number">0
<a href=#278 id=278 data-nosnippet>278</a> </span>})
<a href=#279 id=279 data-nosnippet>279</a> .collect::&lt;HashSet&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#280 id=280 data-nosnippet>280</a>
<a href=#281 id=281 data-nosnippet>281</a> <span class="macro">assert_eq!</span>(
<a href=#282 id=282 data-nosnippet>282</a> [<span class="string">"a"</span>.to_string(), <span class="string">"b"</span>.to_string(), <span class="string">"e.f"</span>.to_string()]
<a href=#283 id=283 data-nosnippet>283</a> .into_iter()
<a href=#284 id=284 data-nosnippet>284</a> .collect::&lt;HashSet&lt;<span class="kw">_</span>&gt;&gt;(),
<a href=#285 id=285 data-nosnippet>285</a> listed
<a href=#286 id=286 data-nosnippet>286</a> );
<a href=#287 id=287 data-nosnippet>287</a> }
<a href=#288 id=288 data-nosnippet>288</a>}
</code></pre></div></section></main></body></html>
@@ -0,0 +1,44 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/copy_query_to.rs`."><title>copy_query_to.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>copy_query_to.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-2"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>common_datasource::file_format::Format;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>common_query::Output;
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_telemetry::{debug, tracing};
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>snafu::ResultExt;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>table::requests::CopyQueryToRequest;
<a href=#20 id=20 data-nosnippet>20</a>
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span><span class="kw">crate</span>::error;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span><span class="kw">crate</span>::error::Result;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#24 id=24 data-nosnippet>24</a>
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">impl </span>StatementExecutor {
<a href=#26 id=26 data-nosnippet>26</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#27 id=27 data-nosnippet>27</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>copy_query_to(
<a href=#28 id=28 data-nosnippet>28</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#29 id=29 data-nosnippet>29</a> req: CopyQueryToRequest,
<a href=#30 id=30 data-nosnippet>30</a> query_output: Output,
<a href=#31 id=31 data-nosnippet>31</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#32 id=32 data-nosnippet>32</a> <span class="kw">let </span>CopyQueryToRequest {
<a href=#33 id=33 data-nosnippet>33</a> with,
<a href=#34 id=34 data-nosnippet>34</a> connection,
<a href=#35 id=35 data-nosnippet>35</a> location,
<a href=#36 id=36 data-nosnippet>36</a> } = <span class="kw-2">&amp;</span>req;
<a href=#37 id=37 data-nosnippet>37</a> <span class="kw">let </span>format = Format::try_from(with).context(error::ParseFileFormatSnafu)<span class="question-mark">?</span>;
<a href=#38 id=38 data-nosnippet>38</a>
<a href=#39 id=39 data-nosnippet>39</a> <span class="macro">debug!</span>(<span class="string">"Copy query to location: {location}"</span>);
<a href=#40 id=40 data-nosnippet>40</a> <span class="self">self</span>.copy_to_file(<span class="kw-2">&amp;</span>format, query_output, location, connection)
<a href=#41 id=41 data-nosnippet>41</a> .<span class="kw">await
<a href=#42 id=42 data-nosnippet>42</a> </span>}
<a href=#43 id=43 data-nosnippet>43</a>}
</code></pre></div></section></main></body></html>
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,672 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/copy_table_to.rs`."><title>copy_table_to.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>copy_table_to.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::collections::HashMap;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>std::sync::Arc;
<a href=#17 id=17 data-nosnippet>17</a>
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>arrow::array::{Array, ArrayRef, AsArray, UInt64Array, make_array};
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>arrow::datatypes::DataType;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>client::OutputData;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>common_base::readable_size::ReadableSize;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>common_datasource::file_format::Format;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>common_datasource::file_format::csv::stream_to_csv;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>common_datasource::file_format::json::stream_to_json;
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>common_datasource::file_format::parquet::stream_to_parquet;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>common_datasource::object_store::build_backend_for_write_with_path;
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>common_datasource::packed_writer::PackedTableWriter;
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>common_datasource::parquet_writer::ParquetFileWriter;
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span>common_query::Output;
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span>common_recordbatch::adapter::DfRecordBatchStreamAdapter;
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span>common_recordbatch::{
<a href=#32 id=32 data-nosnippet>32</a> SendableRecordBatchMapper, SendableRecordBatchStream, map_json_type_to_string,
<a href=#33 id=33 data-nosnippet>33</a> map_json_type_to_string_schema,
<a href=#34 id=34 data-nosnippet>34</a>};
<a href=#35 id=35 data-nosnippet>35</a><span class="kw">use </span>common_telemetry::{debug, tracing};
<a href=#36 id=36 data-nosnippet>36</a><span class="kw">use </span>datafusion::datasource::DefaultTableSource;
<a href=#37 id=37 data-nosnippet>37</a><span class="kw">use </span>datafusion_common::TableReference <span class="kw">as </span>DfTableReference;
<a href=#38 id=38 data-nosnippet>38</a><span class="kw">use </span>datafusion_expr::LogicalPlanBuilder;
<a href=#39 id=39 data-nosnippet>39</a><span class="kw">use </span>futures::StreamExt;
<a href=#40 id=40 data-nosnippet>40</a><span class="kw">use </span>object_store::ObjectStore;
<a href=#41 id=41 data-nosnippet>41</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#42 id=42 data-nosnippet>42</a><span class="kw">use </span>snafu::{OptionExt, ResultExt, ensure};
<a href=#43 id=43 data-nosnippet>43</a><span class="kw">use </span>table::TableRef;
<a href=#44 id=44 data-nosnippet>44</a><span class="kw">use </span>table::requests::CopyTableRequest;
<a href=#45 id=45 data-nosnippet>45</a><span class="kw">use </span>table::table::adapter::DfTableProviderAdapter;
<a href=#46 id=46 data-nosnippet>46</a><span class="kw">use </span>table::table_reference::TableReference;
<a href=#47 id=47 data-nosnippet>47</a><span class="kw">use </span>tokio_util::sync::CancellationToken;
<a href=#48 id=48 data-nosnippet>48</a>
<a href=#49 id=49 data-nosnippet>49</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, BuildDfLogicalPlanSnafu, ExecLogicalPlanSnafu, <span class="prelude-ty">Result</span>};
<a href=#50 id=50 data-nosnippet>50</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#51 id=51 data-nosnippet>51</a><span class="kw">use </span><span class="kw">crate</span>::statement::export_logical_tables::writers::ExportWriteBudget;
<a href=#52 id=52 data-nosnippet>52</a><span class="kw">use </span><span class="kw">crate</span>::statement::export_logical_tables::{
<a href=#53 id=53 data-nosnippet>53</a> expand_export_batch, map_writer_error, rows_within_budget,
<a href=#54 id=54 data-nosnippet>54</a>};
<a href=#55 id=55 data-nosnippet>55</a>
<a href=#56 id=56 data-nosnippet>56</a><span class="comment">// The buffer size should be greater than 5MB (minimum multipart upload size).
<a href=#57 id=57 data-nosnippet>57</a></span><span class="doccomment">/// Buffer size to flush data to object stores.
<a href=#58 id=58 data-nosnippet>58</a></span><span class="kw">const </span>WRITE_BUFFER_THRESHOLD: ReadableSize = ReadableSize::mb(<span class="number">8</span>);
<a href=#59 id=59 data-nosnippet>59</a>
<a href=#60 id=60 data-nosnippet>60</a><span class="doccomment">/// Default number of concurrent write, it only works on object store backend(e.g., S3).
<a href=#61 id=61 data-nosnippet>61</a></span><span class="kw">const </span>WRITE_CONCURRENCY: usize = <span class="number">8</span>;
<a href=#62 id=62 data-nosnippet>62</a>
<a href=#63 id=63 data-nosnippet>63</a><span class="kw">impl </span>StatementExecutor {
<a href=#64 id=64 data-nosnippet>64</a> <span class="kw">async fn </span>stream_to_file(
<a href=#65 id=65 data-nosnippet>65</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#66 id=66 data-nosnippet>66</a> stream: SendableRecordBatchStream,
<a href=#67 id=67 data-nosnippet>67</a> format: <span class="kw-2">&amp;</span>Format,
<a href=#68 id=68 data-nosnippet>68</a> object_store: ObjectStore,
<a href=#69 id=69 data-nosnippet>69</a> path: <span class="kw-2">&amp;</span>str,
<a href=#70 id=70 data-nosnippet>70</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#71 id=71 data-nosnippet>71</a> <span class="kw">let </span>threshold = WRITE_BUFFER_THRESHOLD.as_bytes() <span class="kw">as </span>usize;
<a href=#72 id=72 data-nosnippet>72</a>
<a href=#73 id=73 data-nosnippet>73</a> <span class="kw">let </span>stream = Box::pin(SendableRecordBatchMapper::new(
<a href=#74 id=74 data-nosnippet>74</a> stream,
<a href=#75 id=75 data-nosnippet>75</a> map_json_type_to_string,
<a href=#76 id=76 data-nosnippet>76</a> map_json_type_to_string_schema,
<a href=#77 id=77 data-nosnippet>77</a> ));
<a href=#78 id=78 data-nosnippet>78</a> <span class="kw">match </span>format {
<a href=#79 id=79 data-nosnippet>79</a> Format::Csv(format) =&gt; stream_to_csv(
<a href=#80 id=80 data-nosnippet>80</a> Box::pin(DfRecordBatchStreamAdapter::new(stream)),
<a href=#81 id=81 data-nosnippet>81</a> object_store,
<a href=#82 id=82 data-nosnippet>82</a> path,
<a href=#83 id=83 data-nosnippet>83</a> threshold,
<a href=#84 id=84 data-nosnippet>84</a> WRITE_CONCURRENCY,
<a href=#85 id=85 data-nosnippet>85</a> format,
<a href=#86 id=86 data-nosnippet>86</a> )
<a href=#87 id=87 data-nosnippet>87</a> .<span class="kw">await
<a href=#88 id=88 data-nosnippet>88</a> </span>.context(error::WriteStreamToFileSnafu { path }),
<a href=#89 id=89 data-nosnippet>89</a> Format::Json(format) =&gt; stream_to_json(
<a href=#90 id=90 data-nosnippet>90</a> Box::pin(DfRecordBatchStreamAdapter::new(stream)),
<a href=#91 id=91 data-nosnippet>91</a> object_store,
<a href=#92 id=92 data-nosnippet>92</a> path,
<a href=#93 id=93 data-nosnippet>93</a> threshold,
<a href=#94 id=94 data-nosnippet>94</a> WRITE_CONCURRENCY,
<a href=#95 id=95 data-nosnippet>95</a> format,
<a href=#96 id=96 data-nosnippet>96</a> )
<a href=#97 id=97 data-nosnippet>97</a> .<span class="kw">await
<a href=#98 id=98 data-nosnippet>98</a> </span>.context(error::WriteStreamToFileSnafu { path }),
<a href=#99 id=99 data-nosnippet>99</a> Format::Parquet(<span class="kw">_</span>) =&gt; stream_to_parquet(
<a href=#100 id=100 data-nosnippet>100</a> Box::pin(DfRecordBatchStreamAdapter::new(stream)),
<a href=#101 id=101 data-nosnippet>101</a> object_store,
<a href=#102 id=102 data-nosnippet>102</a> path,
<a href=#103 id=103 data-nosnippet>103</a> WRITE_CONCURRENCY,
<a href=#104 id=104 data-nosnippet>104</a> )
<a href=#105 id=105 data-nosnippet>105</a> .<span class="kw">await
<a href=#106 id=106 data-nosnippet>106</a> </span>.context(error::WriteStreamToFileSnafu { path }),
<a href=#107 id=107 data-nosnippet>107</a> <span class="kw">_ </span>=&gt; error::UnsupportedFormatSnafu {
<a href=#108 id=108 data-nosnippet>108</a> format: format.clone(),
<a href=#109 id=109 data-nosnippet>109</a> }
<a href=#110 id=110 data-nosnippet>110</a> .fail(),
<a href=#111 id=111 data-nosnippet>111</a> }
<a href=#112 id=112 data-nosnippet>112</a> }
<a href=#113 id=113 data-nosnippet>113</a>
<a href=#114 id=114 data-nosnippet>114</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#115 id=115 data-nosnippet>115</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>copy_table_to(
<a href=#116 id=116 data-nosnippet>116</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#117 id=117 data-nosnippet>117</a> req: CopyTableRequest,
<a href=#118 id=118 data-nosnippet>118</a> query_ctx: QueryContextRef,
<a href=#119 id=119 data-nosnippet>119</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#120 id=120 data-nosnippet>120</a> <span class="kw">let </span>table_ref = TableReference::full(<span class="kw-2">&amp;</span>req.catalog_name, <span class="kw-2">&amp;</span>req.schema_name, <span class="kw-2">&amp;</span>req.table_name);
<a href=#121 id=121 data-nosnippet>121</a> <span class="kw">let </span>table = <span class="self">self</span>.get_table(<span class="kw-2">&amp;</span>table_ref).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#122 id=122 data-nosnippet>122</a> <span class="self">self</span>.copy_captured_table_to(table, req, query_ctx).<span class="kw">await
<a href=#123 id=123 data-nosnippet>123</a> </span>}
<a href=#124 id=124 data-nosnippet>124</a>
<a href=#125 id=125 data-nosnippet>125</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>copy_captured_table_to(
<a href=#126 id=126 data-nosnippet>126</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#127 id=127 data-nosnippet>127</a> table: TableRef,
<a href=#128 id=128 data-nosnippet>128</a> req: CopyTableRequest,
<a href=#129 id=129 data-nosnippet>129</a> query_ctx: QueryContextRef,
<a href=#130 id=130 data-nosnippet>130</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#131 id=131 data-nosnippet>131</a> <span class="self">self</span>.copy_captured_table_to_managed(table, req, query_ctx, <span class="prelude-val">None</span>, <span class="prelude-val">None</span>)
<a href=#132 id=132 data-nosnippet>132</a> .<span class="kw">await
<a href=#133 id=133 data-nosnippet>133</a> </span>}
<a href=#134 id=134 data-nosnippet>134</a>
<a href=#135 id=135 data-nosnippet>135</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>copy_captured_table_to_managed(
<a href=#136 id=136 data-nosnippet>136</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#137 id=137 data-nosnippet>137</a> table: TableRef,
<a href=#138 id=138 data-nosnippet>138</a> req: CopyTableRequest,
<a href=#139 id=139 data-nosnippet>139</a> query_ctx: QueryContextRef,
<a href=#140 id=140 data-nosnippet>140</a> managed: <span class="prelude-ty">Option</span>&lt;(<span class="kw-2">&amp;</span>ExportWriteBudget, <span class="kw-2">&amp;</span>CancellationToken)&gt;,
<a href=#141 id=141 data-nosnippet>141</a> packed: <span class="prelude-ty">Option</span>&lt;PackedTableWriter&gt;,
<a href=#142 id=142 data-nosnippet>142</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#143 id=143 data-nosnippet>143</a> <span class="kw">let </span>info = table.table_info();
<a href=#144 id=144 data-nosnippet>144</a> <span class="kw">let </span>table_ref = TableReference::full(<span class="kw-2">&amp;</span>info.catalog_name, <span class="kw-2">&amp;</span>info.schema_name, <span class="kw-2">&amp;</span>info.name);
<a href=#145 id=145 data-nosnippet>145</a> <span class="kw">let </span>table_id = info.table_id();
<a href=#146 id=146 data-nosnippet>146</a> <span class="kw">let </span>format = Format::try_from(<span class="kw-2">&amp;</span>req.with).context(error::ParseFileFormatSnafu)<span class="question-mark">?</span>;
<a href=#147 id=147 data-nosnippet>147</a>
<a href=#148 id=148 data-nosnippet>148</a> <span class="kw">let </span>df_table_ref = DfTableReference::from(table_ref);
<a href=#149 id=149 data-nosnippet>149</a>
<a href=#150 id=150 data-nosnippet>150</a> <span class="kw">let </span>filters = table
<a href=#151 id=151 data-nosnippet>151</a> .schema()
<a href=#152 id=152 data-nosnippet>152</a> .timestamp_column()
<a href=#153 id=153 data-nosnippet>153</a> .and_then(|c| {
<a href=#154 id=154 data-nosnippet>154</a> common_query::logical_plan::build_filter_from_timestamp(
<a href=#155 id=155 data-nosnippet>155</a> <span class="kw-2">&amp;</span>c.name,
<a href=#156 id=156 data-nosnippet>156</a> req.timestamp_range.as_ref(),
<a href=#157 id=157 data-nosnippet>157</a> )
<a href=#158 id=158 data-nosnippet>158</a> })
<a href=#159 id=159 data-nosnippet>159</a> .into_iter()
<a href=#160 id=160 data-nosnippet>160</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#161 id=161 data-nosnippet>161</a>
<a href=#162 id=162 data-nosnippet>162</a> <span class="kw">let </span>table_provider = Arc::new(DfTableProviderAdapter::new(table));
<a href=#163 id=163 data-nosnippet>163</a> <span class="kw">let </span>table_source = Arc::new(DefaultTableSource::new(table_provider));
<a href=#164 id=164 data-nosnippet>164</a>
<a href=#165 id=165 data-nosnippet>165</a> <span class="kw">let </span><span class="kw-2">mut </span>builder = LogicalPlanBuilder::scan_with_filters(
<a href=#166 id=166 data-nosnippet>166</a> df_table_ref,
<a href=#167 id=167 data-nosnippet>167</a> table_source,
<a href=#168 id=168 data-nosnippet>168</a> <span class="prelude-val">None</span>,
<a href=#169 id=169 data-nosnippet>169</a> filters.clone(),
<a href=#170 id=170 data-nosnippet>170</a> )
<a href=#171 id=171 data-nosnippet>171</a> .context(BuildDfLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#172 id=172 data-nosnippet>172</a> <span class="kw">for </span>f <span class="kw">in </span>filters {
<a href=#173 id=173 data-nosnippet>173</a> builder = builder.filter(f).context(BuildDfLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#174 id=174 data-nosnippet>174</a> }
<a href=#175 id=175 data-nosnippet>175</a> <span class="kw">let </span>plan = builder.build().context(BuildDfLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#176 id=176 data-nosnippet>176</a>
<a href=#177 id=177 data-nosnippet>177</a> <span class="kw">let </span>output = <span class="self">self
<a href=#178 id=178 data-nosnippet>178</a> </span>.query_engine
<a href=#179 id=179 data-nosnippet>179</a> .execute(plan, query_ctx)
<a href=#180 id=180 data-nosnippet>180</a> .<span class="kw">await
<a href=#181 id=181 data-nosnippet>181</a> </span>.context(ExecLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#182 id=182 data-nosnippet>182</a>
<a href=#183 id=183 data-nosnippet>183</a> <span class="kw">let </span>CopyTableRequest {
<a href=#184 id=184 data-nosnippet>184</a> location,
<a href=#185 id=185 data-nosnippet>185</a> connection,
<a href=#186 id=186 data-nosnippet>186</a> ..
<a href=#187 id=187 data-nosnippet>187</a> } = <span class="kw-2">&amp;</span>req;
<a href=#188 id=188 data-nosnippet>188</a>
<a href=#189 id=189 data-nosnippet>189</a> <span class="macro">debug!</span>(<span class="string">"Copy table: {table_id} to location: {location}"</span>);
<a href=#190 id=190 data-nosnippet>190</a> <span class="self">self</span>.copy_to_file_managed(<span class="kw-2">&amp;</span>format, output, location, connection, managed, packed)
<a href=#191 id=191 data-nosnippet>191</a> .<span class="kw">await
<a href=#192 id=192 data-nosnippet>192</a> </span>}
<a href=#193 id=193 data-nosnippet>193</a>
<a href=#194 id=194 data-nosnippet>194</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>copy_to_file(
<a href=#195 id=195 data-nosnippet>195</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#196 id=196 data-nosnippet>196</a> format: <span class="kw-2">&amp;</span>Format,
<a href=#197 id=197 data-nosnippet>197</a> output: Output,
<a href=#198 id=198 data-nosnippet>198</a> location: <span class="kw-2">&amp;</span>str,
<a href=#199 id=199 data-nosnippet>199</a> connection: <span class="kw-2">&amp;</span>HashMap&lt;String, String&gt;,
<a href=#200 id=200 data-nosnippet>200</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#201 id=201 data-nosnippet>201</a> <span class="self">self</span>.copy_to_file_managed(format, output, location, connection, <span class="prelude-val">None</span>, <span class="prelude-val">None</span>)
<a href=#202 id=202 data-nosnippet>202</a> .<span class="kw">await
<a href=#203 id=203 data-nosnippet>203</a> </span>}
<a href=#204 id=204 data-nosnippet>204</a>
<a href=#205 id=205 data-nosnippet>205</a> <span class="kw">async fn </span>copy_to_file_managed(
<a href=#206 id=206 data-nosnippet>206</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#207 id=207 data-nosnippet>207</a> format: <span class="kw-2">&amp;</span>Format,
<a href=#208 id=208 data-nosnippet>208</a> output: Output,
<a href=#209 id=209 data-nosnippet>209</a> location: <span class="kw-2">&amp;</span>str,
<a href=#210 id=210 data-nosnippet>210</a> connection: <span class="kw-2">&amp;</span>HashMap&lt;String, String&gt;,
<a href=#211 id=211 data-nosnippet>211</a> managed: <span class="prelude-ty">Option</span>&lt;(<span class="kw-2">&amp;</span>ExportWriteBudget, <span class="kw-2">&amp;</span>CancellationToken)&gt;,
<a href=#212 id=212 data-nosnippet>212</a> packed: <span class="prelude-ty">Option</span>&lt;PackedTableWriter&gt;,
<a href=#213 id=213 data-nosnippet>213</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#214 id=214 data-nosnippet>214</a> <span class="kw">let </span>output = <span class="kw">if </span>managed.is_none() {
<a href=#215 id=215 data-nosnippet>215</a> output
<a href=#216 id=216 data-nosnippet>216</a> .map_dictionary_to_values()
<a href=#217 id=217 data-nosnippet>217</a> .context(error::BuildRecordBatchSnafu)<span class="question-mark">?
<a href=#218 id=218 data-nosnippet>218</a> </span>} <span class="kw">else </span>{
<a href=#219 id=219 data-nosnippet>219</a> output
<a href=#220 id=220 data-nosnippet>220</a> };
<a href=#221 id=221 data-nosnippet>221</a> <span class="kw">let </span>stream = <span class="kw">match </span>output.data {
<a href=#222 id=222 data-nosnippet>222</a> OutputData::Stream(stream) =&gt; stream,
<a href=#223 id=223 data-nosnippet>223</a> OutputData::RecordBatches(record_batches) =&gt; record_batches.as_stream(),
<a href=#224 id=224 data-nosnippet>224</a> <span class="kw">_ </span>=&gt; <span class="macro">unreachable!</span>(),
<a href=#225 id=225 data-nosnippet>225</a> };
<a href=#226 id=226 data-nosnippet>226</a>
<a href=#227 id=227 data-nosnippet>227</a> <span class="kw">let </span>backend =
<a href=#228 id=228 data-nosnippet>228</a> build_backend_for_write_with_path(location, connection, <span class="kw-2">&amp;</span><span class="self">self</span>.local_file_access)
<a href=#229 id=229 data-nosnippet>229</a> .<span class="kw">await
<a href=#230 id=230 data-nosnippet>230</a> </span>.context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#231 id=231 data-nosnippet>231</a> <span class="kw">let </span>filename = backend.object_path.context(error::UnexpectedSnafu {
<a href=#232 id=232 data-nosnippet>232</a> violated: <span class="macro">format!</span>(<span class="string">"Expected filename, path: {location}"</span>),
<a href=#233 id=233 data-nosnippet>233</a> })<span class="question-mark">?</span>;
<a href=#234 id=234 data-nosnippet>234</a> <span class="kw">if let </span><span class="prelude-val">Some</span>((budget, token)) = managed {
<a href=#235 id=235 data-nosnippet>235</a> stream_to_managed_parquet_with_packed(
<a href=#236 id=236 data-nosnippet>236</a> stream,
<a href=#237 id=237 data-nosnippet>237</a> backend.object_store,
<a href=#238 id=238 data-nosnippet>238</a> <span class="kw-2">&amp;</span>filename,
<a href=#239 id=239 data-nosnippet>239</a> budget,
<a href=#240 id=240 data-nosnippet>240</a> token,
<a href=#241 id=241 data-nosnippet>241</a> packed,
<a href=#242 id=242 data-nosnippet>242</a> )
<a href=#243 id=243 data-nosnippet>243</a> .<span class="kw">await
<a href=#244 id=244 data-nosnippet>244</a> </span>} <span class="kw">else </span>{
<a href=#245 id=245 data-nosnippet>245</a> <span class="self">self</span>.stream_to_file(stream, format, backend.object_store, <span class="kw-2">&amp;</span>filename)
<a href=#246 id=246 data-nosnippet>246</a> .<span class="kw">await
<a href=#247 id=247 data-nosnippet>247</a> </span>}
<a href=#248 id=248 data-nosnippet>248</a> }
<a href=#249 id=249 data-nosnippet>249</a>}
<a href=#250 id=250 data-nosnippet>250</a>
<a href=#251 id=251 data-nosnippet>251</a><span class="attr">#[cfg(test)]
<a href=#252 id=252 data-nosnippet>252</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>stream_to_managed_parquet(
<a href=#253 id=253 data-nosnippet>253</a> stream: SendableRecordBatchStream,
<a href=#254 id=254 data-nosnippet>254</a> store: ObjectStore,
<a href=#255 id=255 data-nosnippet>255</a> path: <span class="kw-2">&amp;</span>str,
<a href=#256 id=256 data-nosnippet>256</a> budget: <span class="kw-2">&amp;</span>ExportWriteBudget,
<a href=#257 id=257 data-nosnippet>257</a> token: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#258 id=258 data-nosnippet>258</a>) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#259 id=259 data-nosnippet>259</a> stream_to_managed_parquet_with_packed(stream, store, path, budget, token, <span class="prelude-val">None</span>).<span class="kw">await
<a href=#260 id=260 data-nosnippet>260</a></span>}
<a href=#261 id=261 data-nosnippet>261</a>
<a href=#262 id=262 data-nosnippet>262</a><span class="kw">async fn </span>stream_to_managed_parquet_with_packed(
<a href=#263 id=263 data-nosnippet>263</a> <span class="kw-2">mut </span>stream: SendableRecordBatchStream,
<a href=#264 id=264 data-nosnippet>264</a> store: ObjectStore,
<a href=#265 id=265 data-nosnippet>265</a> path: <span class="kw-2">&amp;</span>str,
<a href=#266 id=266 data-nosnippet>266</a> budget: <span class="kw-2">&amp;</span>ExportWriteBudget,
<a href=#267 id=267 data-nosnippet>267</a> token: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#268 id=268 data-nosnippet>268</a> packed: <span class="prelude-ty">Option</span>&lt;PackedTableWriter&gt;,
<a href=#269 id=269 data-nosnippet>269</a>) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#270 id=270 data-nosnippet>270</a> <span class="kw">use </span>common_recordbatch::{RecordBatch, map_dictionary_to_values_schema};
<a href=#271 id=271 data-nosnippet>271</a> <span class="kw">let </span>original = stream.schema();
<a href=#272 id=272 data-nosnippet>272</a> <span class="kw">let </span>(expanded_schema, expand) = map_dictionary_to_values_schema(original.clone());
<a href=#273 id=273 data-nosnippet>273</a> <span class="kw">let </span>json_columns = expanded_schema
<a href=#274 id=274 data-nosnippet>274</a> .column_schemas()
<a href=#275 id=275 data-nosnippet>275</a> .iter()
<a href=#276 id=276 data-nosnippet>276</a> .enumerate()
<a href=#277 id=277 data-nosnippet>277</a> .filter_map(|(index, column)| {
<a href=#278 id=278 data-nosnippet>278</a> (column.data_type.is_json() &amp;&amp; !column.data_type.is_json2()).then_some(index)
<a href=#279 id=279 data-nosnippet>279</a> })
<a href=#280 id=280 data-nosnippet>280</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#281 id=281 data-nosnippet>281</a> <span class="kw">let </span>(mapped_schema, json) = map_json_type_to_string_schema(expanded_schema.clone());
<a href=#282 id=282 data-nosnippet>282</a> <span class="kw">let </span>output_schema = <span class="kw">if </span>json {
<a href=#283 id=283 data-nosnippet>283</a> mapped_schema
<a href=#284 id=284 data-nosnippet>284</a> } <span class="kw">else </span>{
<a href=#285 id=285 data-nosnippet>285</a> expanded_schema.clone()
<a href=#286 id=286 data-nosnippet>286</a> };
<a href=#287 id=287 data-nosnippet>287</a> <span class="kw">let </span>packed_destination = packed.is_some();
<a href=#288 id=288 data-nosnippet>288</a> <span class="kw">let </span><span class="kw-2">mut </span>writer = <span class="kw">if let </span><span class="prelude-val">Some</span>(packed) = packed {
<a href=#289 id=289 data-nosnippet>289</a> ParquetFileWriter::open_packed(
<a href=#290 id=290 data-nosnippet>290</a> output_schema.arrow_schema().clone(),
<a href=#291 id=291 data-nosnippet>291</a> store,
<a href=#292 id=292 data-nosnippet>292</a> path,
<a href=#293 id=293 data-nosnippet>293</a> <span class="prelude-val">Some</span>(
<a href=#294 id=294 data-nosnippet>294</a> <span class="kw">crate</span>::statement::export_logical_tables::LogicalTableExportLimits::default().writer,
<a href=#295 id=295 data-nosnippet>295</a> ),
<a href=#296 id=296 data-nosnippet>296</a> packed,
<a href=#297 id=297 data-nosnippet>297</a> )
<a href=#298 id=298 data-nosnippet>298</a> } <span class="kw">else </span>{
<a href=#299 id=299 data-nosnippet>299</a> ParquetFileWriter::open(output_schema.arrow_schema().clone(), store, path, <span class="number">1</span>, <span class="prelude-val">None</span>).<span class="kw">await
<a href=#300 id=300 data-nosnippet>300</a> </span>}
<a href=#301 id=301 data-nosnippet>301</a> .context(error::WriteStreamToFileSnafu { path })<span class="question-mark">?</span>;
<a href=#302 id=302 data-nosnippet>302</a> <span class="kw">let </span><span class="kw-2">mut </span>started = <span class="bool-val">false</span>;
<a href=#303 id=303 data-nosnippet>303</a> <span class="kw">let </span>result = <span class="kw">async </span>{
<a href=#304 id=304 data-nosnippet>304</a> <span class="kw">let </span><span class="kw-2">mut </span>rows = <span class="number">0</span>;
<a href=#305 id=305 data-nosnippet>305</a> <span class="kw">loop </span>{
<a href=#306 id=306 data-nosnippet>306</a> <span class="kw">let </span>batch = <span class="macro">tokio::select!</span> {
<a href=#307 id=307 data-nosnippet>307</a> biased;
<a href=#308 id=308 data-nosnippet>308</a> <span class="kw">_ </span>= token.cancelled() =&gt; <span class="kw">return </span>error::LogicalTableExportCancelledSnafu.fail(),
<a href=#309 id=309 data-nosnippet>309</a> batch = stream.next() =&gt; batch,
<a href=#310 id=310 data-nosnippet>310</a> };
<a href=#311 id=311 data-nosnippet>311</a> <span class="kw">let </span><span class="prelude-val">Some</span>(batch) = batch <span class="kw">else </span>{ <span class="kw">break </span>};
<a href=#312 id=312 data-nosnippet>312</a> <span class="kw">let </span>batch = batch
<a href=#313 id=313 data-nosnippet>313</a> .context(error::BuildRecordBatchSnafu)<span class="question-mark">?
<a href=#314 id=314 data-nosnippet>314</a> </span>.into_df_record_batch();
<a href=#315 id=315 data-nosnippet>315</a> <span class="kw">let </span><span class="kw-2">mut </span>offset = <span class="number">0</span>;
<a href=#316 id=316 data-nosnippet>316</a> <span class="kw">while </span>offset &lt; batch.num_rows() {
<a href=#317 id=317 data-nosnippet>317</a> <span class="comment">// The scan batch remains query-owned. Bound the downstream slice and
<a href=#318 id=318 data-nosnippet>318</a> // detach its buffers below if it would retain more than its reservation.
<a href=#319 id=319 data-nosnippet>319</a> </span><span class="kw">let </span>(conversion, retained) =
<a href=#320 id=320 data-nosnippet>320</a> ExportWriteBudget::conversion_budget(<span class="number">0</span>, batch.num_columns(), usize::MAX)<span class="question-mark">?</span>;
<a href=#321 id=321 data-nosnippet>321</a> <span class="kw">let </span>input = batch.clone();
<a href=#322 id=322 data-nosnippet>322</a> <span class="kw">let </span>json_columns = json_columns.clone();
<a href=#323 id=323 data-nosnippet>323</a> <span class="kw">let </span>(len, estimated) = common_runtime::spawn_blocking_global(<span class="kw">move </span>|| {
<a href=#324 id=324 data-nosnippet>324</a> rows_within_budget(<span class="kw-2">&amp;</span>input, offset, input.num_rows(), conversion, <span class="kw-2">&amp;</span>json_columns)
<a href=#325 id=325 data-nosnippet>325</a> })
<a href=#326 id=326 data-nosnippet>326</a> .<span class="kw">await
<a href=#327 id=327 data-nosnippet>327</a> </span>.context(error::JoinTaskSnafu)<span class="question-mark">??</span>;
<a href=#328 id=328 data-nosnippet>328</a> <span class="kw">let </span>reservation = retained.saturating_add(estimated.saturating_mul(<span class="number">4</span>));
<a href=#329 id=329 data-nosnippet>329</a> <span class="kw">let </span>permit = budget.reserve(reservation, token).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#330 id=330 data-nosnippet>330</a> <span class="kw">let </span>batch = batch.clone();
<a href=#331 id=331 data-nosnippet>331</a> <span class="kw">let </span>(original, expanded_schema, output_schema) = (
<a href=#332 id=332 data-nosnippet>332</a> original.clone(),
<a href=#333 id=333 data-nosnippet>333</a> expanded_schema.clone(),
<a href=#334 id=334 data-nosnippet>334</a> output_schema.clone(),
<a href=#335 id=335 data-nosnippet>335</a> );
<a href=#336 id=336 data-nosnippet>336</a> <span class="kw">let </span>(batch, len) = common_runtime::spawn_blocking_global(<span class="kw">move </span>|| {
<a href=#337 id=337 data-nosnippet>337</a> <span class="kw">let </span><span class="kw-2">mut </span>batch = RecordBatch::from_df_record_batch(
<a href=#338 id=338 data-nosnippet>338</a> original.clone(),
<a href=#339 id=339 data-nosnippet>339</a> batch.slice(offset, len),
<a href=#340 id=340 data-nosnippet>340</a> );
<a href=#341 id=341 data-nosnippet>341</a> <span class="kw">if </span>expand {
<a href=#342 id=342 data-nosnippet>342</a> batch = RecordBatch::from_df_record_batch(
<a href=#343 id=343 data-nosnippet>343</a> expanded_schema.clone(),
<a href=#344 id=344 data-nosnippet>344</a> expand_export_batch(
<a href=#345 id=345 data-nosnippet>345</a> <span class="kw-2">&amp;</span>batch.into_df_record_batch(),
<a href=#346 id=346 data-nosnippet>346</a> expanded_schema.arrow_schema().clone(),
<a href=#347 id=347 data-nosnippet>347</a> )<span class="question-mark">?</span>,
<a href=#348 id=348 data-nosnippet>348</a> );
<a href=#349 id=349 data-nosnippet>349</a> }
<a href=#350 id=350 data-nosnippet>350</a> <span class="kw">if </span>json {
<a href=#351 id=351 data-nosnippet>351</a> batch = map_json_type_to_string(batch, <span class="kw-2">&amp;</span>expanded_schema, <span class="kw-2">&amp;</span>output_schema)
<a href=#352 id=352 data-nosnippet>352</a> .context(error::BuildRecordBatchSnafu)<span class="question-mark">?</span>;
<a href=#353 id=353 data-nosnippet>353</a> }
<a href=#354 id=354 data-nosnippet>354</a> <span class="kw">let </span><span class="kw-2">mut </span>batch = batch.into_df_record_batch();
<a href=#355 id=355 data-nosnippet>355</a> <span class="kw">if </span>batch.get_array_memory_size() &gt; reservation {
<a href=#356 id=356 data-nosnippet>356</a> <span class="kw">let </span>indices = UInt64Array::from_iter_values(<span class="number">0</span>..batch.num_rows() <span class="kw">as </span>u64);
<a href=#357 id=357 data-nosnippet>357</a> batch = arrow::compute::take_record_batch(<span class="kw-2">&amp;</span>batch, <span class="kw-2">&amp;</span>indices)
<a href=#358 id=358 data-nosnippet>358</a> .context(error::ComputeArrowSnafu)<span class="question-mark">?</span>;
<a href=#359 id=359 data-nosnippet>359</a> <span class="kw">let </span>arrays = batch
<a href=#360 id=360 data-nosnippet>360</a> .columns()
<a href=#361 id=361 data-nosnippet>361</a> .iter()
<a href=#362 id=362 data-nosnippet>362</a> .map(compact_view_buffers)
<a href=#363 id=363 data-nosnippet>363</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#364 id=364 data-nosnippet>364</a> batch = arrow::record_batch::RecordBatch::try_new(batch.schema(), arrays)
<a href=#365 id=365 data-nosnippet>365</a> .context(error::ComputeArrowSnafu)<span class="question-mark">?</span>;
<a href=#366 id=366 data-nosnippet>366</a> }
<a href=#367 id=367 data-nosnippet>367</a> <span class="macro">ensure!</span>(
<a href=#368 id=368 data-nosnippet>368</a> batch.get_array_memory_size() &lt;= reservation,
<a href=#369 id=369 data-nosnippet>369</a> error::LogicalTableExportResourceSnafu {
<a href=#370 id=370 data-nosnippet>370</a> reason: <span class="string">"converted backing buffers exceed reservation"
<a href=#371 id=371 data-nosnippet>371</a> </span>}
<a href=#372 id=372 data-nosnippet>372</a> );
<a href=#373 id=373 data-nosnippet>373</a> <span class="prelude-val">Ok</span>::&lt;<span class="kw">_</span>, error::Error&gt;((batch, len))
<a href=#374 id=374 data-nosnippet>374</a> })
<a href=#375 id=375 data-nosnippet>375</a> .<span class="kw">await
<a href=#376 id=376 data-nosnippet>376</a> </span>.context(error::JoinTaskSnafu)<span class="question-mark">??</span>;
<a href=#377 id=377 data-nosnippet>377</a> started = <span class="bool-val">true</span>;
<a href=#378 id=378 data-nosnippet>378</a> <span class="kw">let </span>write = writer.write(batch, <span class="prelude-val">Some</span>(token)).<span class="kw">await</span>;
<a href=#379 id=379 data-nosnippet>379</a> drop(permit);
<a href=#380 id=380 data-nosnippet>380</a> write.map_err(|error| map_writer_error(error, path))<span class="question-mark">?</span>;
<a href=#381 id=381 data-nosnippet>381</a> rows += len;
<a href=#382 id=382 data-nosnippet>382</a> offset += len;
<a href=#383 id=383 data-nosnippet>383</a> }
<a href=#384 id=384 data-nosnippet>384</a> }
<a href=#385 id=385 data-nosnippet>385</a> started = <span class="bool-val">true</span>;
<a href=#386 id=386 data-nosnippet>386</a> writer
<a href=#387 id=387 data-nosnippet>387</a> .finish(<span class="prelude-val">Some</span>(token))
<a href=#388 id=388 data-nosnippet>388</a> .<span class="kw">await
<a href=#389 id=389 data-nosnippet>389</a> </span>.map_err(|error| map_writer_error(error, path))<span class="question-mark">?</span>;
<a href=#390 id=390 data-nosnippet>390</a> <span class="prelude-val">Ok</span>(rows)
<a href=#391 id=391 data-nosnippet>391</a> }
<a href=#392 id=392 data-nosnippet>392</a> .<span class="kw">await</span>;
<a href=#393 id=393 data-nosnippet>393</a> <span class="kw">if </span>result.is_err() {
<a href=#394 id=394 data-nosnippet>394</a> token.cancel();
<a href=#395 id=395 data-nosnippet>395</a> <span class="kw">if </span>(started || packed_destination)
<a href=#396 id=396 data-nosnippet>396</a> &amp;&amp; <span class="kw">let </span><span class="prelude-val">Err</span>(error) = writer.abort().<span class="kw">await
<a href=#397 id=397 data-nosnippet>397</a> </span>{
<a href=#398 id=398 data-nosnippet>398</a> <span class="macro">common_telemetry::warn!</span>(error; <span class="string">"Failed to abort ordinary export file"</span>);
<a href=#399 id=399 data-nosnippet>399</a> }
<a href=#400 id=400 data-nosnippet>400</a> }
<a href=#401 id=401 data-nosnippet>401</a> result
<a href=#402 id=402 data-nosnippet>402</a>}
<a href=#403 id=403 data-nosnippet>403</a>
<a href=#404 id=404 data-nosnippet>404</a><span class="comment">// Arrow take copies ordinary buffers but shares view data buffers, including
<a href=#405 id=405 data-nosnippet>405</a>// views nested inside lists or structs. Reclaim those unselected values too.
<a href=#406 id=406 data-nosnippet>406</a></span><span class="kw">fn </span>compact_view_buffers(array: <span class="kw-2">&amp;</span>ArrayRef) -&gt; <span class="prelude-ty">Result</span>&lt;ArrayRef&gt; {
<a href=#407 id=407 data-nosnippet>407</a> <span class="kw">match </span>array.data_type() {
<a href=#408 id=408 data-nosnippet>408</a> DataType::Utf8View =&gt; <span class="prelude-val">Ok</span>(Arc::new(array.as_string_view().gc())),
<a href=#409 id=409 data-nosnippet>409</a> DataType::BinaryView =&gt; <span class="prelude-val">Ok</span>(Arc::new(array.as_binary_view().gc())),
<a href=#410 id=410 data-nosnippet>410</a> <span class="kw">_ </span>=&gt; {
<a href=#411 id=411 data-nosnippet>411</a> <span class="kw">let </span>data = array.to_data();
<a href=#412 id=412 data-nosnippet>412</a> <span class="kw">if </span>data.child_data().is_empty() {
<a href=#413 id=413 data-nosnippet>413</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(array.clone());
<a href=#414 id=414 data-nosnippet>414</a> }
<a href=#415 id=415 data-nosnippet>415</a> <span class="kw">let </span>children = data
<a href=#416 id=416 data-nosnippet>416</a> .child_data()
<a href=#417 id=417 data-nosnippet>417</a> .iter()
<a href=#418 id=418 data-nosnippet>418</a> .map(|child| {
<a href=#419 id=419 data-nosnippet>419</a> compact_view_buffers(<span class="kw-2">&amp;</span>make_array(child.clone())).map(|array| array.to_data())
<a href=#420 id=420 data-nosnippet>420</a> })
<a href=#421 id=421 data-nosnippet>421</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#422 id=422 data-nosnippet>422</a> <span class="prelude-val">Ok</span>(make_array(
<a href=#423 id=423 data-nosnippet>423</a> data.into_builder()
<a href=#424 id=424 data-nosnippet>424</a> .child_data(children)
<a href=#425 id=425 data-nosnippet>425</a> .build()
<a href=#426 id=426 data-nosnippet>426</a> .context(error::ComputeArrowSnafu)<span class="question-mark">?</span>,
<a href=#427 id=427 data-nosnippet>427</a> ))
<a href=#428 id=428 data-nosnippet>428</a> }
<a href=#429 id=429 data-nosnippet>429</a> }
<a href=#430 id=430 data-nosnippet>430</a>}
<a href=#431 id=431 data-nosnippet>431</a>
<a href=#432 id=432 data-nosnippet>432</a><span class="attr">#[cfg(test)]
<a href=#433 id=433 data-nosnippet>433</a></span><span class="kw">mod </span>tests {
<a href=#434 id=434 data-nosnippet>434</a> <span class="kw">use </span>arrow::array::{
<a href=#435 id=435 data-nosnippet>435</a> ArrayRef, BinaryArray, BinaryViewArray, DictionaryArray, Int32Array, StringArray,
<a href=#436 id=436 data-nosnippet>436</a> StringViewArray,
<a href=#437 id=437 data-nosnippet>437</a> };
<a href=#438 id=438 data-nosnippet>438</a> <span class="kw">use </span>arrow::datatypes::Int32Type;
<a href=#439 id=439 data-nosnippet>439</a> <span class="kw">use </span>common_recordbatch::{RecordBatch, RecordBatches};
<a href=#440 id=440 data-nosnippet>440</a> <span class="kw">use </span>datafusion::parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
<a href=#441 id=441 data-nosnippet>441</a> <span class="kw">use </span>datatypes::prelude::ConcreteDataType;
<a href=#442 id=442 data-nosnippet>442</a> <span class="kw">use </span>datatypes::schema::{ColumnSchema, Schema};
<a href=#443 id=443 data-nosnippet>443</a>
<a href=#444 id=444 data-nosnippet>444</a> <span class="kw">use super</span>::<span class="kw-2">*</span>;
<a href=#445 id=445 data-nosnippet>445</a>
<a href=#446 id=446 data-nosnippet>446</a> <span class="attr">#[tokio::test]
<a href=#447 id=447 data-nosnippet>447</a> </span><span class="kw">async fn </span>managed_copy_preserves_existing_file_before_sink_open() {
<a href=#448 id=448 data-nosnippet>448</a> <span class="kw">let </span>temp_dir = common_test_util::temp_dir::create_temp_dir(<span class="string">"managed_copy_existing_file"</span>);
<a href=#449 id=449 data-nosnippet>449</a> <span class="kw">let </span>store = object_store::secure_fs::SecureFsRoot::open(temp_dir.path())
<a href=#450 id=450 data-nosnippet>450</a> .unwrap()
<a href=#451 id=451 data-nosnippet>451</a> .build_operator();
<a href=#452 id=452 data-nosnippet>452</a> <span class="kw">let </span>path = <span class="string">"existing.parquet"</span>;
<a href=#453 id=453 data-nosnippet>453</a> store.write(path, <span class="string">"original"</span>).<span class="kw">await</span>.unwrap();
<a href=#454 id=454 data-nosnippet>454</a> <span class="kw">let </span>schema = Arc::new(Schema::new(<span class="macro">vec!</span>[ColumnSchema::new(
<a href=#455 id=455 data-nosnippet>455</a> <span class="string">"value"</span>,
<a href=#456 id=456 data-nosnippet>456</a> ConcreteDataType::int32_datatype(),
<a href=#457 id=457 data-nosnippet>457</a> <span class="bool-val">false</span>,
<a href=#458 id=458 data-nosnippet>458</a> )]));
<a href=#459 id=459 data-nosnippet>459</a> <span class="kw">let </span>batch = arrow::record_batch::RecordBatch::try_new(
<a href=#460 id=460 data-nosnippet>460</a> schema.arrow_schema().clone(),
<a href=#461 id=461 data-nosnippet>461</a> <span class="macro">vec!</span>[Arc::new(Int32Array::from(<span class="macro">vec!</span>[<span class="number">1</span>]))],
<a href=#462 id=462 data-nosnippet>462</a> )
<a href=#463 id=463 data-nosnippet>463</a> .unwrap();
<a href=#464 id=464 data-nosnippet>464</a> <span class="kw">let </span>batches = <span class="macro">async_stream::stream!</span> {
<a href=#465 id=465 data-nosnippet>465</a> <span class="kw">yield </span><span class="prelude-val">Ok</span>(batch);
<a href=#466 id=466 data-nosnippet>466</a> <span class="kw">yield </span><span class="prelude-val">Err</span>(datafusion::error::DataFusionError::Execution(<span class="string">"source failed"</span>.into()));
<a href=#467 id=467 data-nosnippet>467</a> };
<a href=#468 id=468 data-nosnippet>468</a> <span class="kw">let </span>stream = datafusion::physical_plan::stream::RecordBatchStreamAdapter::new(
<a href=#469 id=469 data-nosnippet>469</a> schema.arrow_schema().clone(),
<a href=#470 id=470 data-nosnippet>470</a> batches,
<a href=#471 id=471 data-nosnippet>471</a> );
<a href=#472 id=472 data-nosnippet>472</a> <span class="kw">let </span>stream =
<a href=#473 id=473 data-nosnippet>473</a> common_recordbatch::adapter::RecordBatchStreamAdapter::try_new(Box::pin(stream))
<a href=#474 id=474 data-nosnippet>474</a> .unwrap();
<a href=#475 id=475 data-nosnippet>475</a>
<a href=#476 id=476 data-nosnippet>476</a> <span class="kw">let </span>result = stream_to_managed_parquet(
<a href=#477 id=477 data-nosnippet>477</a> Box::pin(stream),
<a href=#478 id=478 data-nosnippet>478</a> store.clone(),
<a href=#479 id=479 data-nosnippet>479</a> path,
<a href=#480 id=480 data-nosnippet>480</a> <span class="kw-2">&amp;</span>ExportWriteBudget::new(<span class="number">1</span>),
<a href=#481 id=481 data-nosnippet>481</a> <span class="kw-2">&amp;</span>CancellationToken::new(),
<a href=#482 id=482 data-nosnippet>482</a> )
<a href=#483 id=483 data-nosnippet>483</a> .<span class="kw">await</span>;
<a href=#484 id=484 data-nosnippet>484</a>
<a href=#485 id=485 data-nosnippet>485</a> <span class="macro">assert!</span>(result.is_err());
<a href=#486 id=486 data-nosnippet>486</a> <span class="macro">assert_eq!</span>(
<a href=#487 id=487 data-nosnippet>487</a> store.read(path).<span class="kw">await</span>.unwrap().to_bytes().as_ref(),
<a href=#488 id=488 data-nosnippet>488</a> <span class="string">b"original"
<a href=#489 id=489 data-nosnippet>489</a> </span>);
<a href=#490 id=490 data-nosnippet>490</a> }
<a href=#491 id=491 data-nosnippet>491</a>
<a href=#492 id=492 data-nosnippet>492</a> <span class="attr">#[tokio::test]
<a href=#493 id=493 data-nosnippet>493</a> </span><span class="kw">async fn </span>managed_copy_rechunks_large_backing_buffers() {
<a href=#494 id=494 data-nosnippet>494</a> <span class="kw">let </span>value = <span class="string">"x"</span>.repeat(<span class="number">9 </span>* <span class="number">1024</span>);
<a href=#495 id=495 data-nosnippet>495</a> <span class="kw">for </span>view <span class="kw">in </span>[<span class="bool-val">false</span>, <span class="bool-val">true</span>] {
<a href=#496 id=496 data-nosnippet>496</a> <span class="kw">let </span>data_type = <span class="kw">if </span>view {
<a href=#497 id=497 data-nosnippet>497</a> ConcreteDataType::utf8_view_datatype()
<a href=#498 id=498 data-nosnippet>498</a> } <span class="kw">else </span>{
<a href=#499 id=499 data-nosnippet>499</a> ConcreteDataType::string_datatype()
<a href=#500 id=500 data-nosnippet>500</a> };
<a href=#501 id=501 data-nosnippet>501</a> <span class="kw">let </span>schema = Arc::new(Schema::new(<span class="macro">vec!</span>[ColumnSchema::new(
<a href=#502 id=502 data-nosnippet>502</a> <span class="string">"value"</span>, data_type, <span class="bool-val">false</span>,
<a href=#503 id=503 data-nosnippet>503</a> )]));
<a href=#504 id=504 data-nosnippet>504</a> <span class="kw">let </span>values = <span class="macro">vec!</span>[value.as_str(); <span class="number">8192</span>];
<a href=#505 id=505 data-nosnippet>505</a> <span class="kw">let </span>array: ArrayRef = <span class="kw">if </span>view {
<a href=#506 id=506 data-nosnippet>506</a> Arc::new(StringViewArray::from(values))
<a href=#507 id=507 data-nosnippet>507</a> } <span class="kw">else </span>{
<a href=#508 id=508 data-nosnippet>508</a> Arc::new(StringArray::from(values))
<a href=#509 id=509 data-nosnippet>509</a> };
<a href=#510 id=510 data-nosnippet>510</a> <span class="kw">let </span>batch = arrow::record_batch::RecordBatch::try_new(
<a href=#511 id=511 data-nosnippet>511</a> schema.arrow_schema().clone(),
<a href=#512 id=512 data-nosnippet>512</a> <span class="macro">vec!</span>[array],
<a href=#513 id=513 data-nosnippet>513</a> )
<a href=#514 id=514 data-nosnippet>514</a> .unwrap();
<a href=#515 id=515 data-nosnippet>515</a> <span class="macro">assert!</span>(batch.get_array_memory_size() &gt; <span class="number">64 </span>* <span class="number">1024 </span>* <span class="number">1024</span>);
<a href=#516 id=516 data-nosnippet>516</a> <span class="comment">// Also exercise a small slice that still references the large allocation.
<a href=#517 id=517 data-nosnippet>517</a> </span><span class="kw">for </span>input <span class="kw">in </span>[batch.clone(), batch.slice(<span class="number">200</span>, <span class="number">3</span>)] {
<a href=#518 id=518 data-nosnippet>518</a> <span class="kw">let </span>rows = input.num_rows();
<a href=#519 id=519 data-nosnippet>519</a> <span class="kw">let </span>batches = RecordBatches::try_new(
<a href=#520 id=520 data-nosnippet>520</a> schema.clone(),
<a href=#521 id=521 data-nosnippet>521</a> <span class="macro">vec!</span>[RecordBatch::from_df_record_batch(schema.clone(), input)],
<a href=#522 id=522 data-nosnippet>522</a> )
<a href=#523 id=523 data-nosnippet>523</a> .unwrap();
<a href=#524 id=524 data-nosnippet>524</a> <span class="kw">let </span>store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
<a href=#525 id=525 data-nosnippet>525</a> <span class="kw">let </span>budget = ExportWriteBudget::new(<span class="number">1</span>);
<a href=#526 id=526 data-nosnippet>526</a> <span class="macro">assert_eq!</span>(
<a href=#527 id=527 data-nosnippet>527</a> stream_to_managed_parquet(
<a href=#528 id=528 data-nosnippet>528</a> batches.as_stream(),
<a href=#529 id=529 data-nosnippet>529</a> store.clone(),
<a href=#530 id=530 data-nosnippet>530</a> <span class="string">"large.parquet"</span>,
<a href=#531 id=531 data-nosnippet>531</a> <span class="kw-2">&amp;</span>budget,
<a href=#532 id=532 data-nosnippet>532</a> <span class="kw-2">&amp;</span>CancellationToken::new(),
<a href=#533 id=533 data-nosnippet>533</a> )
<a href=#534 id=534 data-nosnippet>534</a> .<span class="kw">await
<a href=#535 id=535 data-nosnippet>535</a> </span>.unwrap(),
<a href=#536 id=536 data-nosnippet>536</a> rows
<a href=#537 id=537 data-nosnippet>537</a> );
<a href=#538 id=538 data-nosnippet>538</a> <span class="kw">let </span>reader = ParquetRecordBatchReaderBuilder::try_new(
<a href=#539 id=539 data-nosnippet>539</a> store.read(<span class="string">"large.parquet"</span>).<span class="kw">await</span>.unwrap().to_bytes(),
<a href=#540 id=540 data-nosnippet>540</a> )
<a href=#541 id=541 data-nosnippet>541</a> .unwrap()
<a href=#542 id=542 data-nosnippet>542</a> .build()
<a href=#543 id=543 data-nosnippet>543</a> .unwrap();
<a href=#544 id=544 data-nosnippet>544</a> <span class="kw">let </span><span class="kw-2">mut </span>actual_rows = <span class="number">0</span>;
<a href=#545 id=545 data-nosnippet>545</a> <span class="kw">for </span>batch <span class="kw">in </span>reader {
<a href=#546 id=546 data-nosnippet>546</a> <span class="kw">let </span>batch = batch.unwrap();
<a href=#547 id=547 data-nosnippet>547</a> <span class="kw">let </span>strings = arrow::compute::cast(batch.column(<span class="number">0</span>), <span class="kw-2">&amp;</span>DataType::Utf8).unwrap();
<a href=#548 id=548 data-nosnippet>548</a> <span class="macro">assert!</span>(
<a href=#549 id=549 data-nosnippet>549</a> strings
<a href=#550 id=550 data-nosnippet>550</a> .as_string::&lt;i32&gt;()
<a href=#551 id=551 data-nosnippet>551</a> .iter()
<a href=#552 id=552 data-nosnippet>552</a> .all(|item| item == <span class="prelude-val">Some</span>(value.as_str()))
<a href=#553 id=553 data-nosnippet>553</a> );
<a href=#554 id=554 data-nosnippet>554</a> actual_rows += batch.num_rows();
<a href=#555 id=555 data-nosnippet>555</a> }
<a href=#556 id=556 data-nosnippet>556</a> <span class="macro">assert_eq!</span>(actual_rows, rows);
<a href=#557 id=557 data-nosnippet>557</a> <span class="macro">assert_eq!</span>(budget.available(), (<span class="number">1</span>, <span class="number">64 </span>* <span class="number">1024 </span>* <span class="number">1024</span>));
<a href=#558 id=558 data-nosnippet>558</a> }
<a href=#559 id=559 data-nosnippet>559</a> }
<a href=#560 id=560 data-nosnippet>560</a> }
<a href=#561 id=561 data-nosnippet>561</a>
<a href=#562 id=562 data-nosnippet>562</a> <span class="attr">#[tokio::test]
<a href=#563 id=563 data-nosnippet>563</a> </span><span class="kw">async fn </span>managed_copy_preserves_dictionary_json_views_and_empty_schema() {
<a href=#564 id=564 data-nosnippet>564</a> <span class="kw">let </span>schema = Arc::new(Schema::new(<span class="macro">vec!</span>[
<a href=#565 id=565 data-nosnippet>565</a> ColumnSchema::new(
<a href=#566 id=566 data-nosnippet>566</a> <span class="string">"host"</span>,
<a href=#567 id=567 data-nosnippet>567</a> ConcreteDataType::dictionary_datatype(
<a href=#568 id=568 data-nosnippet>568</a> ConcreteDataType::int32_datatype(),
<a href=#569 id=569 data-nosnippet>569</a> ConcreteDataType::string_datatype(),
<a href=#570 id=570 data-nosnippet>570</a> ),
<a href=#571 id=571 data-nosnippet>571</a> <span class="bool-val">true</span>,
<a href=#572 id=572 data-nosnippet>572</a> ),
<a href=#573 id=573 data-nosnippet>573</a> ColumnSchema::new(<span class="string">"json"</span>, ConcreteDataType::json_datatype(), <span class="bool-val">true</span>),
<a href=#574 id=574 data-nosnippet>574</a> ColumnSchema::new(<span class="string">"text_view"</span>, ConcreteDataType::utf8_view_datatype(), <span class="bool-val">true</span>),
<a href=#575 id=575 data-nosnippet>575</a> ColumnSchema::new(
<a href=#576 id=576 data-nosnippet>576</a> <span class="string">"binary_view"</span>,
<a href=#577 id=577 data-nosnippet>577</a> ConcreteDataType::binary_view_datatype(),
<a href=#578 id=578 data-nosnippet>578</a> <span class="bool-val">true</span>,
<a href=#579 id=579 data-nosnippet>579</a> ),
<a href=#580 id=580 data-nosnippet>580</a> ]));
<a href=#581 id=581 data-nosnippet>581</a> <span class="kw">let </span>dictionary = DictionaryArray::&lt;Int32Type&gt;::new(
<a href=#582 id=582 data-nosnippet>582</a> Int32Array::from(<span class="macro">vec!</span>[<span class="prelude-val">Some</span>(<span class="number">0</span>), <span class="prelude-val">None</span>, <span class="prelude-val">Some</span>(<span class="number">1</span>), <span class="prelude-val">Some</span>(<span class="number">1</span>)]),
<a href=#583 id=583 data-nosnippet>583</a> Arc::new(StringArray::from(<span class="macro">vec!</span>[<span class="string">"tag"</span>, <span class="kw-2">&amp;</span><span class="string">"x"</span>.repeat(<span class="number">3 </span>* <span class="number">1024 </span>* <span class="number">1024</span>)])),
<a href=#584 id=584 data-nosnippet>584</a> );
<a href=#585 id=585 data-nosnippet>585</a> <span class="kw">let </span>json = datatypes::types::parse_string_to_jsonb(<span class="kw-2">&amp;</span><span class="macro">format!</span>(
<a href=#586 id=586 data-nosnippet>586</a> <span class="string">r#"{{"value":"escaped\ntext{}","n":123456789}}"#</span>,
<a href=#587 id=587 data-nosnippet>587</a> <span class="string">"y"</span>.repeat(<span class="number">160_000</span>)
<a href=#588 id=588 data-nosnippet>588</a> ))
<a href=#589 id=589 data-nosnippet>589</a> .unwrap();
<a href=#590 id=590 data-nosnippet>590</a> <span class="kw">let </span>arrays = <span class="macro">vec!</span>[
<a href=#591 id=591 data-nosnippet>591</a> Arc::new(dictionary) <span class="kw">as </span>ArrayRef,
<a href=#592 id=592 data-nosnippet>592</a> Arc::new(BinaryArray::from(<span class="macro">vec!</span>[
<a href=#593 id=593 data-nosnippet>593</a> <span class="prelude-val">Some</span>(json.as_slice()),
<a href=#594 id=594 data-nosnippet>594</a> <span class="prelude-val">None</span>,
<a href=#595 id=595 data-nosnippet>595</a> <span class="prelude-val">None</span>,
<a href=#596 id=596 data-nosnippet>596</a> <span class="prelude-val">Some</span>(json.as_slice()),
<a href=#597 id=597 data-nosnippet>597</a> ])),
<a href=#598 id=598 data-nosnippet>598</a> Arc::new(StringViewArray::from(<span class="macro">vec!</span>[
<a href=#599 id=599 data-nosnippet>599</a> <span class="prelude-val">Some</span>(<span class="string">"short"</span>),
<a href=#600 id=600 data-nosnippet>600</a> <span class="prelude-val">None</span>,
<a href=#601 id=601 data-nosnippet>601</a> <span class="prelude-val">Some</span>(<span class="string">"long view"</span>),
<a href=#602 id=602 data-nosnippet>602</a> <span class="prelude-val">Some</span>(<span class="string">"long view"</span>),
<a href=#603 id=603 data-nosnippet>603</a> ])),
<a href=#604 id=604 data-nosnippet>604</a> Arc::new(BinaryViewArray::from(<span class="macro">vec!</span>[
<a href=#605 id=605 data-nosnippet>605</a> <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span><span class="string">b"short"</span>[..]),
<a href=#606 id=606 data-nosnippet>606</a> <span class="prelude-val">None</span>,
<a href=#607 id=607 data-nosnippet>607</a> <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span><span class="string">b"long view"</span>[..]),
<a href=#608 id=608 data-nosnippet>608</a> <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span><span class="string">b"long view"</span>[..]),
<a href=#609 id=609 data-nosnippet>609</a> ])),
<a href=#610 id=610 data-nosnippet>610</a> ];
<a href=#611 id=611 data-nosnippet>611</a> <span class="kw">let </span>batch =
<a href=#612 id=612 data-nosnippet>612</a> arrow::record_batch::RecordBatch::try_new(schema.arrow_schema().clone(), arrays)
<a href=#613 id=613 data-nosnippet>613</a> .unwrap();
<a href=#614 id=614 data-nosnippet>614</a> <span class="kw">let </span>batch = RecordBatch::from_df_record_batch(schema.clone(), batch);
<a href=#615 id=615 data-nosnippet>615</a> <span class="kw">for </span>empty <span class="kw">in </span>[<span class="bool-val">false</span>, <span class="bool-val">true</span>] {
<a href=#616 id=616 data-nosnippet>616</a> <span class="kw">let </span>batches = RecordBatches::try_new(
<a href=#617 id=617 data-nosnippet>617</a> schema.clone(),
<a href=#618 id=618 data-nosnippet>618</a> <span class="kw">if </span>empty { <span class="macro">vec!</span>[] } <span class="kw">else </span>{ <span class="macro">vec!</span>[batch.clone()] },
<a href=#619 id=619 data-nosnippet>619</a> )
<a href=#620 id=620 data-nosnippet>620</a> .unwrap();
<a href=#621 id=621 data-nosnippet>621</a> <span class="kw">let </span>store = ObjectStore::new(object_store::services::Memory::default()).unwrap();
<a href=#622 id=622 data-nosnippet>622</a> <span class="kw">let </span>budget = ExportWriteBudget::new(<span class="number">1</span>);
<a href=#623 id=623 data-nosnippet>623</a> <span class="kw">let </span>token = CancellationToken::new();
<a href=#624 id=624 data-nosnippet>624</a> <span class="kw">let </span>managed = stream_to_managed_parquet(
<a href=#625 id=625 data-nosnippet>625</a> batches.as_stream(),
<a href=#626 id=626 data-nosnippet>626</a> store.clone(),
<a href=#627 id=627 data-nosnippet>627</a> <span class="string">"managed.parquet"</span>,
<a href=#628 id=628 data-nosnippet>628</a> <span class="kw-2">&amp;</span>budget,
<a href=#629 id=629 data-nosnippet>629</a> <span class="kw-2">&amp;</span>token,
<a href=#630 id=630 data-nosnippet>630</a> )
<a href=#631 id=631 data-nosnippet>631</a> .<span class="kw">await
<a href=#632 id=632 data-nosnippet>632</a> </span>.unwrap();
<a href=#633 id=633 data-nosnippet>633</a> <span class="kw">let </span>output = Output::new_with_record_batches(batches)
<a href=#634 id=634 data-nosnippet>634</a> .map_dictionary_to_values()
<a href=#635 id=635 data-nosnippet>635</a> .unwrap();
<a href=#636 id=636 data-nosnippet>636</a> <span class="kw">let </span>OutputData::RecordBatches(batches) = output.data <span class="kw">else </span>{
<a href=#637 id=637 data-nosnippet>637</a> <span class="macro">panic!</span>(<span class="string">"expected batches"</span>);
<a href=#638 id=638 data-nosnippet>638</a> };
<a href=#639 id=639 data-nosnippet>639</a> <span class="kw">let </span>stream = SendableRecordBatchMapper::new(
<a href=#640 id=640 data-nosnippet>640</a> batches.as_stream(),
<a href=#641 id=641 data-nosnippet>641</a> map_json_type_to_string,
<a href=#642 id=642 data-nosnippet>642</a> map_json_type_to_string_schema,
<a href=#643 id=643 data-nosnippet>643</a> );
<a href=#644 id=644 data-nosnippet>644</a> <span class="kw">let </span>ordinary = stream_to_parquet(
<a href=#645 id=645 data-nosnippet>645</a> Box::pin(DfRecordBatchStreamAdapter::new(Box::pin(stream))),
<a href=#646 id=646 data-nosnippet>646</a> store.clone(),
<a href=#647 id=647 data-nosnippet>647</a> <span class="string">"ordinary.parquet"</span>,
<a href=#648 id=648 data-nosnippet>648</a> WRITE_CONCURRENCY,
<a href=#649 id=649 data-nosnippet>649</a> )
<a href=#650 id=650 data-nosnippet>650</a> .<span class="kw">await
<a href=#651 id=651 data-nosnippet>651</a> </span>.unwrap();
<a href=#652 id=652 data-nosnippet>652</a> <span class="macro">assert_eq!</span>(managed, ordinary);
<a href=#653 id=653 data-nosnippet>653</a> <span class="kw">let </span><span class="kw-2">mut </span>outputs = Vec::new();
<a href=#654 id=654 data-nosnippet>654</a> <span class="kw">for </span>name <span class="kw">in </span>[<span class="string">"managed.parquet"</span>, <span class="string">"ordinary.parquet"</span>] {
<a href=#655 id=655 data-nosnippet>655</a> <span class="kw">let </span>reader = ParquetRecordBatchReaderBuilder::try_new(
<a href=#656 id=656 data-nosnippet>656</a> store.read(name).<span class="kw">await</span>.unwrap().to_bytes(),
<a href=#657 id=657 data-nosnippet>657</a> )
<a href=#658 id=658 data-nosnippet>658</a> .unwrap();
<a href=#659 id=659 data-nosnippet>659</a> <span class="kw">let </span>schema = reader.schema().clone();
<a href=#660 id=660 data-nosnippet>660</a> <span class="kw">let </span>values = reader
<a href=#661 id=661 data-nosnippet>661</a> .build()
<a href=#662 id=662 data-nosnippet>662</a> .unwrap()
<a href=#663 id=663 data-nosnippet>663</a> .collect::&lt;std::result::Result&lt;Vec&lt;<span class="kw">_</span>&gt;, <span class="kw">_</span>&gt;&gt;()
<a href=#664 id=664 data-nosnippet>664</a> .unwrap();
<a href=#665 id=665 data-nosnippet>665</a> outputs.push((schema, values));
<a href=#666 id=666 data-nosnippet>666</a> }
<a href=#667 id=667 data-nosnippet>667</a> <span class="macro">assert_eq!</span>(outputs[<span class="number">0</span>], outputs[<span class="number">1</span>]);
<a href=#668 id=668 data-nosnippet>668</a> <span class="macro">assert_eq!</span>(budget.available(), (<span class="number">1</span>, <span class="number">64 </span>* <span class="number">1024 </span>* <span class="number">1024</span>));
<a href=#669 id=669 data-nosnippet>669</a> }
<a href=#670 id=670 data-nosnippet>670</a> }
<a href=#671 id=671 data-nosnippet>671</a>}
</code></pre></div></section></main></body></html>
+99
View File
@@ -0,0 +1,99 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/cursor.rs`."><title>cursor.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>cursor.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-2"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>common_query::{Output, OutputData};
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>common_recordbatch::RecordBatches;
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_recordbatch::cursor::RecordBatchStreamCursor;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_telemetry::tracing;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>query::parser::QueryStatement;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>snafu::ResultExt;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>sql::statements::cursor::{CloseCursor, DeclareCursor, FetchCursor};
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>sql::statements::statement::Statement;
<a href=#24 id=24 data-nosnippet>24</a>
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, <span class="prelude-ty">Result</span>};
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#27 id=27 data-nosnippet>27</a>
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">impl </span>StatementExecutor {
<a href=#29 id=29 data-nosnippet>29</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#30 id=30 data-nosnippet>30</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>declare_cursor(
<a href=#31 id=31 data-nosnippet>31</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#32 id=32 data-nosnippet>32</a> declare_cursor: DeclareCursor,
<a href=#33 id=33 data-nosnippet>33</a> query_ctx: QueryContextRef,
<a href=#34 id=34 data-nosnippet>34</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#35 id=35 data-nosnippet>35</a> <span class="kw">let </span>cursor_name = declare_cursor.cursor_name.to_string();
<a href=#36 id=36 data-nosnippet>36</a>
<a href=#37 id=37 data-nosnippet>37</a> <span class="kw">if </span>query_ctx.get_cursor(<span class="kw-2">&amp;</span>cursor_name).is_some() {
<a href=#38 id=38 data-nosnippet>38</a> error::CursorExistsSnafu {
<a href=#39 id=39 data-nosnippet>39</a> name: cursor_name.clone(),
<a href=#40 id=40 data-nosnippet>40</a> }
<a href=#41 id=41 data-nosnippet>41</a> .fail()<span class="question-mark">?</span>;
<a href=#42 id=42 data-nosnippet>42</a> }
<a href=#43 id=43 data-nosnippet>43</a>
<a href=#44 id=44 data-nosnippet>44</a> <span class="kw">let </span>query_stmt = Statement::Query(declare_cursor.query);
<a href=#45 id=45 data-nosnippet>45</a>
<a href=#46 id=46 data-nosnippet>46</a> <span class="kw">let </span>output = <span class="self">self
<a href=#47 id=47 data-nosnippet>47</a> </span>.plan_exec(QueryStatement::Sql(query_stmt), query_ctx.clone())
<a href=#48 id=48 data-nosnippet>48</a> .<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#49 id=49 data-nosnippet>49</a> <span class="kw">match </span>output.data {
<a href=#50 id=50 data-nosnippet>50</a> OutputData::RecordBatches(rb) =&gt; {
<a href=#51 id=51 data-nosnippet>51</a> <span class="kw">let </span>rbs = rb.as_stream();
<a href=#52 id=52 data-nosnippet>52</a> query_ctx.insert_cursor(cursor_name, RecordBatchStreamCursor::new(rbs));
<a href=#53 id=53 data-nosnippet>53</a> }
<a href=#54 id=54 data-nosnippet>54</a> OutputData::Stream(rbs) =&gt; {
<a href=#55 id=55 data-nosnippet>55</a> query_ctx.insert_cursor(cursor_name, RecordBatchStreamCursor::new(rbs));
<a href=#56 id=56 data-nosnippet>56</a> }
<a href=#57 id=57 data-nosnippet>57</a> <span class="comment">// Should not happen because we have query type ensured from parser.
<a href=#58 id=58 data-nosnippet>58</a> </span>OutputData::AffectedRows(<span class="kw">_</span>) =&gt; error::NotSupportedSnafu {
<a href=#59 id=59 data-nosnippet>59</a> feat: <span class="string">"Non-query statement on cursor"</span>,
<a href=#60 id=60 data-nosnippet>60</a> }
<a href=#61 id=61 data-nosnippet>61</a> .fail()<span class="question-mark">?</span>,
<a href=#62 id=62 data-nosnippet>62</a> }
<a href=#63 id=63 data-nosnippet>63</a>
<a href=#64 id=64 data-nosnippet>64</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="number">0</span>))
<a href=#65 id=65 data-nosnippet>65</a> }
<a href=#66 id=66 data-nosnippet>66</a>
<a href=#67 id=67 data-nosnippet>67</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#68 id=68 data-nosnippet>68</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>fetch_cursor(
<a href=#69 id=69 data-nosnippet>69</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#70 id=70 data-nosnippet>70</a> fetch_cursor: FetchCursor,
<a href=#71 id=71 data-nosnippet>71</a> query_ctx: QueryContextRef,
<a href=#72 id=72 data-nosnippet>72</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#73 id=73 data-nosnippet>73</a> <span class="kw">let </span>cursor_name = fetch_cursor.cursor_name.to_string();
<a href=#74 id=74 data-nosnippet>74</a> <span class="kw">let </span>fetch_size = fetch_cursor.fetch_size;
<a href=#75 id=75 data-nosnippet>75</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(rb) = query_ctx.get_cursor(<span class="kw-2">&amp;</span>cursor_name) {
<a href=#76 id=76 data-nosnippet>76</a> <span class="kw">let </span>record_batch = rb
<a href=#77 id=77 data-nosnippet>77</a> .take(fetch_size <span class="kw">as </span>usize)
<a href=#78 id=78 data-nosnippet>78</a> .<span class="kw">await
<a href=#79 id=79 data-nosnippet>79</a> </span>.context(error::BuildRecordBatchSnafu)<span class="question-mark">?</span>;
<a href=#80 id=80 data-nosnippet>80</a> <span class="kw">let </span>record_batches =
<a href=#81 id=81 data-nosnippet>81</a> RecordBatches::try_new(record_batch.schema.clone(), <span class="macro">vec!</span>[record_batch])
<a href=#82 id=82 data-nosnippet>82</a> .context(error::BuildRecordBatchSnafu)<span class="question-mark">?</span>;
<a href=#83 id=83 data-nosnippet>83</a> <span class="prelude-val">Ok</span>(Output::new_with_record_batches(record_batches))
<a href=#84 id=84 data-nosnippet>84</a> } <span class="kw">else </span>{
<a href=#85 id=85 data-nosnippet>85</a> error::CursorNotFoundSnafu { name: cursor_name }.fail()
<a href=#86 id=86 data-nosnippet>86</a> }
<a href=#87 id=87 data-nosnippet>87</a> }
<a href=#88 id=88 data-nosnippet>88</a>
<a href=#89 id=89 data-nosnippet>89</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#90 id=90 data-nosnippet>90</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>close_cursor(
<a href=#91 id=91 data-nosnippet>91</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#92 id=92 data-nosnippet>92</a> close_cursor: CloseCursor,
<a href=#93 id=93 data-nosnippet>93</a> query_ctx: QueryContextRef,
<a href=#94 id=94 data-nosnippet>94</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#95 id=95 data-nosnippet>95</a> query_ctx.remove_cursor(<span class="kw-2">&amp;</span>close_cursor.cursor_name.to_string());
<a href=#96 id=96 data-nosnippet>96</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="number">0</span>))
<a href=#97 id=97 data-nosnippet>97</a> }
<a href=#98 id=98 data-nosnippet>98</a>}
</code></pre></div></section></main></body></html>
@@ -0,0 +1,333 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/database_copy.rs`."><title>database_copy.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>database_copy.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="doccomment">//! Shared selection and destination rules for database COPY and prepared exports.
<a href=#16 id=16 data-nosnippet>16</a>
<a href=#17 id=17 data-nosnippet>17</a></span><span class="kw">use </span>std::collections::HashMap;
<a href=#18 id=18 data-nosnippet>18</a>
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>common_datasource::object_store::FILE_SCHEMA;
<a href=#20 id=20 data-nosnippet>20</a><span class="attr">#[cfg(windows)]
<a href=#21 id=21 data-nosnippet>21</a></span><span class="kw">use </span>common_datasource::object_store::{FS_SCHEMA, parse_url};
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>common_stat::get_total_cpu_cores;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>snafu::{OptionExt, ResultExt, ensure};
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME};
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>table::TableRef;
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>table::metadata::TableType;
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>table::requests::CopyDatabaseRequest;
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span>table::table_reference::TableReference;
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span>tokio::sync::Semaphore;
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span>url::Url;
<a href=#32 id=32 data-nosnippet>32</a>
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, <span class="prelude-ty">Result</span>};
<a href=#34 id=34 data-nosnippet>34</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#35 id=35 data-nosnippet>35</a>
<a href=#36 id=36 data-nosnippet>36</a><span class="kw">fn </span>is_directory_location(location: <span class="kw-2">&amp;</span>str) -&gt; bool {
<a href=#37 id=37 data-nosnippet>37</a> <span class="kw">if </span>location.ends_with(<span class="string">'/'</span>) {
<a href=#38 id=38 data-nosnippet>38</a> <span class="kw">return </span><span class="bool-val">true</span>;
<a href=#39 id=39 data-nosnippet>39</a> }
<a href=#40 id=40 data-nosnippet>40</a>
<a href=#41 id=41 data-nosnippet>41</a> <span class="attr">#[cfg(windows)]
<a href=#42 id=42 data-nosnippet>42</a> </span>{
<a href=#43 id=43 data-nosnippet>43</a> location.ends_with(std::path::MAIN_SEPARATOR)
<a href=#44 id=44 data-nosnippet>44</a> &amp;&amp; <span class="macro">matches!</span>(
<a href=#45 id=45 data-nosnippet>45</a> parse_url(location),
<a href=#46 id=46 data-nosnippet>46</a> <span class="prelude-val">Ok</span>((schema, <span class="kw">_</span>, <span class="kw">_</span>)) <span class="kw">if </span>schema.eq_ignore_ascii_case(FS_SCHEMA)
<a href=#47 id=47 data-nosnippet>47</a> )
<a href=#48 id=48 data-nosnippet>48</a> }
<a href=#49 id=49 data-nosnippet>49</a>
<a href=#50 id=50 data-nosnippet>50</a> <span class="attr">#[cfg(not(windows))]
<a href=#51 id=51 data-nosnippet>51</a> </span><span class="bool-val">false
<a href=#52 id=52 data-nosnippet>52</a></span>}
<a href=#53 id=53 data-nosnippet>53</a>
<a href=#54 id=54 data-nosnippet>54</a><span class="doccomment">/// Get parallelism from options, default to total CPU cores.
<a href=#55 id=55 data-nosnippet>55</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>parse_parallelism_from_option_map(options: <span class="kw-2">&amp;</span>HashMap&lt;String, String&gt;) -&gt; usize {
<a href=#56 id=56 data-nosnippet>56</a> options
<a href=#57 id=57 data-nosnippet>57</a> .get(<span class="string">"parallelism"</span>)
<a href=#58 id=58 data-nosnippet>58</a> .and_then(|v| v.parse::&lt;usize&gt;().ok())
<a href=#59 id=59 data-nosnippet>59</a> .unwrap_or_else(get_total_cpu_cores)
<a href=#60 id=60 data-nosnippet>60</a> .clamp(<span class="number">1</span>, Semaphore::MAX_PERMITS)
<a href=#61 id=61 data-nosnippet>61</a>}
<a href=#62 id=62 data-nosnippet>62</a>
<a href=#63 id=63 data-nosnippet>63</a><span class="doccomment">/// Rejects layouts unsupported by the per-table export path.
<a href=#64 id=64 data-nosnippet>64</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>validate_database_export_layout(options: <span class="kw-2">&amp;</span>HashMap&lt;String, String&gt;) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#65 id=65 data-nosnippet>65</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(layout) = options.get(<span class="string">"metric_data_layout"</span>) {
<a href=#66 id=66 data-nosnippet>66</a> <span class="kw">return </span>error::InvalidCopyParameterSnafu {
<a href=#67 id=67 data-nosnippet>67</a> key: <span class="string">"metric_data_layout"</span>,
<a href=#68 id=68 data-nosnippet>68</a> value: layout,
<a href=#69 id=69 data-nosnippet>69</a> }
<a href=#70 id=70 data-nosnippet>70</a> .fail();
<a href=#71 id=71 data-nosnippet>71</a> }
<a href=#72 id=72 data-nosnippet>72</a> <span class="prelude-val">Ok</span>(())
<a href=#73 id=73 data-nosnippet>73</a>}
<a href=#74 id=74 data-nosnippet>74</a>
<a href=#75 id=75 data-nosnippet>75</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>validate_database_directory(location: <span class="kw-2">&amp;</span>str) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#76 id=76 data-nosnippet>76</a> <span class="macro">ensure!</span>(
<a href=#77 id=77 data-nosnippet>77</a> is_directory_location(location),
<a href=#78 id=78 data-nosnippet>78</a> error::InvalidCopyDatabasePathSnafu { value: location }
<a href=#79 id=79 data-nosnippet>79</a> );
<a href=#80 id=80 data-nosnippet>80</a> <span class="attr">#[cfg(windows)]
<a href=#81 id=81 data-nosnippet>81</a> </span><span class="kw">if </span>common_datasource::object_store::handle_windows_path(location).is_some() {
<a href=#82 id=82 data-nosnippet>82</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(());
<a href=#83 id=83 data-nosnippet>83</a> }
<a href=#84 id=84 data-nosnippet>84</a> <span class="kw">let </span>parsed_directory = <span class="kw">match </span>Url::parse(location) {
<a href=#85 id=85 data-nosnippet>85</a> <span class="prelude-val">Ok</span>(url) =&gt; {
<a href=#86 id=86 data-nosnippet>86</a> url.query().is_none() &amp;&amp; url.fragment().is_none() &amp;&amp; is_directory_location(url.path())
<a href=#87 id=87 data-nosnippet>87</a> }
<a href=#88 id=88 data-nosnippet>88</a> <span class="prelude-val">Err</span>(<span class="kw">_</span>) =&gt; <span class="bool-val">true</span>,
<a href=#89 id=89 data-nosnippet>89</a> };
<a href=#90 id=90 data-nosnippet>90</a> <span class="macro">ensure!</span>(
<a href=#91 id=91 data-nosnippet>91</a> parsed_directory,
<a href=#92 id=92 data-nosnippet>92</a> error::InvalidCopyDatabasePathSnafu { value: location }
<a href=#93 id=93 data-nosnippet>93</a> );
<a href=#94 id=94 data-nosnippet>94</a> <span class="prelude-val">Ok</span>(())
<a href=#95 id=95 data-nosnippet>95</a>}
<a href=#96 id=96 data-nosnippet>96</a>
<a href=#97 id=97 data-nosnippet>97</a><span class="doccomment">/// The writer key and its externally reported location, resolved together.
<a href=#98 id=98 data-nosnippet>98</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>DatabaseExportFile {
<a href=#99 id=99 data-nosnippet>99</a> <span class="kw">pub</span>(<span class="kw">crate</span>) path: String,
<a href=#100 id=100 data-nosnippet>100</a> <span class="kw">pub</span>(<span class="kw">crate</span>) location: String,
<a href=#101 id=101 data-nosnippet>101</a>}
<a href=#102 id=102 data-nosnippet>102</a>
<a href=#103 id=103 data-nosnippet>103</a><span class="kw">impl </span>DatabaseExportFile {
<a href=#104 id=104 data-nosnippet>104</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>new(directory: <span class="kw-2">&amp;</span>str, name: <span class="kw-2">&amp;</span>str, suffix: <span class="kw-2">&amp;</span>str) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="self">Self</span>&gt; {
<a href=#105 id=105 data-nosnippet>105</a> <span class="kw">let </span>filename = <span class="macro">format!</span>(<span class="string">"{name}{suffix}"</span>);
<a href=#106 id=106 data-nosnippet>106</a> <span class="attr">#[cfg(windows)]
<a href=#107 id=107 data-nosnippet>107</a> </span><span class="kw">if </span>common_datasource::object_store::handle_windows_path(directory).is_some() {
<a href=#108 id=108 data-nosnippet>108</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(<span class="self">Self </span>{
<a href=#109 id=109 data-nosnippet>109</a> location: <span class="macro">format!</span>(<span class="string">"{directory}{filename}"</span>),
<a href=#110 id=110 data-nosnippet>110</a> path: filename,
<a href=#111 id=111 data-nosnippet>111</a> });
<a href=#112 id=112 data-nosnippet>112</a> }
<a href=#113 id=113 data-nosnippet>113</a> <span class="kw">match </span>Url::parse(directory) {
<a href=#114 id=114 data-nosnippet>114</a> <span class="prelude-val">Ok</span>(<span class="kw-2">mut </span>url) =&gt; {
<a href=#115 id=115 data-nosnippet>115</a> url.path_segments_mut()
<a href=#116 id=116 data-nosnippet>116</a> .map_err(|<span class="kw">_</span>| error::InvalidCopyDatabasePathSnafu { value: directory }.build())<span class="question-mark">?
<a href=#117 id=117 data-nosnippet>117</a> </span>.pop_if_empty()
<a href=#118 id=118 data-nosnippet>118</a> .push(<span class="kw-2">&amp;</span>filename);
<a href=#119 id=119 data-nosnippet>119</a> <span class="comment">// File URLs are decoded by the filesystem backend; object-store
<a href=#120 id=120 data-nosnippet>120</a> // backends use the encoded URL path as their key.
<a href=#121 id=121 data-nosnippet>121</a> </span><span class="kw">let </span>path = <span class="kw">if </span>url.scheme().eq_ignore_ascii_case(FILE_SCHEMA) {
<a href=#122 id=122 data-nosnippet>122</a> filename
<a href=#123 id=123 data-nosnippet>123</a> } <span class="kw">else </span>{
<a href=#124 id=124 data-nosnippet>124</a> url.path()
<a href=#125 id=125 data-nosnippet>125</a> .rsplit(<span class="string">'/'</span>)
<a href=#126 id=126 data-nosnippet>126</a> .next()
<a href=#127 id=127 data-nosnippet>127</a> .unwrap_or_default()
<a href=#128 id=128 data-nosnippet>128</a> .to_string()
<a href=#129 id=129 data-nosnippet>129</a> };
<a href=#130 id=130 data-nosnippet>130</a> <span class="prelude-val">Ok</span>(<span class="self">Self </span>{
<a href=#131 id=131 data-nosnippet>131</a> path,
<a href=#132 id=132 data-nosnippet>132</a> location: url.into(),
<a href=#133 id=133 data-nosnippet>133</a> })
<a href=#134 id=134 data-nosnippet>134</a> }
<a href=#135 id=135 data-nosnippet>135</a> <span class="prelude-val">Err</span>(url::ParseError::RelativeUrlWithoutBase) =&gt; <span class="prelude-val">Ok</span>(<span class="self">Self </span>{
<a href=#136 id=136 data-nosnippet>136</a> location: <span class="macro">format!</span>(<span class="string">"{directory}{filename}"</span>),
<a href=#137 id=137 data-nosnippet>137</a> path: filename,
<a href=#138 id=138 data-nosnippet>138</a> }),
<a href=#139 id=139 data-nosnippet>139</a> <span class="prelude-val">Err</span>(source) =&gt; <span class="prelude-val">Err</span>(source)
<a href=#140 id=140 data-nosnippet>140</a> .context(common_datasource::error::InvalidUrlSnafu { url: directory })
<a href=#141 id=141 data-nosnippet>141</a> .context(error::BuildBackendSnafu),
<a href=#142 id=142 data-nosnippet>142</a> }
<a href=#143 id=143 data-nosnippet>143</a> }
<a href=#144 id=144 data-nosnippet>144</a>}
<a href=#145 id=145 data-nosnippet>145</a>
<a href=#146 id=146 data-nosnippet>146</a><span class="doccomment">/// Resolve a listed writer key back to its table name and COPY input location.
<a href=#147 id=147 data-nosnippet>147</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>database_import_source(directory: <span class="kw-2">&amp;</span>str, path: <span class="kw-2">&amp;</span>str) -&gt; <span class="prelude-ty">Result</span>&lt;(String, String)&gt; {
<a href=#148 id=148 data-nosnippet>148</a> <span class="kw">let </span><span class="kw-2">mut </span>filename = path.rsplit(<span class="string">'/'</span>).next().unwrap_or(path).to_string();
<a href=#149 id=149 data-nosnippet>149</a> <span class="kw">let </span><span class="kw-2">mut </span>location = <span class="macro">format!</span>(<span class="string">"{directory}{path}"</span>);
<a href=#150 id=150 data-nosnippet>150</a> <span class="attr">#[cfg(windows)]
<a href=#151 id=151 data-nosnippet>151</a> </span><span class="kw">let </span>literal_path = common_datasource::object_store::handle_windows_path(directory).is_some();
<a href=#152 id=152 data-nosnippet>152</a> <span class="attr">#[cfg(not(windows))]
<a href=#153 id=153 data-nosnippet>153</a> </span><span class="kw">let </span>literal_path = <span class="bool-val">false</span>;
<a href=#154 id=154 data-nosnippet>154</a> <span class="kw">if </span>!literal_path &amp;&amp; <span class="kw">let </span><span class="prelude-val">Ok</span>(<span class="kw-2">mut </span>url) = Url::parse(directory) {
<a href=#155 id=155 data-nosnippet>155</a> <span class="kw">if </span>url.scheme().eq_ignore_ascii_case(FILE_SCHEMA) {
<a href=#156 id=156 data-nosnippet>156</a> url.path_segments_mut()
<a href=#157 id=157 data-nosnippet>157</a> .map_err(|<span class="kw">_</span>| error::InvalidCopyDatabasePathSnafu { value: directory }.build())<span class="question-mark">?
<a href=#158 id=158 data-nosnippet>158</a> </span>.pop_if_empty()
<a href=#159 id=159 data-nosnippet>159</a> .extend(path.split(<span class="string">'/'</span>));
<a href=#160 id=160 data-nosnippet>160</a> } <span class="kw">else </span>{
<a href=#161 id=161 data-nosnippet>161</a> <span class="comment">// Listed object keys already contain the export URL's escaping.
<a href=#162 id=162 data-nosnippet>162</a> </span>url.set_path(<span class="kw-2">&amp;</span><span class="macro">format!</span>(<span class="string">"{}{path}"</span>, url.path()));
<a href=#163 id=163 data-nosnippet>163</a> filename = percent_encoding::percent_decode_str(<span class="kw-2">&amp;</span>filename)
<a href=#164 id=164 data-nosnippet>164</a> .decode_utf8()
<a href=#165 id=165 data-nosnippet>165</a> .ok()
<a href=#166 id=166 data-nosnippet>166</a> .context(error::InvalidTableNameSnafu { table_name: path })<span class="question-mark">?
<a href=#167 id=167 data-nosnippet>167</a> </span>.into_owned();
<a href=#168 id=168 data-nosnippet>168</a> }
<a href=#169 id=169 data-nosnippet>169</a> location = url.into();
<a href=#170 id=170 data-nosnippet>170</a> }
<a href=#171 id=171 data-nosnippet>171</a> <span class="kw">let </span>table_name = filename
<a href=#172 id=172 data-nosnippet>172</a> .rsplit_once(<span class="string">'.'</span>)
<a href=#173 id=173 data-nosnippet>173</a> .map(|(stem, <span class="kw">_</span>)| stem)
<a href=#174 id=174 data-nosnippet>174</a> .filter(|stem| !stem.is_empty())
<a href=#175 id=175 data-nosnippet>175</a> .context(error::InvalidTableNameSnafu { table_name: path })<span class="question-mark">?
<a href=#176 id=176 data-nosnippet>176</a> </span>.to_string();
<a href=#177 id=177 data-nosnippet>177</a> <span class="prelude-val">Ok</span>((table_name, location))
<a href=#178 id=178 data-nosnippet>178</a>}
<a href=#179 id=179 data-nosnippet>179</a>
<a href=#180 id=180 data-nosnippet>180</a><span class="kw">impl </span>StatementExecutor {
<a href=#181 id=181 data-nosnippet>181</a> <span class="doccomment">/// Capture each selected data table once. Views, temporary and Metric physical
<a href=#182 id=182 data-nosnippet>182</a> /// tables do not have data outputs. `None` selects the whole schema.
<a href=#183 id=183 data-nosnippet>183</a> </span><span class="kw">pub async fn </span>capture_database_export_tables(
<a href=#184 id=184 data-nosnippet>184</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#185 id=185 data-nosnippet>185</a> req: <span class="kw-2">&amp;</span>CopyDatabaseRequest,
<a href=#186 id=186 data-nosnippet>186</a> names: <span class="prelude-ty">Option</span>&lt;<span class="kw-2">&amp;</span>[String]&gt;,
<a href=#187 id=187 data-nosnippet>187</a> ctx: <span class="kw-2">&amp;</span>QueryContextRef,
<a href=#188 id=188 data-nosnippet>188</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Vec&lt;TableRef&gt;&gt; {
<a href=#189 id=189 data-nosnippet>189</a> <span class="kw">let </span><span class="kw-2">mut </span>names = <span class="kw">match </span>names {
<a href=#190 id=190 data-nosnippet>190</a> <span class="prelude-val">Some</span>(names) =&gt; names.to_vec(),
<a href=#191 id=191 data-nosnippet>191</a> <span class="prelude-val">None </span>=&gt; <span class="self">self
<a href=#192 id=192 data-nosnippet>192</a> </span>.catalog_manager
<a href=#193 id=193 data-nosnippet>193</a> .table_names(<span class="kw-2">&amp;</span>req.catalog_name, <span class="kw-2">&amp;</span>req.schema_name, <span class="prelude-val">Some</span>(ctx))
<a href=#194 id=194 data-nosnippet>194</a> .<span class="kw">await
<a href=#195 id=195 data-nosnippet>195</a> </span>.context(error::CatalogSnafu)<span class="question-mark">?</span>,
<a href=#196 id=196 data-nosnippet>196</a> };
<a href=#197 id=197 data-nosnippet>197</a> names.sort();
<a href=#198 id=198 data-nosnippet>198</a> <span class="kw">let </span><span class="kw-2">mut </span>tables = Vec::new();
<a href=#199 id=199 data-nosnippet>199</a> <span class="kw">for </span>name <span class="kw">in </span>names {
<a href=#200 id=200 data-nosnippet>200</a> <span class="kw">let </span>table = <span class="self">self
<a href=#201 id=201 data-nosnippet>201</a> </span>.get_table(<span class="kw-2">&amp;</span>TableReference::full(
<a href=#202 id=202 data-nosnippet>202</a> <span class="kw-2">&amp;</span>req.catalog_name,
<a href=#203 id=203 data-nosnippet>203</a> <span class="kw-2">&amp;</span>req.schema_name,
<a href=#204 id=204 data-nosnippet>204</a> <span class="kw-2">&amp;</span>name,
<a href=#205 id=205 data-nosnippet>205</a> ))
<a href=#206 id=206 data-nosnippet>206</a> .<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#207 id=207 data-nosnippet>207</a> <span class="kw">let </span>info = table.table_info();
<a href=#208 id=208 data-nosnippet>208</a> <span class="kw">if </span>table.table_type() == TableType::Base
<a href=#209 id=209 data-nosnippet>209</a> &amp;&amp; (info.meta.engine != METRIC_ENGINE_NAME
<a href=#210 id=210 data-nosnippet>210</a> || info
<a href=#211 id=211 data-nosnippet>211</a> .meta
<a href=#212 id=212 data-nosnippet>212</a> .options
<a href=#213 id=213 data-nosnippet>213</a> .extra_options
<a href=#214 id=214 data-nosnippet>214</a> .contains_key(LOGICAL_TABLE_METADATA_KEY))
<a href=#215 id=215 data-nosnippet>215</a> {
<a href=#216 id=216 data-nosnippet>216</a> tables.push(table);
<a href=#217 id=217 data-nosnippet>217</a> }
<a href=#218 id=218 data-nosnippet>218</a> }
<a href=#219 id=219 data-nosnippet>219</a> <span class="prelude-val">Ok</span>(tables)
<a href=#220 id=220 data-nosnippet>220</a> }
<a href=#221 id=221 data-nosnippet>221</a>}
<a href=#222 id=222 data-nosnippet>222</a>
<a href=#223 id=223 data-nosnippet>223</a><span class="attr">#[cfg(test)]
<a href=#224 id=224 data-nosnippet>224</a></span><span class="kw">mod </span>tests {
<a href=#225 id=225 data-nosnippet>225</a> <span class="kw">use super</span>::<span class="kw-2">*</span>;
<a href=#226 id=226 data-nosnippet>226</a>
<a href=#227 id=227 data-nosnippet>227</a> <span class="attr">#[test]
<a href=#228 id=228 data-nosnippet>228</a> </span><span class="kw">fn </span>directory_url_components_cannot_capture_output_names() {
<a href=#229 id=229 data-nosnippet>229</a> <span class="kw">for </span>location <span class="kw">in </span>[
<a href=#230 id=230 data-nosnippet>230</a> <span class="string">"file:///copy/fresh?attempt=/"</span>,
<a href=#231 id=231 data-nosnippet>231</a> <span class="string">"file:///copy/fresh#attempt/"</span>,
<a href=#232 id=232 data-nosnippet>232</a> <span class="string">"s3://bucket/fresh?attempt=/"</span>,
<a href=#233 id=233 data-nosnippet>233</a> <span class="string">"s3://bucket/fresh#attempt/"</span>,
<a href=#234 id=234 data-nosnippet>234</a> <span class="string">"file:///copy/fresh"</span>,
<a href=#235 id=235 data-nosnippet>235</a> ] {
<a href=#236 id=236 data-nosnippet>236</a> <span class="macro">assert!</span>(validate_database_directory(location).is_err(), <span class="string">"{location}"</span>);
<a href=#237 id=237 data-nosnippet>237</a> }
<a href=#238 id=238 data-nosnippet>238</a> <span class="kw">for </span>location <span class="kw">in </span>[<span class="string">"/copy/fresh/"</span>, <span class="string">"file:///copy/fresh/"</span>, <span class="string">"s3://bucket/fresh/"</span>] {
<a href=#239 id=239 data-nosnippet>239</a> validate_database_directory(location).unwrap();
<a href=#240 id=240 data-nosnippet>240</a> }
<a href=#241 id=241 data-nosnippet>241</a> }
<a href=#242 id=242 data-nosnippet>242</a>
<a href=#243 id=243 data-nosnippet>243</a> <span class="attr">#[cfg(windows)]
<a href=#244 id=244 data-nosnippet>244</a> #[test]
<a href=#245 id=245 data-nosnippet>245</a> </span><span class="kw">fn </span>windows_directory_names_are_literal_paths() {
<a href=#246 id=246 data-nosnippet>246</a> <span class="kw">for </span>location <span class="kw">in </span>[<span class="string">"C:/copy/fresh#1/"</span>, <span class="string">r"C:\copy\fresh#1\"</span>] {
<a href=#247 id=247 data-nosnippet>247</a> validate_database_directory(location).unwrap();
<a href=#248 id=248 data-nosnippet>248</a> }
<a href=#249 id=249 data-nosnippet>249</a> <span class="macro">assert!</span>(validate_database_directory(<span class="string">"C:/copy/fresh#1"</span>).is_err());
<a href=#250 id=250 data-nosnippet>250</a> }
<a href=#251 id=251 data-nosnippet>251</a>
<a href=#252 id=252 data-nosnippet>252</a> <span class="attr">#[tokio::test]
<a href=#253 id=253 data-nosnippet>253</a> </span><span class="kw">async fn </span>output_locations_resolve_to_writer_keys() {
<a href=#254 id=254 data-nosnippet>254</a> <span class="kw">use </span>common_datasource::object_store::{LocalFileAccess, build_backend_for_write_with_path};
<a href=#255 id=255 data-nosnippet>255</a>
<a href=#256 id=256 data-nosnippet>256</a> <span class="kw">let </span>dir = common_test_util::temp_dir::create_temp_dir(<span class="string">"database_export_paths"</span>);
<a href=#257 id=257 data-nosnippet>257</a> <span class="kw">let </span>access = LocalFileAccess::sandboxed(dir.path()).unwrap();
<a href=#258 id=258 data-nosnippet>258</a> <span class="kw">let </span>file_url = Url::from_directory_path(dir.path()).unwrap().to_string();
<a href=#259 id=259 data-nosnippet>259</a> <span class="kw">let </span>connection = HashMap::from([
<a href=#260 id=260 data-nosnippet>260</a> (<span class="string">"region"</span>.into(), <span class="string">"us-east-1"</span>.into()),
<a href=#261 id=261 data-nosnippet>261</a> (<span class="string">"access_key_id"</span>.into(), <span class="string">"test-key"</span>.into()),
<a href=#262 id=262 data-nosnippet>262</a> (<span class="string">"secret_access_key"</span>.into(), <span class="string">"test-secret"</span>.into()),
<a href=#263 id=263 data-nosnippet>263</a> ]);
<a href=#264 id=264 data-nosnippet>264</a> <span class="kw">for </span>directory <span class="kw">in </span>[
<a href=#265 id=265 data-nosnippet>265</a> <span class="macro">format!</span>(<span class="string">"{}/"</span>, dir.path().display()),
<a href=#266 id=266 data-nosnippet>266</a> file_url,
<a href=#267 id=267 data-nosnippet>267</a> <span class="string">"s3://export-bucket/data/"</span>.into(),
<a href=#268 id=268 data-nosnippet>268</a> ] {
<a href=#269 id=269 data-nosnippet>269</a> <span class="kw">for </span>name <span class="kw">in </span>[<span class="string">"a#b"</span>, <span class="string">"a:b"</span>] {
<a href=#270 id=270 data-nosnippet>270</a> <span class="kw">if </span><span class="macro">cfg!</span>(windows) &amp;&amp; name.contains(<span class="string">':'</span>) &amp;&amp; !directory.starts_with(<span class="string">"s3:"</span>) {
<a href=#271 id=271 data-nosnippet>271</a> <span class="kw">continue</span>;
<a href=#272 id=272 data-nosnippet>272</a> }
<a href=#273 id=273 data-nosnippet>273</a> <span class="kw">let </span>file = DatabaseExportFile::new(<span class="kw-2">&amp;</span>directory, name, <span class="string">".parquet"</span>).unwrap();
<a href=#274 id=274 data-nosnippet>274</a> <span class="kw">let </span>backend =
<a href=#275 id=275 data-nosnippet>275</a> build_backend_for_write_with_path(<span class="kw-2">&amp;</span>file.location, <span class="kw-2">&amp;</span>connection, <span class="kw-2">&amp;</span>access)
<a href=#276 id=276 data-nosnippet>276</a> .<span class="kw">await
<a href=#277 id=277 data-nosnippet>277</a> </span>.unwrap();
<a href=#278 id=278 data-nosnippet>278</a> <span class="macro">assert_eq!</span>(backend.object_path.as_deref(), <span class="prelude-val">Some</span>(file.path.as_str()));
<a href=#279 id=279 data-nosnippet>279</a> <span class="kw">let </span>(table_name, input_location) =
<a href=#280 id=280 data-nosnippet>280</a> database_import_source(<span class="kw-2">&amp;</span>directory, <span class="kw-2">&amp;</span>file.path).unwrap();
<a href=#281 id=281 data-nosnippet>281</a> <span class="macro">assert_eq!</span>(table_name, name);
<a href=#282 id=282 data-nosnippet>282</a> <span class="macro">assert_eq!</span>(input_location, file.location);
<a href=#283 id=283 data-nosnippet>283</a> <span class="kw">let </span>(table_name, nested_location) =
<a href=#284 id=284 data-nosnippet>284</a> database_import_source(<span class="kw-2">&amp;</span>directory, <span class="kw-2">&amp;</span><span class="macro">format!</span>(<span class="string">"nested/{}"</span>, file.path)).unwrap();
<a href=#285 id=285 data-nosnippet>285</a> <span class="macro">assert_eq!</span>(table_name, name);
<a href=#286 id=286 data-nosnippet>286</a> <span class="macro">assert_eq!</span>(
<a href=#287 id=287 data-nosnippet>287</a> nested_location,
<a href=#288 id=288 data-nosnippet>288</a> file.location
<a href=#289 id=289 data-nosnippet>289</a> .replace(<span class="kw-2">&amp;</span>directory, <span class="kw-2">&amp;</span><span class="macro">format!</span>(<span class="string">"{directory}nested/"</span>))
<a href=#290 id=290 data-nosnippet>290</a> );
<a href=#291 id=291 data-nosnippet>291</a> <span class="kw">if </span>!directory.starts_with(<span class="string">"s3:"</span>) {
<a href=#292 id=292 data-nosnippet>292</a> backend
<a href=#293 id=293 data-nosnippet>293</a> .object_store
<a href=#294 id=294 data-nosnippet>294</a> .write(<span class="kw-2">&amp;</span>file.path, <span class="string">"test"</span>)
<a href=#295 id=295 data-nosnippet>295</a> .<span class="kw">await
<a href=#296 id=296 data-nosnippet>296</a> </span>.unwrap();
<a href=#297 id=297 data-nosnippet>297</a> <span class="macro">assert_eq!</span>(
<a href=#298 id=298 data-nosnippet>298</a> std::fs::read(dir.path().join(<span class="macro">format!</span>(<span class="string">"{name}.parquet"</span>))).unwrap(),
<a href=#299 id=299 data-nosnippet>299</a> <span class="string">b"test"
<a href=#300 id=300 data-nosnippet>300</a> </span>);
<a href=#301 id=301 data-nosnippet>301</a> } <span class="kw">else </span>{
<a href=#302 id=302 data-nosnippet>302</a> <span class="macro">assert_eq!</span>(
<a href=#303 id=303 data-nosnippet>303</a> file.path,
<a href=#304 id=304 data-nosnippet>304</a> <span class="kw">if </span>name == <span class="string">"a#b" </span>{
<a href=#305 id=305 data-nosnippet>305</a> <span class="string">"a%23b.parquet"
<a href=#306 id=306 data-nosnippet>306</a> </span>} <span class="kw">else </span>{
<a href=#307 id=307 data-nosnippet>307</a> <span class="string">"a:b.parquet"
<a href=#308 id=308 data-nosnippet>308</a> </span>}
<a href=#309 id=309 data-nosnippet>309</a> );
<a href=#310 id=310 data-nosnippet>310</a> }
<a href=#311 id=311 data-nosnippet>311</a> }
<a href=#312 id=312 data-nosnippet>312</a> }
<a href=#313 id=313 data-nosnippet>313</a> }
<a href=#314 id=314 data-nosnippet>314</a>
<a href=#315 id=315 data-nosnippet>315</a> <span class="attr">#[test]
<a href=#316 id=316 data-nosnippet>316</a> </span><span class="kw">fn </span>test_parse_parallelism_from_option_map() {
<a href=#317 id=317 data-nosnippet>317</a> <span class="kw">let </span>options = HashMap::new();
<a href=#318 id=318 data-nosnippet>318</a> <span class="macro">assert_eq!</span>(
<a href=#319 id=319 data-nosnippet>319</a> parse_parallelism_from_option_map(<span class="kw-2">&amp;</span>options),
<a href=#320 id=320 data-nosnippet>320</a> get_total_cpu_cores()
<a href=#321 id=321 data-nosnippet>321</a> );
<a href=#322 id=322 data-nosnippet>322</a>
<a href=#323 id=323 data-nosnippet>323</a> <span class="kw">let </span>options = HashMap::from([(<span class="string">"parallelism"</span>.to_string(), <span class="string">"0"</span>.to_string())]);
<a href=#324 id=324 data-nosnippet>324</a> <span class="macro">assert_eq!</span>(parse_parallelism_from_option_map(<span class="kw-2">&amp;</span>options), <span class="number">1</span>);
<a href=#325 id=325 data-nosnippet>325</a>
<a href=#326 id=326 data-nosnippet>326</a> <span class="kw">let </span>options = HashMap::from([(<span class="string">"parallelism"</span>.to_string(), usize::MAX.to_string())]);
<a href=#327 id=327 data-nosnippet>327</a> <span class="macro">assert_eq!</span>(
<a href=#328 id=328 data-nosnippet>328</a> parse_parallelism_from_option_map(<span class="kw-2">&amp;</span>options),
<a href=#329 id=329 data-nosnippet>329</a> Semaphore::MAX_PERMITS
<a href=#330 id=330 data-nosnippet>330</a> );
<a href=#331 id=331 data-nosnippet>331</a> }
<a href=#332 id=332 data-nosnippet>332</a>}
</code></pre></div></section></main></body></html>
File diff suppressed because it is too large Load Diff
+52
View File
@@ -0,0 +1,52 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/describe.rs`."><title>describe.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>describe.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-2"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>common_error::ext::BoxedError;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>common_query::Output;
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_telemetry::tracing;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>session::table_name::table_idents_to_full_name;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>snafu::{OptionExt, ResultExt};
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>sql::statements::describe::DescribeTable;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>sql::util::format_raw_object_name;
<a href=#23 id=23 data-nosnippet>23</a>
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span><span class="kw">crate</span>::error::{
<a href=#25 id=25 data-nosnippet>25</a> CatalogSnafu, DescribeStatementSnafu, ExternalSnafu, <span class="prelude-ty">Result</span>, TableNotFoundSnafu,
<a href=#26 id=26 data-nosnippet>26</a>};
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#28 id=28 data-nosnippet>28</a>
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">impl </span>StatementExecutor {
<a href=#30 id=30 data-nosnippet>30</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#31 id=31 data-nosnippet>31</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>describe_table(
<a href=#32 id=32 data-nosnippet>32</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#33 id=33 data-nosnippet>33</a> stmt: DescribeTable,
<a href=#34 id=34 data-nosnippet>34</a> query_ctx: QueryContextRef,
<a href=#35 id=35 data-nosnippet>35</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#36 id=36 data-nosnippet>36</a> <span class="kw">let </span>(catalog, schema, table) = table_idents_to_full_name(stmt.name(), <span class="kw-2">&amp;</span>query_ctx)
<a href=#37 id=37 data-nosnippet>37</a> .map_err(BoxedError::new)
<a href=#38 id=38 data-nosnippet>38</a> .context(ExternalSnafu)<span class="question-mark">?</span>;
<a href=#39 id=39 data-nosnippet>39</a>
<a href=#40 id=40 data-nosnippet>40</a> <span class="kw">let </span>table = <span class="self">self
<a href=#41 id=41 data-nosnippet>41</a> </span>.catalog_manager
<a href=#42 id=42 data-nosnippet>42</a> .table(<span class="kw-2">&amp;</span>catalog, <span class="kw-2">&amp;</span>schema, <span class="kw-2">&amp;</span>table, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>query_ctx))
<a href=#43 id=43 data-nosnippet>43</a> .<span class="kw">await
<a href=#44 id=44 data-nosnippet>44</a> </span>.context(CatalogSnafu)<span class="question-mark">?
<a href=#45 id=45 data-nosnippet>45</a> </span>.with_context(|| TableNotFoundSnafu {
<a href=#46 id=46 data-nosnippet>46</a> table_name: format_raw_object_name(stmt.name()),
<a href=#47 id=47 data-nosnippet>47</a> })<span class="question-mark">?</span>;
<a href=#48 id=48 data-nosnippet>48</a>
<a href=#49 id=49 data-nosnippet>49</a> query::sql::describe_table(table).context(DescribeStatementSnafu)
<a href=#50 id=50 data-nosnippet>50</a> }
<a href=#51 id=51 data-nosnippet>51</a>}
</code></pre></div></section></main></body></html>
+73
View File
@@ -0,0 +1,73 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/dml.rs`."><title>dml.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>dml.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-2"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>common_catalog::consts::is_ddl_reserved_table;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>common_error::ext::BoxedError;
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_query::Output;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_telemetry::tracing;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>query::parser::QueryStatement;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>session::table_name::table_idents_to_full_name;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>snafu::ResultExt;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>sql::statements::insert::Insert;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>sql::statements::statement::Statement;
<a href=#25 id=25 data-nosnippet>25</a>
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span><span class="kw">crate</span>::error::{ExternalSnafu, ParseSqlSnafu, <span class="prelude-ty">Result</span>};
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#28 id=28 data-nosnippet>28</a>
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">impl </span>StatementExecutor {
<a href=#30 id=30 data-nosnippet>30</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#31 id=31 data-nosnippet>31</a> </span><span class="kw">pub async fn </span>insert(<span class="kw-2">&amp;</span><span class="self">self</span>, insert: Box&lt;Insert&gt;, query_ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#32 id=32 data-nosnippet>32</a> <span class="self">self</span>.create_ddl_reserved_target_on_demand(<span class="kw-2">&amp;</span>insert, <span class="kw-2">&amp;</span>query_ctx)
<a href=#33 id=33 data-nosnippet>33</a> .<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#34 id=34 data-nosnippet>34</a> <span class="kw">if </span>insert.can_extract_values() {
<a href=#35 id=35 data-nosnippet>35</a> <span class="comment">// Fast path: plain insert ("insert with literal values") is executed directly
<a href=#36 id=36 data-nosnippet>36</a> </span><span class="self">self</span>.inserter
<a href=#37 id=37 data-nosnippet>37</a> .handle_statement_insert(insert.as_ref(), <span class="kw-2">&amp;</span>query_ctx)
<a href=#38 id=38 data-nosnippet>38</a> .<span class="kw">await
<a href=#39 id=39 data-nosnippet>39</a> </span>} <span class="kw">else </span>{
<a href=#40 id=40 data-nosnippet>40</a> <span class="comment">// Slow path: insert with subquery. Execute using query engine.
<a href=#41 id=41 data-nosnippet>41</a> </span><span class="kw">let </span>statement = QueryStatement::Sql(Statement::Insert(insert));
<a href=#42 id=42 data-nosnippet>42</a> <span class="self">self</span>.plan_exec(statement, query_ctx).<span class="kw">await
<a href=#43 id=43 data-nosnippet>43</a> </span>}
<a href=#44 id=44 data-nosnippet>44</a> }
<a href=#45 id=45 data-nosnippet>45</a>
<a href=#46 id=46 data-nosnippet>46</a> <span class="doccomment">/// A DDL-reserved table is defined by the system, not the user: its first
<a href=#47 id=47 data-nosnippet>47</a> /// INSERT creates it here with the canonical schema.
<a href=#48 id=48 data-nosnippet>48</a> </span><span class="kw">async fn </span>create_ddl_reserved_target_on_demand(
<a href=#49 id=49 data-nosnippet>49</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#50 id=50 data-nosnippet>50</a> insert: <span class="kw-2">&amp;</span>Insert,
<a href=#51 id=51 data-nosnippet>51</a> query_ctx: <span class="kw-2">&amp;</span>QueryContextRef,
<a href=#52 id=52 data-nosnippet>52</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#53 id=53 data-nosnippet>53</a> <span class="kw">let </span>table_name = insert.table_name().context(ParseSqlSnafu)<span class="question-mark">?</span>;
<a href=#54 id=54 data-nosnippet>54</a> <span class="kw">let </span>(catalog, schema, table) = table_idents_to_full_name(table_name, query_ctx)
<a href=#55 id=55 data-nosnippet>55</a> .map_err(BoxedError::new)
<a href=#56 id=56 data-nosnippet>56</a> .context(ExternalSnafu)<span class="question-mark">?</span>;
<a href=#57 id=57 data-nosnippet>57</a> <span class="kw">if </span>!is_ddl_reserved_table(<span class="kw-2">&amp;</span>schema, <span class="kw-2">&amp;</span>table) {
<a href=#58 id=58 data-nosnippet>58</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(());
<a href=#59 id=59 data-nosnippet>59</a> }
<a href=#60 id=60 data-nosnippet>60</a> <span class="kw">let </span>exists = <span class="self">self
<a href=#61 id=61 data-nosnippet>61</a> </span>.catalog_manager
<a href=#62 id=62 data-nosnippet>62</a> .table_exists(<span class="kw-2">&amp;</span>catalog, <span class="kw-2">&amp;</span>schema, <span class="kw-2">&amp;</span>table, <span class="prelude-val">Some</span>(query_ctx))
<a href=#63 id=63 data-nosnippet>63</a> .<span class="kw">await
<a href=#64 id=64 data-nosnippet>64</a> </span>.map_err(BoxedError::new)
<a href=#65 id=65 data-nosnippet>65</a> .context(ExternalSnafu)<span class="question-mark">?</span>;
<a href=#66 id=66 data-nosnippet>66</a> <span class="kw">if </span>!exists {
<a href=#67 id=67 data-nosnippet>67</a> <span class="self">self</span>.create_declared_relationships_table(<span class="kw-2">&amp;</span>catalog, query_ctx.clone())
<a href=#68 id=68 data-nosnippet>68</a> .<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#69 id=69 data-nosnippet>69</a> }
<a href=#70 id=70 data-nosnippet>70</a> <span class="prelude-val">Ok</span>(())
<a href=#71 id=71 data-nosnippet>71</a> }
<a href=#72 id=72 data-nosnippet>72</a>}
</code></pre></div></section></main></body></html>
@@ -0,0 +1,476 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/export_database.rs`."><title>export_database.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>export_database.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="doccomment">//! Request-scoped database export preparation. The frontend authorizes the entire
<a href=#16 id=16 data-nosnippet>16</a>//! captured selection before preparation; these trusted methods do not check ACLs.
<a href=#17 id=17 data-nosnippet>17</a>
<a href=#18 id=18 data-nosnippet>18</a></span><span class="kw">use </span>std::collections::{BTreeMap, HashSet};
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>std::future::Future;
<a href=#20 id=20 data-nosnippet>20</a>
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>common_datasource::file_format::Format;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>common_datasource::object_store::{FILE_SCHEMA, FS_SCHEMA, build_backend_for_write, parse_url};
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>common_datasource::packed_writer::{PackedTableWriter, PackedWriter};
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>common_meta::key::table_route::TableRouteValue;
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>futures::StreamExt;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>futures::stream::FuturesUnordered;
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>snafu::{OptionExt, ResultExt, ensure};
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span>store_api::metric_engine_consts::{LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME};
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span>table::TableRef;
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span>table::metadata::TableType;
<a href=#32 id=32 data-nosnippet>32</a><span class="kw">use </span>table::requests::{CopyDatabaseRequest, CopyDirection, CopyTableRequest};
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">use </span>tokio_util::sync::CancellationToken;
<a href=#34 id=34 data-nosnippet>34</a>
<a href=#35 id=35 data-nosnippet>35</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, InvalidDatabaseExportSnafu, <span class="prelude-ty">Result</span>};
<a href=#36 id=36 data-nosnippet>36</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#37 id=37 data-nosnippet>37</a><span class="kw">use </span><span class="kw">crate</span>::statement::database_copy::{
<a href=#38 id=38 data-nosnippet>38</a> DatabaseExportFile, parse_parallelism_from_option_map, validate_database_directory,
<a href=#39 id=39 data-nosnippet>39</a> validate_database_export_layout,
<a href=#40 id=40 data-nosnippet>40</a>};
<a href=#41 id=41 data-nosnippet>41</a><span class="kw">use </span><span class="kw">crate</span>::statement::export_logical_tables::writers::{ExportWriteBudget, retain_error};
<a href=#42 id=42 data-nosnippet>42</a><span class="kw">use </span><span class="kw">crate</span>::statement::export_logical_tables::{LogicalTableExport, LogicalTableExportLimits};
<a href=#43 id=43 data-nosnippet>43</a>
<a href=#44 id=44 data-nosnippet>44</a><span class="doccomment">/// A validated request-scoped selection, not a metadata snapshot or an ACL token.
<a href=#45 id=45 data-nosnippet>45</a></span><span class="kw">pub struct </span>PreparedDatabaseExport {
<a href=#46 id=46 data-nosnippet>46</a> request: CopyDatabaseRequest,
<a href=#47 id=47 data-nosnippet>47</a> jobs: Vec&lt;DatabaseExportJob&gt;,
<a href=#48 id=48 data-nosnippet>48</a>}
<a href=#49 id=49 data-nosnippet>49</a>
<a href=#50 id=50 data-nosnippet>50</a><span class="kw">impl </span>PreparedDatabaseExport {
<a href=#51 id=51 data-nosnippet>51</a> <span class="doccomment">/// Inspect physical grouping through the integration testing adapter.
<a href=#52 id=52 data-nosnippet>52</a> </span><span class="attr">#[cfg(feature = <span class="string">"testing"</span>)]
<a href=#53 id=53 data-nosnippet>53</a> </span><span class="kw">pub fn </span>job_count_for_test(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; usize {
<a href=#54 id=54 data-nosnippet>54</a> <span class="self">self</span>.jobs.len()
<a href=#55 id=55 data-nosnippet>55</a> }
<a href=#56 id=56 data-nosnippet>56</a>}
<a href=#57 id=57 data-nosnippet>57</a>
<a href=#58 id=58 data-nosnippet>58</a><span class="kw">enum </span>DatabaseExportJob {
<a href=#59 id=59 data-nosnippet>59</a> Ordinary {
<a href=#60 id=60 data-nosnippet>60</a> table: TableRef,
<a href=#61 id=61 data-nosnippet>61</a> output: DatabaseExportFile,
<a href=#62 id=62 data-nosnippet>62</a> },
<a href=#63 id=63 data-nosnippet>63</a> Metric(LogicalTableExport),
<a href=#64 id=64 data-nosnippet>64</a>}
<a href=#65 id=65 data-nosnippet>65</a>
<a href=#66 id=66 data-nosnippet>66</a><span class="kw">impl </span>StatementExecutor {
<a href=#67 id=67 data-nosnippet>67</a> <span class="doccomment">/// Validate all jobs and destinations before executing any query or writer.
<a href=#68 id=68 data-nosnippet>68</a> /// Callers must authorize every captured table before calling this method.
<a href=#69 id=69 data-nosnippet>69</a> </span><span class="kw">pub async fn </span>prepare_database_export(
<a href=#70 id=70 data-nosnippet>70</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#71 id=71 data-nosnippet>71</a> req: CopyDatabaseRequest,
<a href=#72 id=72 data-nosnippet>72</a> tables: Vec&lt;TableRef&gt;,
<a href=#73 id=73 data-nosnippet>73</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;PreparedDatabaseExport&gt; {
<a href=#74 id=74 data-nosnippet>74</a> <span class="kw">if </span>req
<a href=#75 id=75 data-nosnippet>75</a> .with
<a href=#76 id=76 data-nosnippet>76</a> .get(<span class="string">"metric_data_layout"</span>)
<a href=#77 id=77 data-nosnippet>77</a> .is_none_or(|v| v != <span class="string">"packed"</span>)
<a href=#78 id=78 data-nosnippet>78</a> {
<a href=#79 id=79 data-nosnippet>79</a> validate_database_export_layout(<span class="kw-2">&amp;</span>req.with)<span class="question-mark">?</span>;
<a href=#80 id=80 data-nosnippet>80</a> }
<a href=#81 id=81 data-nosnippet>81</a> validate_database_directory(<span class="kw-2">&amp;</span>req.location)<span class="question-mark">?</span>;
<a href=#82 id=82 data-nosnippet>82</a> <span class="kw">let </span>format = Format::try_from(<span class="kw-2">&amp;</span>req.with).context(error::ParseFileFormatSnafu)<span class="question-mark">?</span>;
<a href=#83 id=83 data-nosnippet>83</a> <span class="macro">ensure!</span>(
<a href=#84 id=84 data-nosnippet>84</a> <span class="macro">matches!</span>(format, Format::Parquet(<span class="kw">_</span>)),
<a href=#85 id=85 data-nosnippet>85</a> error::UnsupportedFormatSnafu { format }
<a href=#86 id=86 data-nosnippet>86</a> );
<a href=#87 id=87 data-nosnippet>87</a> <span class="kw">let </span>(scheme, <span class="kw">_</span>, <span class="kw">_</span>) = parse_url(<span class="kw-2">&amp;</span>req.location).context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#88 id=88 data-nosnippet>88</a> <span class="kw">let </span>local =
<a href=#89 id=89 data-nosnippet>89</a> scheme.eq_ignore_ascii_case(FS_SCHEMA) || scheme.eq_ignore_ascii_case(FILE_SCHEMA);
<a href=#90 id=90 data-nosnippet>90</a> <span class="kw">let </span><span class="kw-2">mut </span>filenames = HashSet::new();
<a href=#91 id=91 data-nosnippet>91</a> <span class="kw">let </span><span class="kw-2">mut </span>logical = Vec::new();
<a href=#92 id=92 data-nosnippet>92</a> <span class="kw">let </span><span class="kw-2">mut </span>jobs = Vec::new();
<a href=#93 id=93 data-nosnippet>93</a> <span class="kw">for </span>table <span class="kw">in </span>tables {
<a href=#94 id=94 data-nosnippet>94</a> <span class="kw">let </span>info = table.table_info();
<a href=#95 id=95 data-nosnippet>95</a> <span class="kw">let </span>name = <span class="kw-2">&amp;</span>info.name;
<a href=#96 id=96 data-nosnippet>96</a> <span class="kw">let </span>output = DatabaseExportFile::new(<span class="kw-2">&amp;</span>req.location, name, <span class="string">".parquet"</span>)<span class="question-mark">?</span>;
<a href=#97 id=97 data-nosnippet>97</a> <span class="macro">ensure!</span>(
<a href=#98 id=98 data-nosnippet>98</a> filenames.insert(<span class="kw">if </span>local {
<a href=#99 id=99 data-nosnippet>99</a> output.path.to_ascii_lowercase()
<a href=#100 id=100 data-nosnippet>100</a> } <span class="kw">else </span>{
<a href=#101 id=101 data-nosnippet>101</a> output.path.clone()
<a href=#102 id=102 data-nosnippet>102</a> }),
<a href=#103 id=103 data-nosnippet>103</a> InvalidDatabaseExportSnafu {
<a href=#104 id=104 data-nosnippet>104</a> reason: <span class="macro">format!</span>(<span class="string">"duplicate output name: {name}"</span>)
<a href=#105 id=105 data-nosnippet>105</a> }
<a href=#106 id=106 data-nosnippet>106</a> );
<a href=#107 id=107 data-nosnippet>107</a> <span class="macro">ensure!</span>(
<a href=#108 id=108 data-nosnippet>108</a> info.catalog_name == req.catalog_name
<a href=#109 id=109 data-nosnippet>109</a> &amp;&amp; info.schema_name == req.schema_name
<a href=#110 id=110 data-nosnippet>110</a> &amp;&amp; table.table_type() == TableType::Base,
<a href=#111 id=111 data-nosnippet>111</a> InvalidDatabaseExportSnafu {
<a href=#112 id=112 data-nosnippet>112</a> reason: <span class="string">"expected base tables in the selected schema"
<a href=#113 id=113 data-nosnippet>113</a> </span>}
<a href=#114 id=114 data-nosnippet>114</a> );
<a href=#115 id=115 data-nosnippet>115</a> <span class="kw">if </span>info.meta.engine == METRIC_ENGINE_NAME {
<a href=#116 id=116 data-nosnippet>116</a> <span class="macro">ensure!</span>(
<a href=#117 id=117 data-nosnippet>117</a> info.meta
<a href=#118 id=118 data-nosnippet>118</a> .options
<a href=#119 id=119 data-nosnippet>119</a> .extra_options
<a href=#120 id=120 data-nosnippet>120</a> .contains_key(LOGICAL_TABLE_METADATA_KEY),
<a href=#121 id=121 data-nosnippet>121</a> InvalidDatabaseExportSnafu {
<a href=#122 id=122 data-nosnippet>122</a> reason: <span class="string">"expected a Metric logical table"
<a href=#123 id=123 data-nosnippet>123</a> </span>}
<a href=#124 id=124 data-nosnippet>124</a> );
<a href=#125 id=125 data-nosnippet>125</a> logical.push(table);
<a href=#126 id=126 data-nosnippet>126</a> } <span class="kw">else </span>{
<a href=#127 id=127 data-nosnippet>127</a> jobs.push(DatabaseExportJob::Ordinary { table, output });
<a href=#128 id=128 data-nosnippet>128</a> }
<a href=#129 id=129 data-nosnippet>129</a> }
<a href=#130 id=130 data-nosnippet>130</a> <span class="kw">let </span>ids = logical
<a href=#131 id=131 data-nosnippet>131</a> .iter()
<a href=#132 id=132 data-nosnippet>132</a> .map(|t| t.table_info().table_id())
<a href=#133 id=133 data-nosnippet>133</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#134 id=134 data-nosnippet>134</a> <span class="kw">let </span>routes = <span class="self">self
<a href=#135 id=135 data-nosnippet>135</a> </span>.table_metadata_manager
<a href=#136 id=136 data-nosnippet>136</a> .table_route_manager()
<a href=#137 id=137 data-nosnippet>137</a> .table_route_storage()
<a href=#138 id=138 data-nosnippet>138</a> .batch_get(<span class="kw-2">&amp;</span>ids)
<a href=#139 id=139 data-nosnippet>139</a> .<span class="kw">await
<a href=#140 id=140 data-nosnippet>140</a> </span>.context(error::TableMetadataManagerSnafu)<span class="question-mark">?</span>;
<a href=#141 id=141 data-nosnippet>141</a> <span class="kw">let </span><span class="kw-2">mut </span>groups = BTreeMap::&lt;<span class="kw">_</span>, Vec&lt;TableRef&gt;&gt;::new();
<a href=#142 id=142 data-nosnippet>142</a> <span class="kw">for </span>(table, route) <span class="kw">in </span>logical.into_iter().zip(routes) {
<a href=#143 id=143 data-nosnippet>143</a> <span class="kw">let </span><span class="prelude-val">Some</span>(TableRouteValue::Logical(route)) = route <span class="kw">else </span>{
<a href=#144 id=144 data-nosnippet>144</a> <span class="kw">return </span>InvalidDatabaseExportSnafu {
<a href=#145 id=145 data-nosnippet>145</a> reason: <span class="macro">format!</span>(
<a href=#146 id=146 data-nosnippet>146</a> <span class="string">"missing or non-logical route for {}"</span>,
<a href=#147 id=147 data-nosnippet>147</a> table.table_info().table_id()
<a href=#148 id=148 data-nosnippet>148</a> ),
<a href=#149 id=149 data-nosnippet>149</a> }
<a href=#150 id=150 data-nosnippet>150</a> .fail();
<a href=#151 id=151 data-nosnippet>151</a> };
<a href=#152 id=152 data-nosnippet>152</a> groups
<a href=#153 id=153 data-nosnippet>153</a> .entry(route.physical_table_id())
<a href=#154 id=154 data-nosnippet>154</a> .or_default()
<a href=#155 id=155 data-nosnippet>155</a> .push(table);
<a href=#156 id=156 data-nosnippet>156</a> }
<a href=#157 id=157 data-nosnippet>157</a> <span class="kw">let </span>physical_ids = groups.keys().copied().collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#158 id=158 data-nosnippet>158</a> <span class="kw">let </span><span class="kw-2">mut </span>physical = <span class="self">self
<a href=#159 id=159 data-nosnippet>159</a> </span>.catalog_manager
<a href=#160 id=160 data-nosnippet>160</a> .tables_by_ids(<span class="kw-2">&amp;</span>req.catalog_name, <span class="kw-2">&amp;</span>req.schema_name, <span class="kw-2">&amp;</span>physical_ids)
<a href=#161 id=161 data-nosnippet>161</a> .<span class="kw">await
<a href=#162 id=162 data-nosnippet>162</a> </span>.context(error::CatalogSnafu)<span class="question-mark">?
<a href=#163 id=163 data-nosnippet>163</a> </span>.into_iter()
<a href=#164 id=164 data-nosnippet>164</a> .map(|t| (t.table_info().table_id(), t))
<a href=#165 id=165 data-nosnippet>165</a> .collect::&lt;BTreeMap&lt;<span class="kw">_</span>, <span class="kw">_</span>&gt;&gt;();
<a href=#166 id=166 data-nosnippet>166</a> <span class="kw">for </span>(id, tables) <span class="kw">in </span>groups {
<a href=#167 id=167 data-nosnippet>167</a> <span class="kw">let </span>table = physical
<a href=#168 id=168 data-nosnippet>168</a> .remove(<span class="kw-2">&amp;</span>id)
<a href=#169 id=169 data-nosnippet>169</a> .with_context(|| InvalidDatabaseExportSnafu {
<a href=#170 id=170 data-nosnippet>170</a> reason: <span class="macro">format!</span>(<span class="string">"missing physical table {id} in the selected schema"</span>),
<a href=#171 id=171 data-nosnippet>171</a> })<span class="question-mark">?</span>;
<a href=#172 id=172 data-nosnippet>172</a> jobs.push(DatabaseExportJob::Metric(
<a href=#173 id=173 data-nosnippet>173</a> LogicalTableExport::try_new_in_directory(table, <span class="kw-2">&amp;</span>tables, <span class="kw-2">&amp;</span>req.location)<span class="question-mark">?</span>,
<a href=#174 id=174 data-nosnippet>174</a> ));
<a href=#175 id=175 data-nosnippet>175</a> }
<a href=#176 id=176 data-nosnippet>176</a> build_backend_for_write(<span class="kw-2">&amp;</span>req.location, <span class="kw-2">&amp;</span>req.connection, <span class="kw-2">&amp;</span><span class="self">self</span>.local_file_access)
<a href=#177 id=177 data-nosnippet>177</a> .<span class="kw">await
<a href=#178 id=178 data-nosnippet>178</a> </span>.context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#179 id=179 data-nosnippet>179</a> <span class="prelude-val">Ok</span>(PreparedDatabaseExport { request: req, jobs })
<a href=#180 id=180 data-nosnippet>180</a> }
<a href=#181 id=181 data-nosnippet>181</a>}
<a href=#182 id=182 data-nosnippet>182</a>
<a href=#183 id=183 data-nosnippet>183</a><span class="doccomment">/// Returned only after every started job has closed or drained its owned I/O.
<a href=#184 id=184 data-nosnippet>184</a></span><span class="attr">#[derive(Debug)]
<a href=#185 id=185 data-nosnippet>185</a></span><span class="kw">pub struct </span>DatabaseExportSummary {
<a href=#186 id=186 data-nosnippet>186</a> <span class="kw">pub </span>rows: usize,
<a href=#187 id=187 data-nosnippet>187</a> <span class="kw">pub </span>output_files: Vec&lt;String&gt;,
<a href=#188 id=188 data-nosnippet>188</a>}
<a href=#189 id=189 data-nosnippet>189</a>
<a href=#190 id=190 data-nosnippet>190</a><span class="kw">impl </span>StatementExecutor {
<a href=#191 id=191 data-nosnippet>191</a> <span class="doccomment">/// Execute with one job budget. The caller must keep this future alive until
<a href=#192 id=192 data-nosnippet>192</a> /// it returns, including after cooperative cancellation. Closed files remain
<a href=#193 id=193 data-nosnippet>193</a> /// owned by the caller's attempt; this method does not publish completion.
<a href=#194 id=194 data-nosnippet>194</a> </span><span class="kw">pub async fn </span>export_database(
<a href=#195 id=195 data-nosnippet>195</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#196 id=196 data-nosnippet>196</a> plan: PreparedDatabaseExport,
<a href=#197 id=197 data-nosnippet>197</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#198 id=198 data-nosnippet>198</a> ctx: QueryContextRef,
<a href=#199 id=199 data-nosnippet>199</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;DatabaseExportSummary&gt; {
<a href=#200 id=200 data-nosnippet>200</a> <span class="kw">let </span><span class="kw-2">mut </span>output_files = Vec::new();
<a href=#201 id=201 data-nosnippet>201</a> <span class="kw">for </span>job <span class="kw">in </span><span class="kw-2">&amp;</span>plan.jobs {
<a href=#202 id=202 data-nosnippet>202</a> <span class="kw">match </span>job {
<a href=#203 id=203 data-nosnippet>203</a> DatabaseExportJob::Ordinary { output, .. } =&gt; {
<a href=#204 id=204 data-nosnippet>204</a> output_files.push(output.location.clone())
<a href=#205 id=205 data-nosnippet>205</a> }
<a href=#206 id=206 data-nosnippet>206</a> DatabaseExportJob::Metric(unit) =&gt; {
<a href=#207 id=207 data-nosnippet>207</a> output_files.extend(unit.output_files().map(|file| file.location.clone()))
<a href=#208 id=208 data-nosnippet>208</a> }
<a href=#209 id=209 data-nosnippet>209</a> }
<a href=#210 id=210 data-nosnippet>210</a> }
<a href=#211 id=211 data-nosnippet>211</a> output_files.sort();
<a href=#212 id=212 data-nosnippet>212</a> <span class="kw">let </span>req = <span class="kw-2">&amp;</span>plan.request;
<a href=#213 id=213 data-nosnippet>213</a> <span class="kw">let </span>parallelism = parse_parallelism_from_option_map(<span class="kw-2">&amp;</span>req.with);
<a href=#214 id=214 data-nosnippet>214</a> <span class="kw">let </span>budget = ExportWriteBudget::new(parallelism);
<a href=#215 id=215 data-nosnippet>215</a> <span class="kw">let </span>packed = <span class="kw">if </span>req
<a href=#216 id=216 data-nosnippet>216</a> .with
<a href=#217 id=217 data-nosnippet>217</a> .get(<span class="string">"metric_data_layout"</span>)
<a href=#218 id=218 data-nosnippet>218</a> .is_some_and(|v| v == <span class="string">"packed"</span>)
<a href=#219 id=219 data-nosnippet>219</a> {
<a href=#220 id=220 data-nosnippet>220</a> <span class="kw">let </span>store =
<a href=#221 id=221 data-nosnippet>221</a> build_backend_for_write(<span class="kw-2">&amp;</span>req.location, <span class="kw-2">&amp;</span>req.connection, <span class="kw-2">&amp;</span><span class="self">self</span>.local_file_access)
<a href=#222 id=222 data-nosnippet>222</a> .<span class="kw">await
<a href=#223 id=223 data-nosnippet>223</a> </span>.context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#224 id=224 data-nosnippet>224</a> <span class="prelude-val">Some</span>(
<a href=#225 id=225 data-nosnippet>225</a> PackedWriter::new(store).context(error::WriteStreamToFileSnafu {
<a href=#226 id=226 data-nosnippet>226</a> path: <span class="kw-2">&amp;</span>req.location,
<a href=#227 id=227 data-nosnippet>227</a> })<span class="question-mark">?</span>,
<a href=#228 id=228 data-nosnippet>228</a> )
<a href=#229 id=229 data-nosnippet>229</a> } <span class="kw">else </span>{
<a href=#230 id=230 data-nosnippet>230</a> <span class="prelude-val">None
<a href=#231 id=231 data-nosnippet>231</a> </span>};
<a href=#232 id=232 data-nosnippet>232</a> <span class="kw">let </span>rows = run_database_export_jobs(plan.jobs, parallelism, cancellation, |job, token| {
<a href=#233 id=233 data-nosnippet>233</a> <span class="kw">let </span>ctx = ctx.clone();
<a href=#234 id=234 data-nosnippet>234</a> <span class="kw">let </span>budget = budget.clone();
<a href=#235 id=235 data-nosnippet>235</a> <span class="kw">let </span>packed = packed.clone();
<a href=#236 id=236 data-nosnippet>236</a> <span class="kw">async move </span>{
<a href=#237 id=237 data-nosnippet>237</a> <span class="kw">match </span>job {
<a href=#238 id=238 data-nosnippet>238</a> DatabaseExportJob::Metric(unit) =&gt; <span class="self">self
<a href=#239 id=239 data-nosnippet>239</a> </span>.export_logical_tables_managed(
<a href=#240 id=240 data-nosnippet>240</a> <span class="kw-2">&amp;</span>unit,
<a href=#241 id=241 data-nosnippet>241</a> <span class="kw-2">&amp;</span>req.location,
<a href=#242 id=242 data-nosnippet>242</a> <span class="kw-2">&amp;</span>req.connection,
<a href=#243 id=243 data-nosnippet>243</a> req.time_range.as_ref(),
<a href=#244 id=244 data-nosnippet>244</a> LogicalTableExportLimits::default(),
<a href=#245 id=245 data-nosnippet>245</a> <span class="kw-2">&amp;</span>token,
<a href=#246 id=246 data-nosnippet>246</a> ctx,
<a href=#247 id=247 data-nosnippet>247</a> budget,
<a href=#248 id=248 data-nosnippet>248</a> packed,
<a href=#249 id=249 data-nosnippet>249</a> )
<a href=#250 id=250 data-nosnippet>250</a> .<span class="kw">await
<a href=#251 id=251 data-nosnippet>251</a> </span>.map(|summary| summary.rows),
<a href=#252 id=252 data-nosnippet>252</a> DatabaseExportJob::Ordinary { table, output } =&gt; {
<a href=#253 id=253 data-nosnippet>253</a> <span class="kw">let </span>info = table.table_info();
<a href=#254 id=254 data-nosnippet>254</a> <span class="kw">let </span>copy = CopyTableRequest {
<a href=#255 id=255 data-nosnippet>255</a> catalog_name: info.catalog_name.clone(),
<a href=#256 id=256 data-nosnippet>256</a> schema_name: info.schema_name.clone(),
<a href=#257 id=257 data-nosnippet>257</a> table_name: info.name.clone(),
<a href=#258 id=258 data-nosnippet>258</a> location: output.location,
<a href=#259 id=259 data-nosnippet>259</a> with: req.with.clone(),
<a href=#260 id=260 data-nosnippet>260</a> connection: req.connection.clone(),
<a href=#261 id=261 data-nosnippet>261</a> pattern: <span class="prelude-val">None</span>,
<a href=#262 id=262 data-nosnippet>262</a> direction: CopyDirection::Export,
<a href=#263 id=263 data-nosnippet>263</a> timestamp_range: req.time_range,
<a href=#264 id=264 data-nosnippet>264</a> limit: <span class="prelude-val">None</span>,
<a href=#265 id=265 data-nosnippet>265</a> };
<a href=#266 id=266 data-nosnippet>266</a> <span class="kw">let </span>_permit = budget.writer(<span class="kw-2">&amp;</span>token).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#267 id=267 data-nosnippet>267</a> <span class="kw">let </span>destination = packed.map(|shared| {
<a href=#268 id=268 data-nosnippet>268</a> PackedTableWriter::new(shared, info.name.clone(), info.table_id(), <span class="bool-val">true</span>)
<a href=#269 id=269 data-nosnippet>269</a> });
<a href=#270 id=270 data-nosnippet>270</a> <span class="self">self</span>.copy_captured_table_to_managed(
<a href=#271 id=271 data-nosnippet>271</a> table,
<a href=#272 id=272 data-nosnippet>272</a> copy,
<a href=#273 id=273 data-nosnippet>273</a> ctx,
<a href=#274 id=274 data-nosnippet>274</a> <span class="prelude-val">Some</span>((<span class="kw-2">&amp;</span>budget, <span class="kw-2">&amp;</span>token)),
<a href=#275 id=275 data-nosnippet>275</a> destination,
<a href=#276 id=276 data-nosnippet>276</a> )
<a href=#277 id=277 data-nosnippet>277</a> .<span class="kw">await
<a href=#278 id=278 data-nosnippet>278</a> </span>}
<a href=#279 id=279 data-nosnippet>279</a> }
<a href=#280 id=280 data-nosnippet>280</a> }
<a href=#281 id=281 data-nosnippet>281</a> })
<a href=#282 id=282 data-nosnippet>282</a> .<span class="kw">await</span>;
<a href=#283 id=283 data-nosnippet>283</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(packed) = packed {
<a href=#284 id=284 data-nosnippet>284</a> <span class="kw">let </span><span class="kw-2">mut </span>packed = packed.lock().<span class="kw">await</span>;
<a href=#285 id=285 data-nosnippet>285</a> <span class="kw">let </span>result = <span class="kw">match </span>rows {
<a href=#286 id=286 data-nosnippet>286</a> <span class="prelude-val">Ok</span>(rows) =&gt; packed
<a href=#287 id=287 data-nosnippet>287</a> .finish(cancellation)
<a href=#288 id=288 data-nosnippet>288</a> .<span class="kw">await
<a href=#289 id=289 data-nosnippet>289</a> </span>.context(error::WriteStreamToFileSnafu {
<a href=#290 id=290 data-nosnippet>290</a> path: <span class="kw-2">&amp;</span>req.location,
<a href=#291 id=291 data-nosnippet>291</a> })
<a href=#292 id=292 data-nosnippet>292</a> .and_then(|files| {
<a href=#293 id=293 data-nosnippet>293</a> <span class="prelude-val">Ok</span>(DatabaseExportSummary {
<a href=#294 id=294 data-nosnippet>294</a> rows,
<a href=#295 id=295 data-nosnippet>295</a> output_files: files
<a href=#296 id=296 data-nosnippet>296</a> .into_iter()
<a href=#297 id=297 data-nosnippet>297</a> .map(|file| {
<a href=#298 id=298 data-nosnippet>298</a> DatabaseExportFile::new(<span class="kw-2">&amp;</span>req.location, <span class="kw-2">&amp;</span>file, <span class="string">""</span>)
<a href=#299 id=299 data-nosnippet>299</a> .map(|f| f.location)
<a href=#300 id=300 data-nosnippet>300</a> })
<a href=#301 id=301 data-nosnippet>301</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;<span class="kw">_</span>&gt;&gt;()<span class="question-mark">?</span>,
<a href=#302 id=302 data-nosnippet>302</a> })
<a href=#303 id=303 data-nosnippet>303</a> }),
<a href=#304 id=304 data-nosnippet>304</a> <span class="prelude-val">Err</span>(error) =&gt; <span class="prelude-val">Err</span>(error),
<a href=#305 id=305 data-nosnippet>305</a> };
<a href=#306 id=306 data-nosnippet>306</a> <span class="kw">if </span>result.is_err()
<a href=#307 id=307 data-nosnippet>307</a> &amp;&amp; <span class="kw">let </span><span class="prelude-val">Err</span>(error) = packed.abort().<span class="kw">await
<a href=#308 id=308 data-nosnippet>308</a> </span>{
<a href=#309 id=309 data-nosnippet>309</a> <span class="macro">common_telemetry::warn!</span>(error; <span class="string">"Failed to abort Metric pack"</span>);
<a href=#310 id=310 data-nosnippet>310</a> }
<a href=#311 id=311 data-nosnippet>311</a> result
<a href=#312 id=312 data-nosnippet>312</a> } <span class="kw">else </span>{
<a href=#313 id=313 data-nosnippet>313</a> <span class="prelude-val">Ok</span>(DatabaseExportSummary {
<a href=#314 id=314 data-nosnippet>314</a> rows: rows<span class="question-mark">?</span>,
<a href=#315 id=315 data-nosnippet>315</a> output_files,
<a href=#316 id=316 data-nosnippet>316</a> })
<a href=#317 id=317 data-nosnippet>317</a> }
<a href=#318 id=318 data-nosnippet>318</a> }
<a href=#319 id=319 data-nosnippet>319</a>}
<a href=#320 id=320 data-nosnippet>320</a>
<a href=#321 id=321 data-nosnippet>321</a><span class="kw">async fn </span>run_database_export_jobs&lt;J, F: Future&lt;Output = <span class="prelude-ty">Result</span>&lt;usize&gt;&gt;&gt;(
<a href=#322 id=322 data-nosnippet>322</a> jobs: <span class="kw">impl </span>IntoIterator&lt;Item = J&gt;,
<a href=#323 id=323 data-nosnippet>323</a> parallelism: usize,
<a href=#324 id=324 data-nosnippet>324</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#325 id=325 data-nosnippet>325</a> <span class="kw-2">mut </span>run: <span class="kw">impl </span>FnMut(J, CancellationToken) -&gt; F,
<a href=#326 id=326 data-nosnippet>326</a>) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#327 id=327 data-nosnippet>327</a> <span class="kw">let </span>token = CancellationToken::new();
<a href=#328 id=328 data-nosnippet>328</a> <span class="kw">let </span><span class="kw-2">mut </span>jobs = jobs.into_iter();
<a href=#329 id=329 data-nosnippet>329</a> <span class="kw">let </span><span class="kw-2">mut </span>active = FuturesUnordered::new();
<a href=#330 id=330 data-nosnippet>330</a> <span class="kw">let </span><span class="kw-2">mut </span>first_error = <span class="prelude-val">None</span>;
<a href=#331 id=331 data-nosnippet>331</a> <span class="kw">let </span><span class="kw-2">mut </span>rows = <span class="number">0</span>;
<a href=#332 id=332 data-nosnippet>332</a> <span class="kw">loop </span>{
<a href=#333 id=333 data-nosnippet>333</a> <span class="kw">while </span>first_error.is_none()
<a href=#334 id=334 data-nosnippet>334</a> &amp;&amp; !cancellation.is_cancelled()
<a href=#335 id=335 data-nosnippet>335</a> &amp;&amp; !token.is_cancelled()
<a href=#336 id=336 data-nosnippet>336</a> &amp;&amp; active.len() &lt; parallelism.max(<span class="number">1</span>)
<a href=#337 id=337 data-nosnippet>337</a> {
<a href=#338 id=338 data-nosnippet>338</a> <span class="kw">let </span><span class="prelude-val">Some</span>(job) = jobs.next() <span class="kw">else </span>{ <span class="kw">break </span>};
<a href=#339 id=339 data-nosnippet>339</a> active.push(run(job, token.clone()));
<a href=#340 id=340 data-nosnippet>340</a> }
<a href=#341 id=341 data-nosnippet>341</a> <span class="kw">if </span>first_error.is_none() &amp;&amp; cancellation.is_cancelled() {
<a href=#342 id=342 data-nosnippet>342</a> first_error = <span class="prelude-val">Some</span>(error::DatabaseExportCancelledSnafu.build());
<a href=#343 id=343 data-nosnippet>343</a> token.cancel();
<a href=#344 id=344 data-nosnippet>344</a> }
<a href=#345 id=345 data-nosnippet>345</a> <span class="kw">if </span>active.is_empty() {
<a href=#346 id=346 data-nosnippet>346</a> <span class="kw">break</span>;
<a href=#347 id=347 data-nosnippet>347</a> }
<a href=#348 id=348 data-nosnippet>348</a> <span class="kw">let </span>result = <span class="macro">tokio::select!</span> {
<a href=#349 id=349 data-nosnippet>349</a> biased;
<a href=#350 id=350 data-nosnippet>350</a> <span class="kw">_ </span>= cancellation.cancelled(), <span class="kw">if </span>first_error.is_none() =&gt; {
<a href=#351 id=351 data-nosnippet>351</a> first_error = <span class="prelude-val">Some</span>(error::DatabaseExportCancelledSnafu.build());
<a href=#352 id=352 data-nosnippet>352</a> token.cancel();
<a href=#353 id=353 data-nosnippet>353</a> <span class="kw">continue</span>;
<a href=#354 id=354 data-nosnippet>354</a> }
<a href=#355 id=355 data-nosnippet>355</a> result = active.next() =&gt; result,
<a href=#356 id=356 data-nosnippet>356</a> };
<a href=#357 id=357 data-nosnippet>357</a> <span class="kw">match </span>result {
<a href=#358 id=358 data-nosnippet>358</a> <span class="prelude-val">Some</span>(<span class="prelude-val">Ok</span>(count)) =&gt; rows += count,
<a href=#359 id=359 data-nosnippet>359</a> <span class="prelude-val">Some</span>(<span class="prelude-val">Err</span>(err)) =&gt; {
<a href=#360 id=360 data-nosnippet>360</a> <span class="kw">if </span>!cancellation.is_cancelled() || first_error.is_none() {
<a href=#361 id=361 data-nosnippet>361</a> retain_error(<span class="kw-2">&amp;mut </span>first_error, err);
<a href=#362 id=362 data-nosnippet>362</a> }
<a href=#363 id=363 data-nosnippet>363</a> token.cancel();
<a href=#364 id=364 data-nosnippet>364</a> }
<a href=#365 id=365 data-nosnippet>365</a> <span class="prelude-val">None </span>=&gt; <span class="kw">break</span>,
<a href=#366 id=366 data-nosnippet>366</a> }
<a href=#367 id=367 data-nosnippet>367</a> }
<a href=#368 id=368 data-nosnippet>368</a> <span class="kw">match </span>first_error {
<a href=#369 id=369 data-nosnippet>369</a> <span class="prelude-val">Some</span>(err) =&gt; <span class="prelude-val">Err</span>(err),
<a href=#370 id=370 data-nosnippet>370</a> <span class="prelude-val">None </span>=&gt; <span class="prelude-val">Ok</span>(rows),
<a href=#371 id=371 data-nosnippet>371</a> }
<a href=#372 id=372 data-nosnippet>372</a>}
<a href=#373 id=373 data-nosnippet>373</a>
<a href=#374 id=374 data-nosnippet>374</a><span class="attr">#[cfg(test)]
<a href=#375 id=375 data-nosnippet>375</a></span><span class="kw">mod </span>tests {
<a href=#376 id=376 data-nosnippet>376</a> <span class="kw">use </span>std::sync::Arc;
<a href=#377 id=377 data-nosnippet>377</a> <span class="kw">use </span>std::sync::atomic::{AtomicUsize, Ordering};
<a href=#378 id=378 data-nosnippet>378</a>
<a href=#379 id=379 data-nosnippet>379</a> <span class="kw">use </span>tokio::sync::{Semaphore, mpsc};
<a href=#380 id=380 data-nosnippet>380</a>
<a href=#381 id=381 data-nosnippet>381</a> <span class="kw">use super</span>::<span class="kw-2">*</span>;
<a href=#382 id=382 data-nosnippet>382</a>
<a href=#383 id=383 data-nosnippet>383</a> <span class="attr">#[tokio::test]
<a href=#384 id=384 data-nosnippet>384</a> </span><span class="kw">async fn </span>bounded_admission_and_drain() {
<a href=#385 id=385 data-nosnippet>385</a> <span class="kw">for </span>cancel <span class="kw">in </span>[<span class="bool-val">false</span>, <span class="bool-val">true</span>] {
<a href=#386 id=386 data-nosnippet>386</a> <span class="kw">let </span>cancellation = CancellationToken::new();
<a href=#387 id=387 data-nosnippet>387</a> <span class="kw">let </span>start_error = Arc::new(Semaphore::new(<span class="number">0</span>));
<a href=#388 id=388 data-nosnippet>388</a> <span class="kw">let </span>finish_io = Arc::new(Semaphore::new(<span class="number">0</span>));
<a href=#389 id=389 data-nosnippet>389</a> <span class="kw">let </span>(started, <span class="kw-2">mut </span>receiver) = mpsc::unbounded_channel();
<a href=#390 id=390 data-nosnippet>390</a> <span class="kw">let </span>finished = Arc::new(AtomicUsize::new(<span class="number">0</span>));
<a href=#391 id=391 data-nosnippet>391</a> <span class="kw">let </span>task = tokio::spawn({
<a href=#392 id=392 data-nosnippet>392</a> <span class="kw">let </span>cancellation = cancellation.clone();
<a href=#393 id=393 data-nosnippet>393</a> <span class="kw">let </span>start_error = start_error.clone();
<a href=#394 id=394 data-nosnippet>394</a> <span class="kw">let </span>finish_io = finish_io.clone();
<a href=#395 id=395 data-nosnippet>395</a> <span class="kw">let </span>finished = finished.clone();
<a href=#396 id=396 data-nosnippet>396</a> <span class="kw">async move </span>{
<a href=#397 id=397 data-nosnippet>397</a> run_database_export_jobs(<span class="number">0</span>..<span class="number">4</span>, <span class="number">2</span>, <span class="kw-2">&amp;</span>cancellation, |job, token| {
<a href=#398 id=398 data-nosnippet>398</a> <span class="kw">let </span>started = started.clone();
<a href=#399 id=399 data-nosnippet>399</a> <span class="kw">let </span>start_error = start_error.clone();
<a href=#400 id=400 data-nosnippet>400</a> <span class="kw">let </span>finish_io = finish_io.clone();
<a href=#401 id=401 data-nosnippet>401</a> <span class="kw">let </span>finished = finished.clone();
<a href=#402 id=402 data-nosnippet>402</a> <span class="kw">async move </span>{
<a href=#403 id=403 data-nosnippet>403</a> started.send(job).unwrap();
<a href=#404 id=404 data-nosnippet>404</a> <span class="kw">if </span>job == <span class="number">0 </span>{
<a href=#405 id=405 data-nosnippet>405</a> <span class="kw">if </span>cancel {
<a href=#406 id=406 data-nosnippet>406</a> token.cancelled().<span class="kw">await</span>;
<a href=#407 id=407 data-nosnippet>407</a> } <span class="kw">else </span>{
<a href=#408 id=408 data-nosnippet>408</a> start_error.acquire().<span class="kw">await</span>.unwrap().forget();
<a href=#409 id=409 data-nosnippet>409</a> }
<a href=#410 id=410 data-nosnippet>410</a> <span class="kw">return </span>InvalidDatabaseExportSnafu {
<a href=#411 id=411 data-nosnippet>411</a> reason: <span class="string">"first error"</span>,
<a href=#412 id=412 data-nosnippet>412</a> }
<a href=#413 id=413 data-nosnippet>413</a> .fail();
<a href=#414 id=414 data-nosnippet>414</a> }
<a href=#415 id=415 data-nosnippet>415</a> <span class="comment">// Already-started ordinary I/O drains after cancellation.
<a href=#416 id=416 data-nosnippet>416</a> </span>token.cancelled().<span class="kw">await</span>;
<a href=#417 id=417 data-nosnippet>417</a> started.send(<span class="number">10</span>).unwrap();
<a href=#418 id=418 data-nosnippet>418</a> finish_io.acquire().<span class="kw">await</span>.unwrap().forget();
<a href=#419 id=419 data-nosnippet>419</a> finished.fetch_add(<span class="number">1</span>, Ordering::SeqCst);
<a href=#420 id=420 data-nosnippet>420</a> InvalidDatabaseExportSnafu {
<a href=#421 id=421 data-nosnippet>421</a> reason: <span class="string">"drain error"</span>,
<a href=#422 id=422 data-nosnippet>422</a> }
<a href=#423 id=423 data-nosnippet>423</a> .fail()
<a href=#424 id=424 data-nosnippet>424</a> }
<a href=#425 id=425 data-nosnippet>425</a> })
<a href=#426 id=426 data-nosnippet>426</a> .<span class="kw">await
<a href=#427 id=427 data-nosnippet>427</a> </span>}
<a href=#428 id=428 data-nosnippet>428</a> });
<a href=#429 id=429 data-nosnippet>429</a> <span class="macro">assert_eq!</span>(receiver.recv().<span class="kw">await</span>, <span class="prelude-val">Some</span>(<span class="number">0</span>));
<a href=#430 id=430 data-nosnippet>430</a> <span class="macro">assert_eq!</span>(receiver.recv().<span class="kw">await</span>, <span class="prelude-val">Some</span>(<span class="number">1</span>));
<a href=#431 id=431 data-nosnippet>431</a> <span class="macro">assert!</span>(receiver.try_recv().is_err());
<a href=#432 id=432 data-nosnippet>432</a> <span class="kw">if </span>cancel {
<a href=#433 id=433 data-nosnippet>433</a> cancellation.cancel();
<a href=#434 id=434 data-nosnippet>434</a> } <span class="kw">else </span>{
<a href=#435 id=435 data-nosnippet>435</a> start_error.add_permits(<span class="number">1</span>);
<a href=#436 id=436 data-nosnippet>436</a> }
<a href=#437 id=437 data-nosnippet>437</a> <span class="macro">assert_eq!</span>(receiver.recv().<span class="kw">await</span>, <span class="prelude-val">Some</span>(<span class="number">10</span>));
<a href=#438 id=438 data-nosnippet>438</a> <span class="macro">assert!</span>(!task.is_finished());
<a href=#439 id=439 data-nosnippet>439</a> finish_io.add_permits(<span class="number">1</span>);
<a href=#440 id=440 data-nosnippet>440</a> <span class="kw">let </span>err = task.<span class="kw">await</span>.unwrap().unwrap_err();
<a href=#441 id=441 data-nosnippet>441</a> <span class="kw">if </span>cancel {
<a href=#442 id=442 data-nosnippet>442</a> <span class="macro">assert!</span>(<span class="macro">matches!</span>(err, error::Error::DatabaseExportCancelled { .. }));
<a href=#443 id=443 data-nosnippet>443</a> } <span class="kw">else </span>{
<a href=#444 id=444 data-nosnippet>444</a> <span class="macro">assert!</span>(
<a href=#445 id=445 data-nosnippet>445</a> <span class="macro">matches!</span>(err, error::Error::InvalidDatabaseExport { reason } <span class="kw">if </span>reason == <span class="string">"first error"</span>)
<a href=#446 id=446 data-nosnippet>446</a> );
<a href=#447 id=447 data-nosnippet>447</a> }
<a href=#448 id=448 data-nosnippet>448</a> <span class="macro">assert_eq!</span>(finished.load(Ordering::SeqCst), <span class="number">1</span>);
<a href=#449 id=449 data-nosnippet>449</a> <span class="macro">assert_eq!</span>(receiver.recv().<span class="kw">await</span>, <span class="prelude-val">None</span>);
<a href=#450 id=450 data-nosnippet>450</a> }
<a href=#451 id=451 data-nosnippet>451</a> }
<a href=#452 id=452 data-nosnippet>452</a>
<a href=#453 id=453 data-nosnippet>453</a> <span class="attr">#[tokio::test]
<a href=#454 id=454 data-nosnippet>454</a> </span><span class="kw">async fn </span>cancellation_before_admission_and_successful_refill() {
<a href=#455 id=455 data-nosnippet>455</a> <span class="kw">let </span>token = CancellationToken::new();
<a href=#456 id=456 data-nosnippet>456</a> token.cancel();
<a href=#457 id=457 data-nosnippet>457</a> <span class="kw">let </span>result = run_database_export_jobs(<span class="number">0</span>..<span class="number">4</span>, <span class="number">2</span>, <span class="kw-2">&amp;</span>token, |<span class="kw">_</span>, <span class="kw">_</span>| <span class="kw">async </span>{
<a href=#458 id=458 data-nosnippet>458</a> <span class="macro">panic!</span>(<span class="string">"cancelled job admitted"</span>)
<a href=#459 id=459 data-nosnippet>459</a> })
<a href=#460 id=460 data-nosnippet>460</a> .<span class="kw">await</span>;
<a href=#461 id=461 data-nosnippet>461</a> <span class="macro">assert!</span>(<span class="macro">matches!</span>(
<a href=#462 id=462 data-nosnippet>462</a> result,
<a href=#463 id=463 data-nosnippet>463</a> <span class="prelude-val">Err</span>(error::Error::DatabaseExportCancelled { .. })
<a href=#464 id=464 data-nosnippet>464</a> ));
<a href=#465 id=465 data-nosnippet>465</a> <span class="kw">let </span>started = AtomicUsize::new(<span class="number">0</span>);
<a href=#466 id=466 data-nosnippet>466</a> <span class="kw">let </span>result = run_database_export_jobs(<span class="number">0</span>..<span class="number">7</span>, <span class="number">2</span>, <span class="kw-2">&amp;</span>CancellationToken::new(), |job, <span class="kw">_</span>| {
<a href=#467 id=467 data-nosnippet>467</a> started.fetch_add(<span class="number">1</span>, Ordering::SeqCst);
<a href=#468 id=468 data-nosnippet>468</a> <span class="kw">async move </span>{ <span class="prelude-val">Ok</span>(job) }
<a href=#469 id=469 data-nosnippet>469</a> })
<a href=#470 id=470 data-nosnippet>470</a> .<span class="kw">await
<a href=#471 id=471 data-nosnippet>471</a> </span>.unwrap();
<a href=#472 id=472 data-nosnippet>472</a> <span class="macro">assert_eq!</span>(result, <span class="number">21</span>);
<a href=#473 id=473 data-nosnippet>473</a> <span class="macro">assert_eq!</span>(started.load(Ordering::SeqCst), <span class="number">7</span>);
<a href=#474 id=474 data-nosnippet>474</a> }
<a href=#475 id=475 data-nosnippet>475</a>}
</code></pre></div></section></main></body></html>
@@ -0,0 +1,932 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/export_logical_tables.rs`."><title>export_logical_tables.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>export_logical_tables.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="doccomment">//! Internal physical-table export into logical-table Parquet files.
<a href=#16 id=16 data-nosnippet>16</a>//!
<a href=#17 id=17 data-nosnippet>17</a>//! Callers must authorize the captured tables and destination before invoking
<a href=#18 id=18 data-nosnippet>18</a>//! this component, and keep the source stable. Each destination must be owned
<a href=#19 id=19 data-nosnippet>19</a>//! exclusively by the export attempt. Completion and cleanup of closed files
<a href=#20 id=20 data-nosnippet>20</a>//! belong to the caller's chunk protocol, not individual file completion.
<a href=#21 id=21 data-nosnippet>21</a>
<a href=#22 id=22 data-nosnippet>22</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">mod </span>writers;
<a href=#23 id=23 data-nosnippet>23</a>
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>std::collections::{BTreeMap, BTreeSet, HashMap};
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>std::sync::Arc;
<a href=#26 id=26 data-nosnippet>26</a>
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>arrow::array::{Array, ArrayRef, AsArray, ListArray, StructArray, UInt32Array};
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>arrow::buffer::OffsetBuffer;
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span>arrow::compute::{cast, take};
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span>arrow::datatypes::{DataType, SchemaRef};
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span>arrow::downcast_dictionary_array;
<a href=#32 id=32 data-nosnippet>32</a><span class="kw">use </span>arrow::record_batch::RecordBatch;
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">use </span>common_datasource::object_store::build_backend_for_write;
<a href=#34 id=34 data-nosnippet>34</a><span class="kw">use </span>common_datasource::packed_writer::{PackedTableWriter, PackedWriterRef};
<a href=#35 id=35 data-nosnippet>35</a><span class="kw">use </span>common_datasource::parquet_writer::{
<a href=#36 id=36 data-nosnippet>36</a> ParquetCreationPolicy, ParquetFileWriter, ParquetWriterLimits,
<a href=#37 id=37 data-nosnippet>37</a>};
<a href=#38 id=38 data-nosnippet>38</a><span class="kw">use </span>common_meta::key::table_route::{TableRouteManager, TableRouteValue};
<a href=#39 id=39 data-nosnippet>39</a><span class="kw">use </span>common_query::OutputData;
<a href=#40 id=40 data-nosnippet>40</a><span class="kw">use </span>common_recordbatch::SendableRecordBatchStream;
<a href=#41 id=41 data-nosnippet>41</a><span class="kw">use </span>common_time::range::TimestampRange;
<a href=#42 id=42 data-nosnippet>42</a><span class="kw">use </span>datafusion::datasource::DefaultTableSource;
<a href=#43 id=43 data-nosnippet>43</a><span class="kw">use </span>datafusion_common::TableReference <span class="kw">as </span>DfTableReference;
<a href=#44 id=44 data-nosnippet>44</a><span class="kw">use </span>datafusion_expr::{LogicalPlan, LogicalPlanBuilder, col, lit};
<a href=#45 id=45 data-nosnippet>45</a><span class="kw">use </span>futures::StreamExt;
<a href=#46 id=46 data-nosnippet>46</a><span class="kw">use </span>object_store::ObjectStore;
<a href=#47 id=47 data-nosnippet>47</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#48 id=48 data-nosnippet>48</a><span class="kw">use </span>snafu::{IntoError, OptionExt, ResultExt, ensure};
<a href=#49 id=49 data-nosnippet>49</a><span class="kw">use </span>store_api::metric_engine_consts::{
<a href=#50 id=50 data-nosnippet>50</a> DATA_SCHEMA_TABLE_ID_COLUMN_NAME <span class="kw">as </span>TABLE_ID, DATA_SCHEMA_TSID_COLUMN_NAME <span class="kw">as </span>TSID,
<a href=#51 id=51 data-nosnippet>51</a> LOGICAL_TABLE_METADATA_KEY, METRIC_ENGINE_NAME, PHYSICAL_TABLE_METADATA_KEY,
<a href=#52 id=52 data-nosnippet>52</a>};
<a href=#53 id=53 data-nosnippet>53</a><span class="kw">use </span>table::TableRef;
<a href=#54 id=54 data-nosnippet>54</a><span class="kw">use </span>table::metadata::TableId;
<a href=#55 id=55 data-nosnippet>55</a><span class="kw">use </span>table::table::adapter::DfTableProviderAdapter;
<a href=#56 id=56 data-nosnippet>56</a><span class="kw">use </span>tokio::sync::OwnedSemaphorePermit;
<a href=#57 id=57 data-nosnippet>57</a><span class="kw">use </span>tokio_util::sync::CancellationToken;
<a href=#58 id=58 data-nosnippet>58</a>
<a href=#59 id=59 data-nosnippet>59</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, InvalidLogicalTableExportSnafu, LogicalTableExportResourceSnafu, <span class="prelude-ty">Result</span>};
<a href=#60 id=60 data-nosnippet>60</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#61 id=61 data-nosnippet>61</a><span class="kw">use </span><span class="kw">crate</span>::statement::database_copy::DatabaseExportFile;
<a href=#62 id=62 data-nosnippet>62</a><span class="kw">use </span><span class="kw">crate</span>::statement::export_logical_tables::writers::{ExportWriteBudget, Payload, TableWriters};
<a href=#63 id=63 data-nosnippet>63</a>
<a href=#64 id=64 data-nosnippet>64</a><span class="doccomment">/// Export preprocessing and per-file limits. Query memory and spill remain
<a href=#65 id=65 data-nosnippet>65</a>/// governed by the query engine.
<a href=#66 id=66 data-nosnippet>66</a>/// Batch checks happen after allocation; flush thresholds are not hard RSS caps.
<a href=#67 id=67 data-nosnippet>67</a></span><span class="attr">#[derive(Clone, Copy, Debug)]
<a href=#68 id=68 data-nosnippet>68</a></span><span class="kw">pub struct </span>LogicalTableExportLimits {
<a href=#69 id=69 data-nosnippet>69</a> <span class="doccomment">/// Maximum retained scan batch before routing.
<a href=#70 id=70 data-nosnippet>70</a> </span><span class="kw">pub </span>input_batch_bytes: usize,
<a href=#71 id=71 data-nosnippet>71</a> <span class="doccomment">/// Maximum estimated expanded logical values in one conversion slice.
<a href=#72 id=72 data-nosnippet>72</a> </span><span class="kw">pub </span>conversion_bytes: usize,
<a href=#73 id=73 data-nosnippet>73</a> <span class="doccomment">/// Limits owned by the single-file Parquet writer.
<a href=#74 id=74 data-nosnippet>74</a> </span><span class="kw">pub </span>writer: ParquetWriterLimits,
<a href=#75 id=75 data-nosnippet>75</a>}
<a href=#76 id=76 data-nosnippet>76</a>
<a href=#77 id=77 data-nosnippet>77</a><span class="kw">impl </span>Default <span class="kw">for </span>LogicalTableExportLimits {
<a href=#78 id=78 data-nosnippet>78</a> <span class="kw">fn </span>default() -&gt; <span class="self">Self </span>{
<a href=#79 id=79 data-nosnippet>79</a> <span class="self">Self </span>{
<a href=#80 id=80 data-nosnippet>80</a> input_batch_bytes: <span class="number">64 </span>* <span class="number">1024 </span>* <span class="number">1024</span>,
<a href=#81 id=81 data-nosnippet>81</a> conversion_bytes: <span class="number">1024 </span>* <span class="number">1024</span>,
<a href=#82 id=82 data-nosnippet>82</a> writer: ParquetWriterLimits {
<a href=#83 id=83 data-nosnippet>83</a> row_group_rows: <span class="number">1024 </span>* <span class="number">1024</span>,
<a href=#84 id=84 data-nosnippet>84</a> flush_threshold_bytes: <span class="number">8 </span>* <span class="number">1024 </span>* <span class="number">1024</span>,
<a href=#85 id=85 data-nosnippet>85</a> max_row_groups: <span class="number">4096</span>,
<a href=#86 id=86 data-nosnippet>86</a> },
<a href=#87 id=87 data-nosnippet>87</a> }
<a href=#88 id=88 data-nosnippet>88</a> }
<a href=#89 id=89 data-nosnippet>89</a>}
<a href=#90 id=90 data-nosnippet>90</a>
<a href=#91 id=91 data-nosnippet>91</a><span class="kw">impl </span>LogicalTableExportLimits {
<a href=#92 id=92 data-nosnippet>92</a> <span class="kw">fn </span>validate(<span class="self">self</span>) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#93 id=93 data-nosnippet>93</a> <span class="macro">ensure!</span>(
<a href=#94 id=94 data-nosnippet>94</a> <span class="self">self</span>.input_batch_bytes &gt; <span class="number">0 </span>&amp;&amp; <span class="self">self</span>.conversion_bytes &gt; <span class="number">0</span>,
<a href=#95 id=95 data-nosnippet>95</a> InvalidLogicalTableExportSnafu {
<a href=#96 id=96 data-nosnippet>96</a> reason: <span class="string">"export limits must be positive"
<a href=#97 id=97 data-nosnippet>97</a> </span>}
<a href=#98 id=98 data-nosnippet>98</a> );
<a href=#99 id=99 data-nosnippet>99</a> <span class="prelude-val">Ok</span>(())
<a href=#100 id=100 data-nosnippet>100</a> }
<a href=#101 id=101 data-nosnippet>101</a>}
<a href=#102 id=102 data-nosnippet>102</a>
<a href=#103 id=103 data-nosnippet>103</a><span class="doccomment">/// Metadata captured for one physical table and its selected logical tables.
<a href=#104 id=104 data-nosnippet>104</a>/// Construct one unit per physical table, schema and time chunk. Capturing these
<a href=#105 id=105 data-nosnippet>105</a>/// references supplies no snapshot isolation or locking guarantee.
<a href=#106 id=106 data-nosnippet>106</a></span><span class="kw">pub struct </span>LogicalTableExport {
<a href=#107 id=107 data-nosnippet>107</a> physical_table: TableRef,
<a href=#108 id=108 data-nosnippet>108</a> scan_projection: Vec&lt;usize&gt;,
<a href=#109 id=109 data-nosnippet>109</a> logical_tables: BTreeMap&lt;TableId, LogicalTableProjection&gt;,
<a href=#110 id=110 data-nosnippet>110</a>}
<a href=#111 id=111 data-nosnippet>111</a>
<a href=#112 id=112 data-nosnippet>112</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>LogicalTableProjection {
<a href=#113 id=113 data-nosnippet>113</a> name: String,
<a href=#114 id=114 data-nosnippet>114</a> output: DatabaseExportFile,
<a href=#115 id=115 data-nosnippet>115</a> schema: SchemaRef,
<a href=#116 id=116 data-nosnippet>116</a> projection: Vec&lt;usize&gt;,
<a href=#117 id=117 data-nosnippet>117</a>}
<a href=#118 id=118 data-nosnippet>118</a>
<a href=#119 id=119 data-nosnippet>119</a><span class="kw">impl </span>LogicalTableExport {
<a href=#120 id=120 data-nosnippet>120</a> <span class="doccomment">/// Capture schemas from selected Metric table references.
<a href=#121 id=121 data-nosnippet>121</a> /// Export validates their physical-table association against table routes.
<a href=#122 id=122 data-nosnippet>122</a> </span><span class="kw">pub fn </span>try_new(physical: TableRef, tables: <span class="kw-2">&amp;</span>[TableRef]) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="self">Self</span>&gt; {
<a href=#123 id=123 data-nosnippet>123</a> <span class="self">Self</span>::try_new_in_directory(physical, tables, <span class="string">""</span>)
<a href=#124 id=124 data-nosnippet>124</a> }
<a href=#125 id=125 data-nosnippet>125</a>
<a href=#126 id=126 data-nosnippet>126</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>try_new_in_directory(
<a href=#127 id=127 data-nosnippet>127</a> physical: TableRef,
<a href=#128 id=128 data-nosnippet>128</a> tables: <span class="kw-2">&amp;</span>[TableRef],
<a href=#129 id=129 data-nosnippet>129</a> directory: <span class="kw-2">&amp;</span>str,
<a href=#130 id=130 data-nosnippet>130</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="self">Self</span>&gt; {
<a href=#131 id=131 data-nosnippet>131</a> <span class="kw">let </span>physical_info = physical.table_info();
<a href=#132 id=132 data-nosnippet>132</a> <span class="macro">ensure!</span>(
<a href=#133 id=133 data-nosnippet>133</a> physical_info.meta.engine == METRIC_ENGINE_NAME
<a href=#134 id=134 data-nosnippet>134</a> &amp;&amp; physical_info
<a href=#135 id=135 data-nosnippet>135</a> .meta
<a href=#136 id=136 data-nosnippet>136</a> .options
<a href=#137 id=137 data-nosnippet>137</a> .extra_options
<a href=#138 id=138 data-nosnippet>138</a> .contains_key(PHYSICAL_TABLE_METADATA_KEY)
<a href=#139 id=139 data-nosnippet>139</a> &amp;&amp; !tables.is_empty(),
<a href=#140 id=140 data-nosnippet>140</a> InvalidLogicalTableExportSnafu {
<a href=#141 id=141 data-nosnippet>141</a> reason: <span class="string">"expected a Metric physical table and selected logical tables"
<a href=#142 id=142 data-nosnippet>142</a> </span>}
<a href=#143 id=143 data-nosnippet>143</a> );
<a href=#144 id=144 data-nosnippet>144</a> <span class="kw">let </span>physical_schema = physical.schema();
<a href=#145 id=145 data-nosnippet>145</a> <span class="kw">let </span>id_index = physical_schema.column_index_by_name(TABLE_ID).context(
<a href=#146 id=146 data-nosnippet>146</a> InvalidLogicalTableExportSnafu {
<a href=#147 id=147 data-nosnippet>147</a> reason: <span class="string">"physical schema has no __table_id"</span>,
<a href=#148 id=148 data-nosnippet>148</a> },
<a href=#149 id=149 data-nosnippet>149</a> )<span class="question-mark">?</span>;
<a href=#150 id=150 data-nosnippet>150</a> <span class="kw">let </span><span class="kw-2">mut </span>scan_projection = BTreeSet::from([id_index]);
<a href=#151 id=151 data-nosnippet>151</a> <span class="kw">let </span><span class="kw-2">mut </span>logical_tables = BTreeMap::new();
<a href=#152 id=152 data-nosnippet>152</a> <span class="kw">for </span>table <span class="kw">in </span>tables {
<a href=#153 id=153 data-nosnippet>153</a> <span class="kw">let </span>info = table.table_info();
<a href=#154 id=154 data-nosnippet>154</a> <span class="kw">let </span>name = <span class="kw-2">&amp;</span>info.name;
<a href=#155 id=155 data-nosnippet>155</a> <span class="macro">ensure!</span>(
<a href=#156 id=156 data-nosnippet>156</a> info.catalog_name == physical_info.catalog_name
<a href=#157 id=157 data-nosnippet>157</a> &amp;&amp; info.schema_name == physical_info.schema_name
<a href=#158 id=158 data-nosnippet>158</a> &amp;&amp; info.meta.engine == METRIC_ENGINE_NAME
<a href=#159 id=159 data-nosnippet>159</a> &amp;&amp; info
<a href=#160 id=160 data-nosnippet>160</a> .meta
<a href=#161 id=161 data-nosnippet>161</a> .options
<a href=#162 id=162 data-nosnippet>162</a> .extra_options
<a href=#163 id=163 data-nosnippet>163</a> .contains_key(LOGICAL_TABLE_METADATA_KEY),
<a href=#164 id=164 data-nosnippet>164</a> InvalidLogicalTableExportSnafu {
<a href=#165 id=165 data-nosnippet>165</a> reason: <span class="macro">format!</span>(
<a href=#166 id=166 data-nosnippet>166</a> <span class="string">"{name} is not a Metric logical table in the physical table schema"
<a href=#167 id=167 data-nosnippet>167</a> </span>)
<a href=#168 id=168 data-nosnippet>168</a> }
<a href=#169 id=169 data-nosnippet>169</a> );
<a href=#170 id=170 data-nosnippet>170</a> <span class="kw">let </span>schema = table.schema().arrow_schema().clone();
<a href=#171 id=171 data-nosnippet>171</a> <span class="kw">let </span>indices = schema
<a href=#172 id=172 data-nosnippet>172</a> .fields()
<a href=#173 id=173 data-nosnippet>173</a> .iter()
<a href=#174 id=174 data-nosnippet>174</a> .map(|field| {
<a href=#175 id=175 data-nosnippet>175</a> <span class="macro">ensure!</span>(
<a href=#176 id=176 data-nosnippet>176</a> field.name() != TABLE_ID &amp;&amp; field.name() != TSID,
<a href=#177 id=177 data-nosnippet>177</a> InvalidLogicalTableExportSnafu {
<a href=#178 id=178 data-nosnippet>178</a> reason: <span class="string">"logical schema contains internal Metric columns"
<a href=#179 id=179 data-nosnippet>179</a> </span>}
<a href=#180 id=180 data-nosnippet>180</a> );
<a href=#181 id=181 data-nosnippet>181</a> <span class="kw">let </span>index = physical_schema
<a href=#182 id=182 data-nosnippet>182</a> .column_index_by_name(field.name())
<a href=#183 id=183 data-nosnippet>183</a> .with_context(|| InvalidLogicalTableExportSnafu {
<a href=#184 id=184 data-nosnippet>184</a> reason: <span class="macro">format!</span>(<span class="string">"physical schema lacks {}"</span>, field.name()),
<a href=#185 id=185 data-nosnippet>185</a> })<span class="question-mark">?</span>;
<a href=#186 id=186 data-nosnippet>186</a> <span class="macro">ensure!</span>(
<a href=#187 id=187 data-nosnippet>187</a> physical_schema.arrow_schema().field(index).data_type()
<a href=#188 id=188 data-nosnippet>188</a> == field.data_type(),
<a href=#189 id=189 data-nosnippet>189</a> InvalidLogicalTableExportSnafu {
<a href=#190 id=190 data-nosnippet>190</a> reason: <span class="macro">format!</span>(<span class="string">"physical/logical type mismatch for {}"</span>, field.name())
<a href=#191 id=191 data-nosnippet>191</a> }
<a href=#192 id=192 data-nosnippet>192</a> );
<a href=#193 id=193 data-nosnippet>193</a> <span class="prelude-val">Ok</span>(index)
<a href=#194 id=194 data-nosnippet>194</a> })
<a href=#195 id=195 data-nosnippet>195</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#196 id=196 data-nosnippet>196</a> scan_projection.extend(indices.iter().copied());
<a href=#197 id=197 data-nosnippet>197</a> <span class="macro">ensure!</span>(
<a href=#198 id=198 data-nosnippet>198</a> logical_tables
<a href=#199 id=199 data-nosnippet>199</a> .insert(
<a href=#200 id=200 data-nosnippet>200</a> info.table_id(),
<a href=#201 id=201 data-nosnippet>201</a> LogicalTableProjection {
<a href=#202 id=202 data-nosnippet>202</a> name: name.clone(),
<a href=#203 id=203 data-nosnippet>203</a> output: DatabaseExportFile::new(directory, name, <span class="string">".parquet"</span>)<span class="question-mark">?</span>,
<a href=#204 id=204 data-nosnippet>204</a> schema,
<a href=#205 id=205 data-nosnippet>205</a> projection: indices,
<a href=#206 id=206 data-nosnippet>206</a> }
<a href=#207 id=207 data-nosnippet>207</a> )
<a href=#208 id=208 data-nosnippet>208</a> .is_none(),
<a href=#209 id=209 data-nosnippet>209</a> InvalidLogicalTableExportSnafu {
<a href=#210 id=210 data-nosnippet>210</a> reason: <span class="string">"duplicate logical table"
<a href=#211 id=211 data-nosnippet>211</a> </span>}
<a href=#212 id=212 data-nosnippet>212</a> );
<a href=#213 id=213 data-nosnippet>213</a> }
<a href=#214 id=214 data-nosnippet>214</a> <span class="kw">let </span>scan_projection = scan_projection.into_iter().collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#215 id=215 data-nosnippet>215</a> <span class="kw">for </span>file <span class="kw">in </span>logical_tables.values_mut() {
<a href=#216 id=216 data-nosnippet>216</a> <span class="kw">for </span>index <span class="kw">in </span><span class="kw-2">&amp;mut </span>file.projection {
<a href=#217 id=217 data-nosnippet>217</a> <span class="kw-2">*</span>index = scan_projection.binary_search(index).map_err(|<span class="kw">_</span>| {
<a href=#218 id=218 data-nosnippet>218</a> error::UnexpectedSnafu {
<a href=#219 id=219 data-nosnippet>219</a> violated: <span class="string">"logical column missing from physical projection"</span>,
<a href=#220 id=220 data-nosnippet>220</a> }
<a href=#221 id=221 data-nosnippet>221</a> .build()
<a href=#222 id=222 data-nosnippet>222</a> })<span class="question-mark">?</span>;
<a href=#223 id=223 data-nosnippet>223</a> }
<a href=#224 id=224 data-nosnippet>224</a> }
<a href=#225 id=225 data-nosnippet>225</a> <span class="prelude-val">Ok</span>(<span class="self">Self </span>{
<a href=#226 id=226 data-nosnippet>226</a> physical_table: physical,
<a href=#227 id=227 data-nosnippet>227</a> scan_projection,
<a href=#228 id=228 data-nosnippet>228</a> logical_tables,
<a href=#229 id=229 data-nosnippet>229</a> })
<a href=#230 id=230 data-nosnippet>230</a> }
<a href=#231 id=231 data-nosnippet>231</a>
<a href=#232 id=232 data-nosnippet>232</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>output_files(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; <span class="kw">impl </span>Iterator&lt;Item = <span class="kw-2">&amp;</span>DatabaseExportFile&gt; {
<a href=#233 id=233 data-nosnippet>233</a> <span class="self">self</span>.logical_tables.values().map(|table| <span class="kw-2">&amp;</span>table.output)
<a href=#234 id=234 data-nosnippet>234</a> }
<a href=#235 id=235 data-nosnippet>235</a>
<a href=#236 id=236 data-nosnippet>236</a> <span class="kw">async fn </span>validate_table_routes(<span class="kw-2">&amp;</span><span class="self">self</span>, manager: <span class="kw-2">&amp;</span>TableRouteManager) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#237 id=237 data-nosnippet>237</a> <span class="kw">let </span>table_ids = <span class="self">self</span>.logical_tables.keys().copied().collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#238 id=238 data-nosnippet>238</a> <span class="kw">let </span>routes = manager
<a href=#239 id=239 data-nosnippet>239</a> .table_route_storage()
<a href=#240 id=240 data-nosnippet>240</a> .batch_get(<span class="kw-2">&amp;</span>table_ids)
<a href=#241 id=241 data-nosnippet>241</a> .<span class="kw">await
<a href=#242 id=242 data-nosnippet>242</a> </span>.context(error::TableMetadataManagerSnafu)<span class="question-mark">?</span>;
<a href=#243 id=243 data-nosnippet>243</a> <span class="kw">let </span>physical_table_id = <span class="self">self</span>.physical_table.table_info().table_id();
<a href=#244 id=244 data-nosnippet>244</a> <span class="kw">for </span>(table_id, route) <span class="kw">in </span>table_ids.into_iter().zip(routes) {
<a href=#245 id=245 data-nosnippet>245</a> <span class="macro">ensure!</span>(
<a href=#246 id=246 data-nosnippet>246</a> <span class="macro">matches!</span>(route, <span class="prelude-val">Some</span>(TableRouteValue::Logical(route)) <span class="kw">if </span>route.physical_table_id() == physical_table_id),
<a href=#247 id=247 data-nosnippet>247</a> InvalidLogicalTableExportSnafu {
<a href=#248 id=248 data-nosnippet>248</a> reason: <span class="macro">format!</span>(
<a href=#249 id=249 data-nosnippet>249</a> <span class="string">"logical table {table_id} does not belong to physical table {physical_table_id}"
<a href=#250 id=250 data-nosnippet>250</a> </span>)
<a href=#251 id=251 data-nosnippet>251</a> }
<a href=#252 id=252 data-nosnippet>252</a> );
<a href=#253 id=253 data-nosnippet>253</a> }
<a href=#254 id=254 data-nosnippet>254</a> <span class="prelude-val">Ok</span>(())
<a href=#255 id=255 data-nosnippet>255</a> }
<a href=#256 id=256 data-nosnippet>256</a>
<a href=#257 id=257 data-nosnippet>257</a> <span class="kw">fn </span>build_plan(<span class="kw-2">&amp;</span><span class="self">self</span>, time_range: <span class="prelude-ty">Option</span>&lt;<span class="kw-2">&amp;</span>TimestampRange&gt;) -&gt; <span class="prelude-ty">Result</span>&lt;LogicalPlan&gt; {
<a href=#258 id=258 data-nosnippet>258</a> <span class="kw">let </span>info = <span class="self">self</span>.physical_table.table_info();
<a href=#259 id=259 data-nosnippet>259</a> <span class="kw">let </span><span class="kw-2">mut </span>filters = <span class="self">self
<a href=#260 id=260 data-nosnippet>260</a> </span>.physical_table
<a href=#261 id=261 data-nosnippet>261</a> .schema()
<a href=#262 id=262 data-nosnippet>262</a> .timestamp_column()
<a href=#263 id=263 data-nosnippet>263</a> .and_then(|column| {
<a href=#264 id=264 data-nosnippet>264</a> common_query::logical_plan::build_filter_from_timestamp(<span class="kw-2">&amp;</span>column.name, time_range)
<a href=#265 id=265 data-nosnippet>265</a> })
<a href=#266 id=266 data-nosnippet>266</a> .into_iter()
<a href=#267 id=267 data-nosnippet>267</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#268 id=268 data-nosnippet>268</a> <span class="kw">if </span><span class="self">self</span>.logical_tables.len() == <span class="number">1
<a href=#269 id=269 data-nosnippet>269</a> </span>&amp;&amp; <span class="kw">let </span><span class="prelude-val">Some</span>((<span class="kw-2">&amp;</span>id, <span class="kw">_</span>)) = <span class="self">self</span>.logical_tables.first_key_value()
<a href=#270 id=270 data-nosnippet>270</a> {
<a href=#271 id=271 data-nosnippet>271</a> filters.push(col(TABLE_ID).eq(lit(id)));
<a href=#272 id=272 data-nosnippet>272</a> }
<a href=#273 id=273 data-nosnippet>273</a> <span class="kw">let </span>source = Arc::new(DefaultTableSource::new(Arc::new(
<a href=#274 id=274 data-nosnippet>274</a> DfTableProviderAdapter::new(<span class="self">self</span>.physical_table.clone()),
<a href=#275 id=275 data-nosnippet>275</a> )));
<a href=#276 id=276 data-nosnippet>276</a> <span class="kw">let </span><span class="kw-2">mut </span>builder = LogicalPlanBuilder::scan_with_filters(
<a href=#277 id=277 data-nosnippet>277</a> DfTableReference::full(
<a href=#278 id=278 data-nosnippet>278</a> info.catalog_name.clone(),
<a href=#279 id=279 data-nosnippet>279</a> info.schema_name.clone(),
<a href=#280 id=280 data-nosnippet>280</a> info.name.clone(),
<a href=#281 id=281 data-nosnippet>281</a> ),
<a href=#282 id=282 data-nosnippet>282</a> source,
<a href=#283 id=283 data-nosnippet>283</a> <span class="prelude-val">Some</span>(<span class="self">self</span>.scan_projection.clone()),
<a href=#284 id=284 data-nosnippet>284</a> filters.clone(),
<a href=#285 id=285 data-nosnippet>285</a> )
<a href=#286 id=286 data-nosnippet>286</a> .context(error::BuildDfLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#287 id=287 data-nosnippet>287</a> <span class="kw">for </span>filter <span class="kw">in </span>filters {
<a href=#288 id=288 data-nosnippet>288</a> builder = builder
<a href=#289 id=289 data-nosnippet>289</a> .filter(filter)
<a href=#290 id=290 data-nosnippet>290</a> .context(error::BuildDfLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#291 id=291 data-nosnippet>291</a> }
<a href=#292 id=292 data-nosnippet>292</a> <span class="comment">// A single selected table has a constant routing key after filtering.
<a href=#293 id=293 data-nosnippet>293</a> </span><span class="kw">if </span><span class="self">self</span>.logical_tables.len() &gt; <span class="number">1 </span>{
<a href=#294 id=294 data-nosnippet>294</a> builder = builder
<a href=#295 id=295 data-nosnippet>295</a> .sort(<span class="macro">vec!</span>[col(TABLE_ID).sort(<span class="bool-val">true</span>, <span class="bool-val">false</span>)])
<a href=#296 id=296 data-nosnippet>296</a> .context(error::BuildDfLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#297 id=297 data-nosnippet>297</a> }
<a href=#298 id=298 data-nosnippet>298</a> builder.build().context(error::BuildDfLogicalPlanSnafu)
<a href=#299 id=299 data-nosnippet>299</a> }
<a href=#300 id=300 data-nosnippet>300</a>}
<a href=#301 id=301 data-nosnippet>301</a>
<a href=#302 id=302 data-nosnippet>302</a><span class="doccomment">/// Counts returned only after every selected logical file has been closed.
<a href=#303 id=303 data-nosnippet>303</a></span><span class="attr">#[derive(Debug, Default, PartialEq, Eq)]
<a href=#304 id=304 data-nosnippet>304</a></span><span class="kw">pub struct </span>LogicalTableExportSummary {
<a href=#305 id=305 data-nosnippet>305</a> <span class="kw">pub </span>rows: usize,
<a href=#306 id=306 data-nosnippet>306</a> <span class="kw">pub </span>skipped_rows: usize,
<a href=#307 id=307 data-nosnippet>307</a> <span class="kw">pub </span>files: usize,
<a href=#308 id=308 data-nosnippet>308</a>}
<a href=#309 id=309 data-nosnippet>309</a>
<a href=#310 id=310 data-nosnippet>310</a><span class="kw">impl </span>StatementExecutor {
<a href=#311 id=311 data-nosnippet>311</a> <span class="doccomment">/// Export selected logical tables to a fresh directory using one physical query.
<a href=#312 id=312 data-nosnippet>312</a> /// The caller bounds concurrent units and owns retry/cleanup of this directory.
<a href=#313 id=313 data-nosnippet>313</a> /// Cancellation waits for in-flight I/O before aborting the active upload.
<a href=#314 id=314 data-nosnippet>314</a> /// Closed files remain an incomplete chunk until the caller publishes its
<a href=#315 id=315 data-nosnippet>315</a> /// completion state.
<a href=#316 id=316 data-nosnippet>316</a> </span><span class="attr">#[allow(clippy::too_many_arguments)]
<a href=#317 id=317 data-nosnippet>317</a> </span><span class="kw">pub async fn </span>export_logical_tables(
<a href=#318 id=318 data-nosnippet>318</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#319 id=319 data-nosnippet>319</a> unit: <span class="kw-2">&amp;</span>LogicalTableExport,
<a href=#320 id=320 data-nosnippet>320</a> directory: <span class="kw-2">&amp;</span>str,
<a href=#321 id=321 data-nosnippet>321</a> connection: <span class="kw-2">&amp;</span>HashMap&lt;String, String&gt;,
<a href=#322 id=322 data-nosnippet>322</a> time_range: <span class="prelude-ty">Option</span>&lt;<span class="kw-2">&amp;</span>TimestampRange&gt;,
<a href=#323 id=323 data-nosnippet>323</a> limits: LogicalTableExportLimits,
<a href=#324 id=324 data-nosnippet>324</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#325 id=325 data-nosnippet>325</a> query_ctx: QueryContextRef,
<a href=#326 id=326 data-nosnippet>326</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;LogicalTableExportSummary&gt; {
<a href=#327 id=327 data-nosnippet>327</a> <span class="self">self</span>.export_logical_tables_managed(
<a href=#328 id=328 data-nosnippet>328</a> unit,
<a href=#329 id=329 data-nosnippet>329</a> directory,
<a href=#330 id=330 data-nosnippet>330</a> connection,
<a href=#331 id=331 data-nosnippet>331</a> time_range,
<a href=#332 id=332 data-nosnippet>332</a> limits,
<a href=#333 id=333 data-nosnippet>333</a> <span class="kw-2">&amp;</span>cancellation.child_token(),
<a href=#334 id=334 data-nosnippet>334</a> query_ctx,
<a href=#335 id=335 data-nosnippet>335</a> ExportWriteBudget::new(<span class="number">1</span>),
<a href=#336 id=336 data-nosnippet>336</a> <span class="prelude-val">None</span>,
<a href=#337 id=337 data-nosnippet>337</a> )
<a href=#338 id=338 data-nosnippet>338</a> .<span class="kw">await
<a href=#339 id=339 data-nosnippet>339</a> </span>}
<a href=#340 id=340 data-nosnippet>340</a>
<a href=#341 id=341 data-nosnippet>341</a> <span class="attr">#[allow(clippy::too_many_arguments)]
<a href=#342 id=342 data-nosnippet>342</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>export_logical_tables_managed(
<a href=#343 id=343 data-nosnippet>343</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#344 id=344 data-nosnippet>344</a> unit: <span class="kw-2">&amp;</span>LogicalTableExport,
<a href=#345 id=345 data-nosnippet>345</a> directory: <span class="kw-2">&amp;</span>str,
<a href=#346 id=346 data-nosnippet>346</a> connection: <span class="kw-2">&amp;</span>HashMap&lt;String, String&gt;,
<a href=#347 id=347 data-nosnippet>347</a> time_range: <span class="prelude-ty">Option</span>&lt;<span class="kw-2">&amp;</span>TimestampRange&gt;,
<a href=#348 id=348 data-nosnippet>348</a> limits: LogicalTableExportLimits,
<a href=#349 id=349 data-nosnippet>349</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#350 id=350 data-nosnippet>350</a> query_ctx: QueryContextRef,
<a href=#351 id=351 data-nosnippet>351</a> budget: Arc&lt;ExportWriteBudget&gt;,
<a href=#352 id=352 data-nosnippet>352</a> packed: <span class="prelude-ty">Option</span>&lt;PackedWriterRef&gt;,
<a href=#353 id=353 data-nosnippet>353</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;LogicalTableExportSummary&gt; {
<a href=#354 id=354 data-nosnippet>354</a> limits.validate()<span class="question-mark">?</span>;
<a href=#355 id=355 data-nosnippet>355</a> <span class="kw">let </span>(store, stream) = <span class="macro">tokio::select!</span> {
<a href=#356 id=356 data-nosnippet>356</a> biased;
<a href=#357 id=357 data-nosnippet>357</a> <span class="kw">_ </span>= cancellation.cancelled() =&gt; <span class="kw">return </span>error::LogicalTableExportCancelledSnafu.fail(),
<a href=#358 id=358 data-nosnippet>358</a> result = <span class="kw">async </span>{
<a href=#359 id=359 data-nosnippet>359</a> unit.validate_table_routes(<span class="self">self</span>.table_metadata_manager.table_route_manager()).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#360 id=360 data-nosnippet>360</a> <span class="kw">let </span>store = build_backend_for_write(<span class="kw-2">&amp;</span><span class="macro">format!</span>(<span class="string">"{}/"</span>, directory.trim_end_matches(<span class="string">'/'</span>)), connection, <span class="kw-2">&amp;</span><span class="self">self</span>.local_file_access)
<a href=#361 id=361 data-nosnippet>361</a> .<span class="kw">await</span>.context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#362 id=362 data-nosnippet>362</a> <span class="kw">let </span>output = <span class="self">self</span>.query_engine.execute(unit.build_plan(time_range)<span class="question-mark">?</span>, query_ctx)
<a href=#363 id=363 data-nosnippet>363</a> .<span class="kw">await</span>.context(error::ExecLogicalPlanSnafu)<span class="question-mark">?</span>;
<a href=#364 id=364 data-nosnippet>364</a> <span class="kw">let </span>stream = <span class="kw">match </span>output.data {
<a href=#365 id=365 data-nosnippet>365</a> OutputData::Stream(stream) =&gt; stream,
<a href=#366 id=366 data-nosnippet>366</a> OutputData::RecordBatches(batches) =&gt; batches.as_stream(),
<a href=#367 id=367 data-nosnippet>367</a> <span class="kw">_ </span>=&gt; <span class="kw">return </span>error::UnexpectedSnafu { violated: <span class="string">"expected physical query rows" </span>}.fail(),
<a href=#368 id=368 data-nosnippet>368</a> };
<a href=#369 id=369 data-nosnippet>369</a> <span class="prelude-val">Ok</span>((store, stream))
<a href=#370 id=370 data-nosnippet>370</a> } =&gt; result<span class="question-mark">?</span>,
<a href=#371 id=371 data-nosnippet>371</a> };
<a href=#372 id=372 data-nosnippet>372</a> export_stream_managed(unit, stream, <span class="kw-2">&amp;</span>store, limits, cancellation, budget, packed).<span class="kw">await
<a href=#373 id=373 data-nosnippet>373</a> </span>}
<a href=#374 id=374 data-nosnippet>374</a>}
<a href=#375 id=375 data-nosnippet>375</a>
<a href=#376 id=376 data-nosnippet>376</a><span class="attr">#[cfg(test)]
<a href=#377 id=377 data-nosnippet>377</a></span><span class="kw">async fn </span>export_stream(
<a href=#378 id=378 data-nosnippet>378</a> unit: <span class="kw-2">&amp;</span>LogicalTableExport,
<a href=#379 id=379 data-nosnippet>379</a> stream: SendableRecordBatchStream,
<a href=#380 id=380 data-nosnippet>380</a> store: <span class="kw-2">&amp;</span>ObjectStore,
<a href=#381 id=381 data-nosnippet>381</a> limits: LogicalTableExportLimits,
<a href=#382 id=382 data-nosnippet>382</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#383 id=383 data-nosnippet>383</a>) -&gt; <span class="prelude-ty">Result</span>&lt;LogicalTableExportSummary&gt; {
<a href=#384 id=384 data-nosnippet>384</a> export_stream_managed(
<a href=#385 id=385 data-nosnippet>385</a> unit,
<a href=#386 id=386 data-nosnippet>386</a> stream,
<a href=#387 id=387 data-nosnippet>387</a> store,
<a href=#388 id=388 data-nosnippet>388</a> limits,
<a href=#389 id=389 data-nosnippet>389</a> <span class="kw-2">&amp;</span>cancellation.child_token(),
<a href=#390 id=390 data-nosnippet>390</a> ExportWriteBudget::new(<span class="number">1</span>),
<a href=#391 id=391 data-nosnippet>391</a> <span class="prelude-val">None</span>,
<a href=#392 id=392 data-nosnippet>392</a> )
<a href=#393 id=393 data-nosnippet>393</a> .<span class="kw">await
<a href=#394 id=394 data-nosnippet>394</a></span>}
<a href=#395 id=395 data-nosnippet>395</a>
<a href=#396 id=396 data-nosnippet>396</a><span class="kw">async fn </span>export_stream_managed(
<a href=#397 id=397 data-nosnippet>397</a> unit: <span class="kw-2">&amp;</span>LogicalTableExport,
<a href=#398 id=398 data-nosnippet>398</a> stream: SendableRecordBatchStream,
<a href=#399 id=399 data-nosnippet>399</a> store: <span class="kw-2">&amp;</span>ObjectStore,
<a href=#400 id=400 data-nosnippet>400</a> limits: LogicalTableExportLimits,
<a href=#401 id=401 data-nosnippet>401</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#402 id=402 data-nosnippet>402</a> budget: Arc&lt;ExportWriteBudget&gt;,
<a href=#403 id=403 data-nosnippet>403</a> packed: <span class="prelude-ty">Option</span>&lt;PackedWriterRef&gt;,
<a href=#404 id=404 data-nosnippet>404</a>) -&gt; <span class="prelude-ty">Result</span>&lt;LogicalTableExportSummary&gt; {
<a href=#405 id=405 data-nosnippet>405</a> <span class="kw">let </span><span class="kw-2">mut </span>writers = TableWriters::new(budget.clone());
<a href=#406 id=406 data-nosnippet>406</a> writers.packed = packed;
<a href=#407 id=407 data-nosnippet>407</a> <span class="kw">let </span>result = write_tables(
<a href=#408 id=408 data-nosnippet>408</a> unit,
<a href=#409 id=409 data-nosnippet>409</a> stream,
<a href=#410 id=410 data-nosnippet>410</a> store,
<a href=#411 id=411 data-nosnippet>411</a> limits,
<a href=#412 id=412 data-nosnippet>412</a> cancellation,
<a href=#413 id=413 data-nosnippet>413</a> <span class="kw-2">&amp;mut </span>writers,
<a href=#414 id=414 data-nosnippet>414</a> <span class="kw-2">&amp;</span>budget,
<a href=#415 id=415 data-nosnippet>415</a> )
<a href=#416 id=416 data-nosnippet>416</a> .<span class="kw">await</span>;
<a href=#417 id=417 data-nosnippet>417</a> <span class="kw">let </span>(summary, result) = <span class="kw">match </span>result {
<a href=#418 id=418 data-nosnippet>418</a> <span class="prelude-val">Ok</span>(summary) =&gt; (summary, <span class="prelude-val">Ok</span>(())),
<a href=#419 id=419 data-nosnippet>419</a> <span class="prelude-val">Err</span>(error) =&gt; (LogicalTableExportSummary::default(), <span class="prelude-val">Err</span>(error)),
<a href=#420 id=420 data-nosnippet>420</a> };
<a href=#421 id=421 data-nosnippet>421</a> writers.drain(result, cancellation).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#422 id=422 data-nosnippet>422</a> <span class="prelude-val">Ok</span>(summary)
<a href=#423 id=423 data-nosnippet>423</a>}
<a href=#424 id=424 data-nosnippet>424</a>
<a href=#425 id=425 data-nosnippet>425</a><span class="kw">fn </span>check_cancelled(cancellation: <span class="kw-2">&amp;</span>CancellationToken) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#426 id=426 data-nosnippet>426</a> <span class="macro">ensure!</span>(
<a href=#427 id=427 data-nosnippet>427</a> !cancellation.is_cancelled(),
<a href=#428 id=428 data-nosnippet>428</a> error::LogicalTableExportCancelledSnafu
<a href=#429 id=429 data-nosnippet>429</a> );
<a href=#430 id=430 data-nosnippet>430</a> <span class="prelude-val">Ok</span>(())
<a href=#431 id=431 data-nosnippet>431</a>}
<a href=#432 id=432 data-nosnippet>432</a>
<a href=#433 id=433 data-nosnippet>433</a><span class="kw">async fn </span>write_tables(
<a href=#434 id=434 data-nosnippet>434</a> unit: <span class="kw-2">&amp;</span>LogicalTableExport,
<a href=#435 id=435 data-nosnippet>435</a> <span class="kw-2">mut </span>stream: SendableRecordBatchStream,
<a href=#436 id=436 data-nosnippet>436</a> store: <span class="kw-2">&amp;</span>ObjectStore,
<a href=#437 id=437 data-nosnippet>437</a> limits: LogicalTableExportLimits,
<a href=#438 id=438 data-nosnippet>438</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#439 id=439 data-nosnippet>439</a> writers: <span class="kw-2">&amp;mut </span>TableWriters,
<a href=#440 id=440 data-nosnippet>440</a> budget: <span class="kw-2">&amp;</span>ExportWriteBudget,
<a href=#441 id=441 data-nosnippet>441</a>) -&gt; <span class="prelude-ty">Result</span>&lt;LogicalTableExportSummary&gt; {
<a href=#442 id=442 data-nosnippet>442</a> <span class="kw">let </span>id_index =
<a href=#443 id=443 data-nosnippet>443</a> stream
<a href=#444 id=444 data-nosnippet>444</a> .schema()
<a href=#445 id=445 data-nosnippet>445</a> .column_index_by_name(TABLE_ID)
<a href=#446 id=446 data-nosnippet>446</a> .context(error::UnexpectedSnafu {
<a href=#447 id=447 data-nosnippet>447</a> violated: <span class="string">"physical query omitted __table_id"</span>,
<a href=#448 id=448 data-nosnippet>448</a> })<span class="question-mark">?</span>;
<a href=#449 id=449 data-nosnippet>449</a> <span class="kw">let </span><span class="kw-2">mut </span>summary = LogicalTableExportSummary::default();
<a href=#450 id=450 data-nosnippet>450</a> <span class="kw">let </span><span class="kw-2">mut </span>previous = <span class="prelude-val">None</span>;
<a href=#451 id=451 data-nosnippet>451</a> <span class="kw">let </span><span class="kw-2">mut </span>written = BTreeSet::new();
<a href=#452 id=452 data-nosnippet>452</a> <span class="kw">loop </span>{
<a href=#453 id=453 data-nosnippet>453</a> <span class="kw">let </span>batch = <span class="macro">tokio::select!</span> {
<a href=#454 id=454 data-nosnippet>454</a> biased;
<a href=#455 id=455 data-nosnippet>455</a> <span class="kw">_ </span>= cancellation.cancelled() =&gt; <span class="kw">return </span>error::LogicalTableExportCancelledSnafu.fail(),
<a href=#456 id=456 data-nosnippet>456</a> batch = stream.next() =&gt; batch,
<a href=#457 id=457 data-nosnippet>457</a> };
<a href=#458 id=458 data-nosnippet>458</a> <span class="kw">let </span><span class="prelude-val">Some</span>(batch) = batch <span class="kw">else </span>{
<a href=#459 id=459 data-nosnippet>459</a> <span class="kw">break</span>;
<a href=#460 id=460 data-nosnippet>460</a> };
<a href=#461 id=461 data-nosnippet>461</a> <span class="kw">let </span>batch = batch
<a href=#462 id=462 data-nosnippet>462</a> .context(error::BuildRecordBatchSnafu)<span class="question-mark">?
<a href=#463 id=463 data-nosnippet>463</a> </span>.into_df_record_batch();
<a href=#464 id=464 data-nosnippet>464</a> <span class="macro">ensure!</span>(
<a href=#465 id=465 data-nosnippet>465</a> batch.get_array_memory_size() &lt;= limits.input_batch_bytes,
<a href=#466 id=466 data-nosnippet>466</a> LogicalTableExportResourceSnafu {
<a href=#467 id=467 data-nosnippet>467</a> reason: <span class="string">"scan batch exceeds input byte budget"
<a href=#468 id=468 data-nosnippet>468</a> </span>}
<a href=#469 id=469 data-nosnippet>469</a> );
<a href=#470 id=470 data-nosnippet>470</a> <span class="kw">let </span>ids = batch
<a href=#471 id=471 data-nosnippet>471</a> .column(id_index)
<a href=#472 id=472 data-nosnippet>472</a> .as_any()
<a href=#473 id=473 data-nosnippet>473</a> .downcast_ref::&lt;UInt32Array&gt;()
<a href=#474 id=474 data-nosnippet>474</a> .context(error::UnexpectedSnafu {
<a href=#475 id=475 data-nosnippet>475</a> violated: <span class="string">"__table_id must be UInt32"</span>,
<a href=#476 id=476 data-nosnippet>476</a> })<span class="question-mark">?</span>;
<a href=#477 id=477 data-nosnippet>477</a> <span class="macro">ensure!</span>(
<a href=#478 id=478 data-nosnippet>478</a> ids.null_count() == <span class="number">0</span>,
<a href=#479 id=479 data-nosnippet>479</a> error::UnexpectedSnafu {
<a href=#480 id=480 data-nosnippet>480</a> violated: <span class="string">"null __table_id"
<a href=#481 id=481 data-nosnippet>481</a> </span>}
<a href=#482 id=482 data-nosnippet>482</a> );
<a href=#483 id=483 data-nosnippet>483</a> <span class="kw">let </span><span class="kw-2">mut </span>start = <span class="number">0</span>;
<a href=#484 id=484 data-nosnippet>484</a> <span class="kw">while </span>start &lt; batch.num_rows() {
<a href=#485 id=485 data-nosnippet>485</a> check_cancelled(cancellation)<span class="question-mark">?</span>;
<a href=#486 id=486 data-nosnippet>486</a> <span class="kw">let </span>id = ids.value(start);
<a href=#487 id=487 data-nosnippet>487</a> <span class="macro">ensure!</span>(
<a href=#488 id=488 data-nosnippet>488</a> previous.is_none_or(|last| last &lt;= id),
<a href=#489 id=489 data-nosnippet>489</a> error::UnexpectedSnafu {
<a href=#490 id=490 data-nosnippet>490</a> violated: <span class="string">"physical query is not ordered by __table_id"
<a href=#491 id=491 data-nosnippet>491</a> </span>}
<a href=#492 id=492 data-nosnippet>492</a> );
<a href=#493 id=493 data-nosnippet>493</a> previous = <span class="prelude-val">Some</span>(id);
<a href=#494 id=494 data-nosnippet>494</a> <span class="kw">let </span><span class="kw-2">mut </span>end = start + <span class="number">1</span>;
<a href=#495 id=495 data-nosnippet>495</a> <span class="kw">while </span>end &lt; batch.num_rows() &amp;&amp; ids.value(end) == id {
<a href=#496 id=496 data-nosnippet>496</a> end += <span class="number">1</span>;
<a href=#497 id=497 data-nosnippet>497</a> }
<a href=#498 id=498 data-nosnippet>498</a> <span class="kw">if </span>writers.table_id().is_some_and(|last| last != id) {
<a href=#499 id=499 data-nosnippet>499</a> writers.close_input();
<a href=#500 id=500 data-nosnippet>500</a> }
<a href=#501 id=501 data-nosnippet>501</a> check_cancelled(cancellation)<span class="question-mark">?</span>;
<a href=#502 id=502 data-nosnippet>502</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(file) = unit.logical_tables.get(<span class="kw-2">&amp;</span>id) {
<a href=#503 id=503 data-nosnippet>503</a> <span class="kw">if </span>writers.table_id().is_none() {
<a href=#504 id=504 data-nosnippet>504</a> writers.open(id, file, store, limits, cancellation).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#505 id=505 data-nosnippet>505</a> written.insert(id);
<a href=#506 id=506 data-nosnippet>506</a> summary.files += <span class="number">1</span>;
<a href=#507 id=507 data-nosnippet>507</a> }
<a href=#508 id=508 data-nosnippet>508</a> <span class="kw">let </span>projected = batch
<a href=#509 id=509 data-nosnippet>509</a> .project(<span class="kw-2">&amp;</span>file.projection)
<a href=#510 id=510 data-nosnippet>510</a> .context(error::ProjectSchemaSnafu)<span class="question-mark">?</span>;
<a href=#511 id=511 data-nosnippet>511</a> <span class="kw">let </span><span class="kw-2">mut </span>offset = start;
<a href=#512 id=512 data-nosnippet>512</a> <span class="kw">while </span>offset &lt; end {
<a href=#513 id=513 data-nosnippet>513</a> <span class="kw">let </span>(expanded, consumed) = expand_bounded_slice(
<a href=#514 id=514 data-nosnippet>514</a> projected.clone(),
<a href=#515 id=515 data-nosnippet>515</a> file.schema.clone(),
<a href=#516 id=516 data-nosnippet>516</a> offset,
<a href=#517 id=517 data-nosnippet>517</a> end,
<a href=#518 id=518 data-nosnippet>518</a> limits.conversion_bytes,
<a href=#519 id=519 data-nosnippet>519</a> budget,
<a href=#520 id=520 data-nosnippet>520</a> cancellation,
<a href=#521 id=521 data-nosnippet>521</a> )
<a href=#522 id=522 data-nosnippet>522</a> .<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#523 id=523 data-nosnippet>523</a> check_cancelled(cancellation)<span class="question-mark">?</span>;
<a href=#524 id=524 data-nosnippet>524</a> writers.send(expanded, cancellation).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#525 id=525 data-nosnippet>525</a> check_cancelled(cancellation)<span class="question-mark">?</span>;
<a href=#526 id=526 data-nosnippet>526</a> offset += consumed;
<a href=#527 id=527 data-nosnippet>527</a> summary.rows += consumed;
<a href=#528 id=528 data-nosnippet>528</a> }
<a href=#529 id=529 data-nosnippet>529</a> } <span class="kw">else </span>{
<a href=#530 id=530 data-nosnippet>530</a> summary.skipped_rows += end - start;
<a href=#531 id=531 data-nosnippet>531</a> }
<a href=#532 id=532 data-nosnippet>532</a> start = end;
<a href=#533 id=533 data-nosnippet>533</a> }
<a href=#534 id=534 data-nosnippet>534</a> }
<a href=#535 id=535 data-nosnippet>535</a> writers.close_input();
<a href=#536 id=536 data-nosnippet>536</a> <span class="kw">for </span>(<span class="kw-2">&amp;</span>id, file) <span class="kw">in </span><span class="kw-2">&amp;</span>unit.logical_tables {
<a href=#537 id=537 data-nosnippet>537</a> <span class="kw">if </span>!written.contains(<span class="kw-2">&amp;</span>id) {
<a href=#538 id=538 data-nosnippet>538</a> check_cancelled(cancellation)<span class="question-mark">?</span>;
<a href=#539 id=539 data-nosnippet>539</a> writers.open(id, file, store, limits, cancellation).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#540 id=540 data-nosnippet>540</a> writers.close_input();
<a href=#541 id=541 data-nosnippet>541</a> summary.files += <span class="number">1</span>;
<a href=#542 id=542 data-nosnippet>542</a> }
<a href=#543 id=543 data-nosnippet>543</a> }
<a href=#544 id=544 data-nosnippet>544</a> check_cancelled(cancellation)<span class="question-mark">?</span>;
<a href=#545 id=545 data-nosnippet>545</a> <span class="prelude-val">Ok</span>(summary)
<a href=#546 id=546 data-nosnippet>546</a>}
<a href=#547 id=547 data-nosnippet>547</a>
<a href=#548 id=548 data-nosnippet>548</a><span class="kw">struct </span>ActiveWriter {
<a href=#549 id=549 data-nosnippet>549</a> path: String,
<a href=#550 id=550 data-nosnippet>550</a> writer: ParquetFileWriter,
<a href=#551 id=551 data-nosnippet>551</a>}
<a href=#552 id=552 data-nosnippet>552</a>
<a href=#553 id=553 data-nosnippet>553</a><span class="kw">impl </span>ActiveWriter {
<a href=#554 id=554 data-nosnippet>554</a> <span class="kw">fn </span>open_packed(
<a href=#555 id=555 data-nosnippet>555</a> table: <span class="kw-2">&amp;</span>LogicalTableProjection,
<a href=#556 id=556 data-nosnippet>556</a> id: u32,
<a href=#557 id=557 data-nosnippet>557</a> store: <span class="kw-2">&amp;</span>ObjectStore,
<a href=#558 id=558 data-nosnippet>558</a> limits: LogicalTableExportLimits,
<a href=#559 id=559 data-nosnippet>559</a> shared: PackedWriterRef,
<a href=#560 id=560 data-nosnippet>560</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="self">Self</span>&gt; {
<a href=#561 id=561 data-nosnippet>561</a> <span class="kw">let </span>path = <span class="macro">format!</span>(<span class="string">"table-{id}.parquet"</span>);
<a href=#562 id=562 data-nosnippet>562</a> <span class="kw">let </span>packed = PackedTableWriter::new(shared, table.name.clone(), id, <span class="bool-val">false</span>);
<a href=#563 id=563 data-nosnippet>563</a> <span class="kw">let </span>writer = ParquetFileWriter::open_packed(
<a href=#564 id=564 data-nosnippet>564</a> table.schema.clone(),
<a href=#565 id=565 data-nosnippet>565</a> store.clone(),
<a href=#566 id=566 data-nosnippet>566</a> <span class="kw-2">&amp;</span>path,
<a href=#567 id=567 data-nosnippet>567</a> <span class="prelude-val">Some</span>(limits.writer),
<a href=#568 id=568 data-nosnippet>568</a> packed,
<a href=#569 id=569 data-nosnippet>569</a> )
<a href=#570 id=570 data-nosnippet>570</a> .map_err(|e| map_writer_error(e, <span class="kw-2">&amp;</span>path))<span class="question-mark">?</span>;
<a href=#571 id=571 data-nosnippet>571</a> <span class="prelude-val">Ok</span>(<span class="self">Self </span>{ path, writer })
<a href=#572 id=572 data-nosnippet>572</a> }
<a href=#573 id=573 data-nosnippet>573</a>
<a href=#574 id=574 data-nosnippet>574</a> <span class="kw">async fn </span>open(
<a href=#575 id=575 data-nosnippet>575</a> table: <span class="kw-2">&amp;</span>LogicalTableProjection,
<a href=#576 id=576 data-nosnippet>576</a> store: <span class="kw-2">&amp;</span>ObjectStore,
<a href=#577 id=577 data-nosnippet>577</a> limits: LogicalTableExportLimits,
<a href=#578 id=578 data-nosnippet>578</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="self">Self</span>&gt; {
<a href=#579 id=579 data-nosnippet>579</a> <span class="kw">let </span>path = table.output.path.clone();
<a href=#580 id=580 data-nosnippet>580</a> <span class="kw">let </span>conditional = store.info().capability().write_with_if_not_exists;
<a href=#581 id=581 data-nosnippet>581</a> <span class="kw">if </span>!conditional {
<a href=#582 id=582 data-nosnippet>582</a> <span class="macro">ensure!</span>(
<a href=#583 id=583 data-nosnippet>583</a> !store
<a href=#584 id=584 data-nosnippet>584</a> .exists(<span class="kw-2">&amp;</span>path)
<a href=#585 id=585 data-nosnippet>585</a> .<span class="kw">await
<a href=#586 id=586 data-nosnippet>586</a> </span>.context(error::ReadObjectSnafu { path: <span class="kw-2">&amp;</span>path })<span class="question-mark">?</span>,
<a href=#587 id=587 data-nosnippet>587</a> InvalidLogicalTableExportSnafu {
<a href=#588 id=588 data-nosnippet>588</a> reason: <span class="macro">format!</span>(<span class="string">"output already exists: {path}"</span>)
<a href=#589 id=589 data-nosnippet>589</a> }
<a href=#590 id=590 data-nosnippet>590</a> );
<a href=#591 id=591 data-nosnippet>591</a> }
<a href=#592 id=592 data-nosnippet>592</a> <span class="kw">let </span>writer = ParquetFileWriter::open_with_creation(
<a href=#593 id=593 data-nosnippet>593</a> table.schema.clone(),
<a href=#594 id=594 data-nosnippet>594</a> store.clone(),
<a href=#595 id=595 data-nosnippet>595</a> <span class="kw-2">&amp;</span>path,
<a href=#596 id=596 data-nosnippet>596</a> <span class="number">1</span>,
<a href=#597 id=597 data-nosnippet>597</a> <span class="prelude-val">Some</span>(limits.writer),
<a href=#598 id=598 data-nosnippet>598</a> <span class="kw">if </span>conditional {
<a href=#599 id=599 data-nosnippet>599</a> ParquetCreationPolicy::IfNotExists
<a href=#600 id=600 data-nosnippet>600</a> } <span class="kw">else </span>{
<a href=#601 id=601 data-nosnippet>601</a> ParquetCreationPolicy::Overwrite
<a href=#602 id=602 data-nosnippet>602</a> },
<a href=#603 id=603 data-nosnippet>603</a> )
<a href=#604 id=604 data-nosnippet>604</a> .<span class="kw">await
<a href=#605 id=605 data-nosnippet>605</a> </span>.map_err(|error| map_writer_error(error, <span class="kw-2">&amp;</span>path))<span class="question-mark">?</span>;
<a href=#606 id=606 data-nosnippet>606</a> <span class="prelude-val">Ok</span>(<span class="self">Self </span>{ path, writer })
<a href=#607 id=607 data-nosnippet>607</a> }
<a href=#608 id=608 data-nosnippet>608</a>}
<a href=#609 id=609 data-nosnippet>609</a>
<a href=#610 id=610 data-nosnippet>610</a><span class="kw">async fn </span>expand_bounded_slice(
<a href=#611 id=611 data-nosnippet>611</a> batch: RecordBatch,
<a href=#612 id=612 data-nosnippet>612</a> schema: SchemaRef,
<a href=#613 id=613 data-nosnippet>613</a> start: usize,
<a href=#614 id=614 data-nosnippet>614</a> end: usize,
<a href=#615 id=615 data-nosnippet>615</a> requested: usize,
<a href=#616 id=616 data-nosnippet>616</a> budget: <span class="kw-2">&amp;</span>ExportWriteBudget,
<a href=#617 id=617 data-nosnippet>617</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#618 id=618 data-nosnippet>618</a>) -&gt; <span class="prelude-ty">Result</span>&lt;(Payload, usize)&gt; {
<a href=#619 id=619 data-nosnippet>619</a> <span class="kw">let </span>(conversion, retained) = ExportWriteBudget::conversion_budget(
<a href=#620 id=620 data-nosnippet>620</a> batch.get_array_memory_size(),
<a href=#621 id=621 data-nosnippet>621</a> batch.num_columns(),
<a href=#622 id=622 data-nosnippet>622</a> requested,
<a href=#623 id=623 data-nosnippet>623</a> )<span class="question-mark">?</span>;
<a href=#624 id=624 data-nosnippet>624</a> <span class="comment">// Parquet accepts compatible string dictionaries without expanding their values.
<a href=#625 id=625 data-nosnippet>625</a> </span><span class="kw">let </span>fields = schema
<a href=#626 id=626 data-nosnippet>626</a> .fields()
<a href=#627 id=627 data-nosnippet>627</a> .iter()
<a href=#628 id=628 data-nosnippet>628</a> .zip(batch.columns())
<a href=#629 id=629 data-nosnippet>629</a> .map(|(field, array)| {
<a href=#630 id=630 data-nosnippet>630</a> <span class="kw">if let </span>DataType::Dictionary(<span class="kw">_</span>, value) = array.data_type()
<a href=#631 id=631 data-nosnippet>631</a> &amp;&amp; value.as_ref() == field.data_type()
<a href=#632 id=632 data-nosnippet>632</a> &amp;&amp; <span class="macro">matches!</span>(value.as_ref(), DataType::Utf8 | DataType::LargeUtf8)
<a href=#633 id=633 data-nosnippet>633</a> {
<a href=#634 id=634 data-nosnippet>634</a> Arc::new(
<a href=#635 id=635 data-nosnippet>635</a> field
<a href=#636 id=636 data-nosnippet>636</a> .as_ref()
<a href=#637 id=637 data-nosnippet>637</a> .clone()
<a href=#638 id=638 data-nosnippet>638</a> .with_data_type(array.data_type().clone()),
<a href=#639 id=639 data-nosnippet>639</a> )
<a href=#640 id=640 data-nosnippet>640</a> } <span class="kw">else </span>{
<a href=#641 id=641 data-nosnippet>641</a> field.clone()
<a href=#642 id=642 data-nosnippet>642</a> }
<a href=#643 id=643 data-nosnippet>643</a> })
<a href=#644 id=644 data-nosnippet>644</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;();
<a href=#645 id=645 data-nosnippet>645</a> <span class="kw">let </span>schema = Arc::new(arrow::datatypes::Schema::new_with_metadata(
<a href=#646 id=646 data-nosnippet>646</a> fields,
<a href=#647 id=647 data-nosnippet>647</a> schema.metadata().clone(),
<a href=#648 id=648 data-nosnippet>648</a> ));
<a href=#649 id=649 data-nosnippet>649</a> <span class="comment">// Sharing compatible arrays needs no CPU conversion. Keep the size check on
<a href=#650 id=650 data-nosnippet>650</a> // this task only when its dictionary inspection is bounded.
<a href=#651 id=651 data-nosnippet>651</a> </span><span class="kw">if </span>start &lt; end
<a href=#652 id=652 data-nosnippet>652</a> &amp;&amp; batch
<a href=#653 id=653 data-nosnippet>653</a> .columns()
<a href=#654 id=654 data-nosnippet>654</a> .iter()
<a href=#655 id=655 data-nosnippet>655</a> .zip(schema.fields())
<a href=#656 id=656 data-nosnippet>656</a> .all(|(array, field)| array.data_type() == field.data_type())
<a href=#657 id=657 data-nosnippet>657</a> &amp;&amp; <span class="kw">let </span><span class="prelude-val">Some</span>(estimated) = estimate_batch_bytes_upper_bound(<span class="kw-2">&amp;</span>batch, end - start, <span class="kw-2">&amp;</span>[])<span class="question-mark">?
<a href=#658 id=658 data-nosnippet>658</a> </span>&amp;&amp; estimated &lt;= conversion
<a href=#659 id=659 data-nosnippet>659</a> {
<a href=#660 id=660 data-nosnippet>660</a> <span class="kw">let </span>reservation = retained.saturating_add(estimated.saturating_mul(<span class="number">4</span>));
<a href=#661 id=661 data-nosnippet>661</a> <span class="kw">let </span>permit = budget.reserve(reservation, cancellation).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#662 id=662 data-nosnippet>662</a> <span class="kw">return </span>convert_slice_payload(batch, schema, start, end - start, reservation, permit);
<a href=#663 id=663 data-nosnippet>663</a> }
<a href=#664 id=664 data-nosnippet>664</a> <span class="kw">let </span>input = batch.clone();
<a href=#665 id=665 data-nosnippet>665</a> <span class="kw">let </span>(len, estimated) = common_runtime::spawn_blocking_global(<span class="kw">move </span>|| {
<a href=#666 id=666 data-nosnippet>666</a> rows_within_budget(<span class="kw-2">&amp;</span>input, start, end, conversion, <span class="kw-2">&amp;</span>[])
<a href=#667 id=667 data-nosnippet>667</a> })
<a href=#668 id=668 data-nosnippet>668</a> .<span class="kw">await
<a href=#669 id=669 data-nosnippet>669</a> </span>.context(error::JoinTaskSnafu)<span class="question-mark">??</span>;
<a href=#670 id=670 data-nosnippet>670</a> <span class="kw">let </span>reservation = retained.saturating_add(estimated.saturating_mul(<span class="number">4</span>));
<a href=#671 id=671 data-nosnippet>671</a> <span class="kw">let </span>permit = budget.reserve(reservation, cancellation).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#672 id=672 data-nosnippet>672</a> common_runtime::spawn_blocking_global(<span class="kw">move </span>|| {
<a href=#673 id=673 data-nosnippet>673</a> convert_slice_payload(batch, schema, start, len, reservation, permit)
<a href=#674 id=674 data-nosnippet>674</a> })
<a href=#675 id=675 data-nosnippet>675</a> .<span class="kw">await
<a href=#676 id=676 data-nosnippet>676</a> </span>.context(error::JoinTaskSnafu)<span class="question-mark">?
<a href=#677 id=677 data-nosnippet>677</a></span>}
<a href=#678 id=678 data-nosnippet>678</a>
<a href=#679 id=679 data-nosnippet>679</a><span class="kw">fn </span>convert_slice_payload(
<a href=#680 id=680 data-nosnippet>680</a> batch: RecordBatch,
<a href=#681 id=681 data-nosnippet>681</a> schema: SchemaRef,
<a href=#682 id=682 data-nosnippet>682</a> start: usize,
<a href=#683 id=683 data-nosnippet>683</a> len: usize,
<a href=#684 id=684 data-nosnippet>684</a> reservation: usize,
<a href=#685 id=685 data-nosnippet>685</a> permit: OwnedSemaphorePermit,
<a href=#686 id=686 data-nosnippet>686</a>) -&gt; <span class="prelude-ty">Result</span>&lt;(Payload, usize)&gt; {
<a href=#687 id=687 data-nosnippet>687</a> <span class="kw">let </span>expanded = expand_export_batch(<span class="kw-2">&amp;</span>batch.slice(start, len), schema)<span class="question-mark">?</span>;
<a href=#688 id=688 data-nosnippet>688</a> <span class="macro">ensure!</span>(
<a href=#689 id=689 data-nosnippet>689</a> expanded.get_array_memory_size() &lt;= reservation,
<a href=#690 id=690 data-nosnippet>690</a> LogicalTableExportResourceSnafu {
<a href=#691 id=691 data-nosnippet>691</a> reason: <span class="string">"converted backing buffers exceed reservation"
<a href=#692 id=692 data-nosnippet>692</a> </span>}
<a href=#693 id=693 data-nosnippet>693</a> );
<a href=#694 id=694 data-nosnippet>694</a> <span class="prelude-val">Ok</span>((
<a href=#695 id=695 data-nosnippet>695</a> Payload {
<a href=#696 id=696 data-nosnippet>696</a> batch: expanded,
<a href=#697 id=697 data-nosnippet>697</a> permit,
<a href=#698 id=698 data-nosnippet>698</a> },
<a href=#699 id=699 data-nosnippet>699</a> len,
<a href=#700 id=700 data-nosnippet>700</a> ))
<a href=#701 id=701 data-nosnippet>701</a>}
<a href=#702 id=702 data-nosnippet>702</a>
<a href=#703 id=703 data-nosnippet>703</a><span class="doccomment">/// Limits nested casts to selected children: Arrow's List cast otherwise expands
<a href=#704 id=704 data-nosnippet>704</a>/// the entire values array, even when the parent has been sliced.
<a href=#705 id=705 data-nosnippet>705</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>expand_export_batch(batch: <span class="kw-2">&amp;</span>RecordBatch, schema: SchemaRef) -&gt; <span class="prelude-ty">Result</span>&lt;RecordBatch&gt; {
<a href=#706 id=706 data-nosnippet>706</a> <span class="kw">let </span>arrays = batch
<a href=#707 id=707 data-nosnippet>707</a> .columns()
<a href=#708 id=708 data-nosnippet>708</a> .iter()
<a href=#709 id=709 data-nosnippet>709</a> .zip(schema.fields())
<a href=#710 id=710 data-nosnippet>710</a> .map(|(array, field)| expand_export_array(array, field.data_type()))
<a href=#711 id=711 data-nosnippet>711</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#712 id=712 data-nosnippet>712</a> RecordBatch::try_new(schema, arrays).context(error::ComputeArrowSnafu)
<a href=#713 id=713 data-nosnippet>713</a>}
<a href=#714 id=714 data-nosnippet>714</a>
<a href=#715 id=715 data-nosnippet>715</a><span class="kw">fn </span>expand_export_array(array: <span class="kw-2">&amp;</span>ArrayRef, target: <span class="kw-2">&amp;</span>DataType) -&gt; <span class="prelude-ty">Result</span>&lt;ArrayRef&gt; {
<a href=#716 id=716 data-nosnippet>716</a> <span class="kw">if </span>array.data_type() == target {
<a href=#717 id=717 data-nosnippet>717</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(array.clone());
<a href=#718 id=718 data-nosnippet>718</a> }
<a href=#719 id=719 data-nosnippet>719</a> <span class="kw">match </span>(array.data_type(), target) {
<a href=#720 id=720 data-nosnippet>720</a> (DataType::Dictionary(<span class="kw">_</span>, <span class="kw">_</span>), <span class="kw">_</span>) =&gt; {
<a href=#721 id=721 data-nosnippet>721</a> <span class="macro">downcast_dictionary_array!</span> {
<a href=#722 id=722 data-nosnippet>722</a> array =&gt; {
<a href=#723 id=723 data-nosnippet>723</a> <span class="comment">// Select dictionary entries before recursively converting nested values.
<a href=#724 id=724 data-nosnippet>724</a> </span><span class="kw">let </span>selected = take(array.values().as_ref(), array.keys(), <span class="prelude-val">None</span>)
<a href=#725 id=725 data-nosnippet>725</a> .context(error::ComputeArrowSnafu)<span class="question-mark">?</span>;
<a href=#726 id=726 data-nosnippet>726</a> expand_export_array(<span class="kw-2">&amp;</span>selected, target)
<a href=#727 id=727 data-nosnippet>727</a> },
<a href=#728 id=728 data-nosnippet>728</a> <span class="kw">_ </span>=&gt; error::UnexpectedSnafu { violated: <span class="string">"invalid dictionary array" </span>}.fail(),
<a href=#729 id=729 data-nosnippet>729</a> }
<a href=#730 id=730 data-nosnippet>730</a> }
<a href=#731 id=731 data-nosnippet>731</a> (DataType::List(<span class="kw">_</span>), DataType::List(field)) =&gt; {
<a href=#732 id=732 data-nosnippet>732</a> <span class="kw">let </span>list = array.as_list::&lt;i32&gt;();
<a href=#733 id=733 data-nosnippet>733</a> <span class="kw">let </span>offsets = list.value_offsets();
<a href=#734 id=734 data-nosnippet>734</a> <span class="kw">let </span>start = offsets[<span class="number">0</span>];
<a href=#735 id=735 data-nosnippet>735</a> <span class="kw">let </span>end = offsets[list.len()];
<a href=#736 id=736 data-nosnippet>736</a> <span class="kw">let </span>values = list.values().slice(start <span class="kw">as </span>usize, (end - start) <span class="kw">as </span>usize);
<a href=#737 id=737 data-nosnippet>737</a> <span class="kw">let </span>values = expand_export_array(<span class="kw-2">&amp;</span>values, field.data_type())<span class="question-mark">?</span>;
<a href=#738 id=738 data-nosnippet>738</a> <span class="kw">let </span>offsets = OffsetBuffer::new(
<a href=#739 id=739 data-nosnippet>739</a> offsets
<a href=#740 id=740 data-nosnippet>740</a> .iter()
<a href=#741 id=741 data-nosnippet>741</a> .map(|offset| offset - start)
<a href=#742 id=742 data-nosnippet>742</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;()
<a href=#743 id=743 data-nosnippet>743</a> .into(),
<a href=#744 id=744 data-nosnippet>744</a> );
<a href=#745 id=745 data-nosnippet>745</a> <span class="prelude-val">Ok</span>(Arc::new(
<a href=#746 id=746 data-nosnippet>746</a> ListArray::try_new(field.clone(), offsets, values, list.nulls().cloned())
<a href=#747 id=747 data-nosnippet>747</a> .context(error::ComputeArrowSnafu)<span class="question-mark">?</span>,
<a href=#748 id=748 data-nosnippet>748</a> ))
<a href=#749 id=749 data-nosnippet>749</a> }
<a href=#750 id=750 data-nosnippet>750</a> (DataType::Struct(<span class="kw">_</span>), DataType::Struct(fields)) =&gt; {
<a href=#751 id=751 data-nosnippet>751</a> <span class="kw">let </span>array = array.as_struct();
<a href=#752 id=752 data-nosnippet>752</a> <span class="kw">let </span>columns = array
<a href=#753 id=753 data-nosnippet>753</a> .columns()
<a href=#754 id=754 data-nosnippet>754</a> .iter()
<a href=#755 id=755 data-nosnippet>755</a> .zip(fields)
<a href=#756 id=756 data-nosnippet>756</a> .map(|(array, field)| expand_export_array(array, field.data_type()))
<a href=#757 id=757 data-nosnippet>757</a> .collect::&lt;<span class="prelude-ty">Result</span>&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;&gt;()<span class="question-mark">?</span>;
<a href=#758 id=758 data-nosnippet>758</a> <span class="prelude-val">Ok</span>(Arc::new(
<a href=#759 id=759 data-nosnippet>759</a> StructArray::try_new(fields.clone(), columns, array.nulls().cloned())
<a href=#760 id=760 data-nosnippet>760</a> .context(error::ComputeArrowSnafu)<span class="question-mark">?</span>,
<a href=#761 id=761 data-nosnippet>761</a> ))
<a href=#762 id=762 data-nosnippet>762</a> }
<a href=#763 id=763 data-nosnippet>763</a> <span class="kw">_ </span>=&gt; cast(array, target).context(error::ComputeArrowSnafu),
<a href=#764 id=764 data-nosnippet>764</a> }
<a href=#765 id=765 data-nosnippet>765</a>}
<a href=#766 id=766 data-nosnippet>766</a>
<a href=#767 id=767 data-nosnippet>767</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>map_writer_error(
<a href=#768 id=768 data-nosnippet>768</a> source: common_datasource::error::Error,
<a href=#769 id=769 data-nosnippet>769</a> path: <span class="kw-2">&amp;</span>str,
<a href=#770 id=770 data-nosnippet>770</a>) -&gt; error::Error {
<a href=#771 id=771 data-nosnippet>771</a> <span class="kw">match </span>source {
<a href=#772 id=772 data-nosnippet>772</a> common_datasource::error::Error::ParquetWriteCancelled {} =&gt; {
<a href=#773 id=773 data-nosnippet>773</a> error::LogicalTableExportCancelledSnafu.build()
<a href=#774 id=774 data-nosnippet>774</a> }
<a href=#775 id=775 data-nosnippet>775</a> common_datasource::error::Error::InvalidParquetWriterLimits {} =&gt; {
<a href=#776 id=776 data-nosnippet>776</a> InvalidLogicalTableExportSnafu {
<a href=#777 id=777 data-nosnippet>777</a> reason: <span class="string">"Parquet writer limits must be positive"</span>,
<a href=#778 id=778 data-nosnippet>778</a> }
<a href=#779 id=779 data-nosnippet>779</a> .build()
<a href=#780 id=780 data-nosnippet>780</a> }
<a href=#781 id=781 data-nosnippet>781</a> common_datasource::error::Error::ParquetWriterResource { reason } =&gt; {
<a href=#782 id=782 data-nosnippet>782</a> LogicalTableExportResourceSnafu { reason }.build()
<a href=#783 id=783 data-nosnippet>783</a> }
<a href=#784 id=784 data-nosnippet>784</a> common_datasource::error::Error::WriteObject { error, .. }
<a href=#785 id=785 data-nosnippet>785</a> <span class="kw">if </span>error.kind() == object_store::ErrorKind::ConditionNotMatch =&gt;
<a href=#786 id=786 data-nosnippet>786</a> {
<a href=#787 id=787 data-nosnippet>787</a> <span class="kw">let </span>reason = <span class="macro">format!</span>(<span class="string">"output already exists: {path}"</span>);
<a href=#788 id=788 data-nosnippet>788</a> InvalidLogicalTableExportSnafu { reason }.build()
<a href=#789 id=789 data-nosnippet>789</a> }
<a href=#790 id=790 data-nosnippet>790</a> source =&gt; error::WriteStreamToFileSnafu { path }.into_error(source),
<a href=#791 id=791 data-nosnippet>791</a> }
<a href=#792 id=792 data-nosnippet>792</a>}
<a href=#793 id=793 data-nosnippet>793</a>
<a href=#794 id=794 data-nosnippet>794</a><span class="comment">// Count only selected logical values before dictionary expansion, including nested
<a href=#795 id=795 data-nosnippet>795</a>// histogram lists/structs. Offset and validity overhead is charged per value.
<a href=#796 id=796 data-nosnippet>796</a></span><span class="kw">fn </span>estimate_value_size(array: <span class="kw-2">&amp;</span><span class="kw">dyn </span>Array, row: usize) -&gt; <span class="prelude-ty">Result</span>&lt;usize&gt; {
<a href=#797 id=797 data-nosnippet>797</a> <span class="kw">if </span>array.is_null(row) &amp;&amp; !<span class="macro">matches!</span>(array.data_type(), DataType::Struct(<span class="kw">_</span>) | DataType::List(<span class="kw">_</span>)) {
<a href=#798 id=798 data-nosnippet>798</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(<span class="number">32</span>);
<a href=#799 id=799 data-nosnippet>799</a> }
<a href=#800 id=800 data-nosnippet>800</a> <span class="kw">let </span>bytes = <span class="kw">match </span>array.data_type() {
<a href=#801 id=801 data-nosnippet>801</a> DataType::Boolean =&gt; <span class="number">1</span>,
<a href=#802 id=802 data-nosnippet>802</a> DataType::Null =&gt; <span class="number">0</span>,
<a href=#803 id=803 data-nosnippet>803</a> DataType::Utf8 =&gt; array.as_string::&lt;i32&gt;().value(row).len(),
<a href=#804 id=804 data-nosnippet>804</a> DataType::LargeUtf8 =&gt; array.as_string::&lt;i64&gt;().value(row).len(),
<a href=#805 id=805 data-nosnippet>805</a> DataType::Utf8View =&gt; array.as_string_view().value(row).len(),
<a href=#806 id=806 data-nosnippet>806</a> DataType::Binary =&gt; array.as_binary::&lt;i32&gt;().value(row).len(),
<a href=#807 id=807 data-nosnippet>807</a> DataType::LargeBinary =&gt; array.as_binary::&lt;i64&gt;().value(row).len(),
<a href=#808 id=808 data-nosnippet>808</a> DataType::BinaryView =&gt; array.as_binary_view().value(row).len(),
<a href=#809 id=809 data-nosnippet>809</a> DataType::Struct(<span class="kw">_</span>) =&gt; {
<a href=#810 id=810 data-nosnippet>810</a> array
<a href=#811 id=811 data-nosnippet>811</a> .as_struct()
<a href=#812 id=812 data-nosnippet>812</a> .columns()
<a href=#813 id=813 data-nosnippet>813</a> .iter()
<a href=#814 id=814 data-nosnippet>814</a> .try_fold(<span class="number">0usize</span>, |sum, child| {
<a href=#815 id=815 data-nosnippet>815</a> <span class="prelude-val">Ok</span>::&lt;<span class="kw">_</span>, error::Error&gt;(
<a href=#816 id=816 data-nosnippet>816</a> sum.saturating_add(estimate_value_size(child.as_ref(), row)<span class="question-mark">?</span>),
<a href=#817 id=817 data-nosnippet>817</a> )
<a href=#818 id=818 data-nosnippet>818</a> })<span class="question-mark">?
<a href=#819 id=819 data-nosnippet>819</a> </span>}
<a href=#820 id=820 data-nosnippet>820</a> DataType::List(<span class="kw">_</span>) =&gt; {
<a href=#821 id=821 data-nosnippet>821</a> <span class="kw">let </span>list = array.as_list::&lt;i32&gt;();
<a href=#822 id=822 data-nosnippet>822</a> <span class="kw">let </span>offsets = list.value_offsets();
<a href=#823 id=823 data-nosnippet>823</a> (offsets[row] <span class="kw">as </span>usize..offsets[row + <span class="number">1</span>] <span class="kw">as </span>usize).try_fold(<span class="number">0usize</span>, |sum, index| {
<a href=#824 id=824 data-nosnippet>824</a> <span class="prelude-val">Ok</span>::&lt;<span class="kw">_</span>, error::Error&gt;(
<a href=#825 id=825 data-nosnippet>825</a> sum.saturating_add(estimate_value_size(list.values().as_ref(), index)<span class="question-mark">?</span>),
<a href=#826 id=826 data-nosnippet>826</a> )
<a href=#827 id=827 data-nosnippet>827</a> })<span class="question-mark">?
<a href=#828 id=828 data-nosnippet>828</a> </span>}
<a href=#829 id=829 data-nosnippet>829</a> DataType::Dictionary(<span class="kw">_</span>, <span class="kw">_</span>) =&gt; {
<a href=#830 id=830 data-nosnippet>830</a> <span class="macro">downcast_dictionary_array!</span> {
<a href=#831 id=831 data-nosnippet>831</a> array =&gt; {
<a href=#832 id=832 data-nosnippet>832</a> <span class="kw">match </span>array.key(row) {
<a href=#833 id=833 data-nosnippet>833</a> <span class="prelude-val">Some</span>(index) =&gt; estimate_value_size(array.values().as_ref(), index)<span class="question-mark">?</span>,
<a href=#834 id=834 data-nosnippet>834</a> <span class="prelude-val">None </span>=&gt; <span class="number">0</span>,
<a href=#835 id=835 data-nosnippet>835</a> }
<a href=#836 id=836 data-nosnippet>836</a> },
<a href=#837 id=837 data-nosnippet>837</a> <span class="kw">_ </span>=&gt; <span class="kw">return </span>error::UnexpectedSnafu { violated: <span class="string">"invalid dictionary array" </span>}.fail(),
<a href=#838 id=838 data-nosnippet>838</a> }
<a href=#839 id=839 data-nosnippet>839</a> }
<a href=#840 id=840 data-nosnippet>840</a> other =&gt; other
<a href=#841 id=841 data-nosnippet>841</a> .primitive_width()
<a href=#842 id=842 data-nosnippet>842</a> .with_context(|| InvalidLogicalTableExportSnafu {
<a href=#843 id=843 data-nosnippet>843</a> reason: <span class="macro">format!</span>(<span class="string">"unsupported Metric Parquet type: {other}"</span>),
<a href=#844 id=844 data-nosnippet>844</a> })<span class="question-mark">?</span>,
<a href=#845 id=845 data-nosnippet>845</a> };
<a href=#846 id=846 data-nosnippet>846</a> <span class="prelude-val">Ok</span>(bytes.saturating_add(<span class="number">16</span>))
<a href=#847 id=847 data-nosnippet>847</a>}
<a href=#848 id=848 data-nosnippet>848</a>
<a href=#849 id=849 data-nosnippet>849</a><span class="doccomment">/// Estimates a conservative byte upper bound with bounded work on an async task.
<a href=#850 id=850 data-nosnippet>850</a>/// `None` means this path cannot estimate the batch, not that it exceeds a budget.
<a href=#851 id=851 data-nosnippet>851</a></span><span class="kw">fn </span>estimate_batch_bytes_upper_bound(
<a href=#852 id=852 data-nosnippet>852</a> batch: <span class="kw-2">&amp;</span>RecordBatch,
<a href=#853 id=853 data-nosnippet>853</a> rows: usize,
<a href=#854 id=854 data-nosnippet>854</a> json_columns: <span class="kw-2">&amp;</span>[usize],
<a href=#855 id=855 data-nosnippet>855</a>) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="prelude-ty">Option</span>&lt;usize&gt;&gt; {
<a href=#856 id=856 data-nosnippet>856</a> <span class="kw">let </span><span class="kw-2">mut </span>dictionary_values = <span class="number">1024usize</span>;
<a href=#857 id=857 data-nosnippet>857</a> <span class="kw">let </span><span class="kw-2">mut </span>max_row_bytes = <span class="number">0usize</span>;
<a href=#858 id=858 data-nosnippet>858</a> <span class="kw">for </span>(index, array) <span class="kw">in </span>batch.columns().iter().enumerate() {
<a href=#859 id=859 data-nosnippet>859</a> <span class="kw">let </span>size = <span class="kw">if let </span><span class="prelude-val">Some</span>(width) = array.data_type().primitive_width() {
<a href=#860 id=860 data-nosnippet>860</a> (width + <span class="number">16</span>).max(<span class="number">32</span>)
<a href=#861 id=861 data-nosnippet>861</a> } <span class="kw">else if let </span><span class="prelude-val">Some</span>(dictionary) = array.as_any_dictionary_opt()
<a href=#862 id=862 data-nosnippet>862</a> &amp;&amp; dictionary.values().len() &lt; rows
<a href=#863 id=863 data-nosnippet>863</a> &amp;&amp; dictionary.values().len() &lt;= dictionary_values
<a href=#864 id=864 data-nosnippet>864</a> &amp;&amp; <span class="macro">matches!</span>(
<a href=#865 id=865 data-nosnippet>865</a> dictionary.values().data_type(),
<a href=#866 id=866 data-nosnippet>866</a> DataType::Utf8 | DataType::LargeUtf8 | DataType::Utf8View
<a href=#867 id=867 data-nosnippet>867</a> )
<a href=#868 id=868 data-nosnippet>868</a> {
<a href=#869 id=869 data-nosnippet>869</a> <span class="kw">let </span>values = dictionary.values();
<a href=#870 id=870 data-nosnippet>870</a> dictionary_values -= values.len();
<a href=#871 id=871 data-nosnippet>871</a> <span class="kw">let </span>max = (<span class="number">0</span>..values.len()).try_fold(<span class="number">0usize</span>, |max, row| {
<a href=#872 id=872 data-nosnippet>872</a> <span class="prelude-val">Ok</span>::&lt;<span class="kw">_</span>, error::Error&gt;(max.max(estimate_value_size(values.as_ref(), row)<span class="question-mark">?</span>))
<a href=#873 id=873 data-nosnippet>873</a> })<span class="question-mark">?</span>;
<a href=#874 id=874 data-nosnippet>874</a> max.saturating_add(<span class="number">16</span>).max(<span class="number">32</span>)
<a href=#875 id=875 data-nosnippet>875</a> } <span class="kw">else </span>{
<a href=#876 id=876 data-nosnippet>876</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(<span class="prelude-val">None</span>);
<a href=#877 id=877 data-nosnippet>877</a> };
<a href=#878 id=878 data-nosnippet>878</a> <span class="kw">let </span>expansion = <span class="kw">if </span>json_columns.contains(<span class="kw-2">&amp;</span>index) { <span class="number">8 </span>} <span class="kw">else </span>{ <span class="number">1 </span>};
<a href=#879 id=879 data-nosnippet>879</a> max_row_bytes = max_row_bytes.saturating_add(size.saturating_mul(expansion));
<a href=#880 id=880 data-nosnippet>880</a> }
<a href=#881 id=881 data-nosnippet>881</a> <span class="prelude-val">Ok</span>(max_row_bytes.checked_mul(rows))
<a href=#882 id=882 data-nosnippet>882</a>}
<a href=#883 id=883 data-nosnippet>883</a>
<a href=#884 id=884 data-nosnippet>884</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>rows_within_budget(
<a href=#885 id=885 data-nosnippet>885</a> batch: <span class="kw-2">&amp;</span>RecordBatch,
<a href=#886 id=886 data-nosnippet>886</a> start: usize,
<a href=#887 id=887 data-nosnippet>887</a> end: usize,
<a href=#888 id=888 data-nosnippet>888</a> budget: usize,
<a href=#889 id=889 data-nosnippet>889</a> json_columns: <span class="kw-2">&amp;</span>[usize],
<a href=#890 id=890 data-nosnippet>890</a>) -&gt; <span class="prelude-ty">Result</span>&lt;(usize, usize)&gt; {
<a href=#891 id=891 data-nosnippet>891</a> <span class="kw">if </span>start &lt; end
<a href=#892 id=892 data-nosnippet>892</a> &amp;&amp; <span class="kw">let </span><span class="prelude-val">Some</span>(bytes) = estimate_batch_bytes_upper_bound(batch, end - start, json_columns)<span class="question-mark">?
<a href=#893 id=893 data-nosnippet>893</a> </span>&amp;&amp; bytes &lt;= budget
<a href=#894 id=894 data-nosnippet>894</a> {
<a href=#895 id=895 data-nosnippet>895</a> <span class="kw">return </span><span class="prelude-val">Ok</span>((end - start, bytes));
<a href=#896 id=896 data-nosnippet>896</a> }
<a href=#897 id=897 data-nosnippet>897</a> <span class="kw">let </span><span class="kw-2">mut </span>bytes = <span class="number">0usize</span>;
<a href=#898 id=898 data-nosnippet>898</a> <span class="kw">let </span><span class="kw-2">mut </span>row = start;
<a href=#899 id=899 data-nosnippet>899</a> <span class="kw">while </span>row &lt; end {
<a href=#900 id=900 data-nosnippet>900</a> <span class="kw">let </span>row_bytes =
<a href=#901 id=901 data-nosnippet>901</a> batch
<a href=#902 id=902 data-nosnippet>902</a> .columns()
<a href=#903 id=903 data-nosnippet>903</a> .iter()
<a href=#904 id=904 data-nosnippet>904</a> .enumerate()
<a href=#905 id=905 data-nosnippet>905</a> .try_fold(<span class="number">0usize</span>, |sum, (index, array)| {
<a href=#906 id=906 data-nosnippet>906</a> <span class="comment">// JSON escaping and number formatting only expand JSON values.
<a href=#907 id=907 data-nosnippet>907</a> </span><span class="kw">let </span>expansion = <span class="kw">if </span>json_columns.contains(<span class="kw-2">&amp;</span>index) &amp;&amp; !array.is_null(row) {
<a href=#908 id=908 data-nosnippet>908</a> <span class="number">8
<a href=#909 id=909 data-nosnippet>909</a> </span>} <span class="kw">else </span>{
<a href=#910 id=910 data-nosnippet>910</a> <span class="number">1
<a href=#911 id=911 data-nosnippet>911</a> </span>};
<a href=#912 id=912 data-nosnippet>912</a> <span class="kw">let </span>size = estimate_value_size(array.as_ref(), row)<span class="question-mark">?</span>;
<a href=#913 id=913 data-nosnippet>913</a> <span class="prelude-val">Ok</span>::&lt;<span class="kw">_</span>, error::Error&gt;(sum.saturating_add(size.saturating_mul(expansion)))
<a href=#914 id=914 data-nosnippet>914</a> })<span class="question-mark">?</span>;
<a href=#915 id=915 data-nosnippet>915</a> <span class="kw">if </span>row_bytes &gt; budget.saturating_sub(bytes) {
<a href=#916 id=916 data-nosnippet>916</a> <span class="kw">break</span>;
<a href=#917 id=917 data-nosnippet>917</a> }
<a href=#918 id=918 data-nosnippet>918</a> bytes += row_bytes;
<a href=#919 id=919 data-nosnippet>919</a> row += <span class="number">1</span>;
<a href=#920 id=920 data-nosnippet>920</a> }
<a href=#921 id=921 data-nosnippet>921</a> <span class="macro">ensure!</span>(
<a href=#922 id=922 data-nosnippet>922</a> row &gt; start,
<a href=#923 id=923 data-nosnippet>923</a> LogicalTableExportResourceSnafu {
<a href=#924 id=924 data-nosnippet>924</a> reason: <span class="string">"one expanded logical row exceeds conversion byte budget"
<a href=#925 id=925 data-nosnippet>925</a> </span>}
<a href=#926 id=926 data-nosnippet>926</a> );
<a href=#927 id=927 data-nosnippet>927</a> <span class="prelude-val">Ok</span>((row - start, bytes))
<a href=#928 id=928 data-nosnippet>928</a>}
<a href=#929 id=929 data-nosnippet>929</a>
<a href=#930 id=930 data-nosnippet>930</a><span class="attr">#[cfg(test)]
<a href=#931 id=931 data-nosnippet>931</a></span><span class="kw">mod </span>tests;
</code></pre></div></section></main></body></html>
@@ -0,0 +1,318 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/export_logical_tables/writers.rs`."><title>writers.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../../" data-static-root-path="../../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../../static.files/src-script-813739b1.js"></script><script defer src="../../../../src-files.js"></script><script defer src="../../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/export_logical_tables/</div>writers.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::sync::Arc;
<a href=#16 id=16 data-nosnippet>16</a>
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>arrow::record_batch::RecordBatch;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_datasource::packed_writer::PackedWriterRef;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>futures::stream::FuturesUnordered;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>futures::{FutureExt, StreamExt};
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>object_store::ObjectStore;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>snafu::{OptionExt, ResultExt, ensure};
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>tokio::sync::{OwnedSemaphorePermit, Semaphore, mpsc};
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>tokio::task::JoinHandle;
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>tokio_util::sync::CancellationToken;
<a href=#26 id=26 data-nosnippet>26</a>
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, <span class="prelude-ty">Result</span>};
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span><span class="kw">crate</span>::statement::export_logical_tables::{
<a href=#29 id=29 data-nosnippet>29</a> ActiveWriter, LogicalTableExportLimits, LogicalTableProjection, check_cancelled,
<a href=#30 id=30 data-nosnippet>30</a> map_writer_error,
<a href=#31 id=31 data-nosnippet>31</a>};
<a href=#32 id=32 data-nosnippet>32</a>
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">const </span>PAYLOAD_BYTES: usize = <span class="number">64 </span>* <span class="number">1024 </span>* <span class="number">1024</span>;
<a href=#34 id=34 data-nosnippet>34</a>
<a href=#35 id=35 data-nosnippet>35</a><span class="doccomment">/// Request-wide downstream payload and writer admission, separate from query memory.
<a href=#36 id=36 data-nosnippet>36</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>ExportWriteBudget {
<a href=#37 id=37 data-nosnippet>37</a> writers: Arc&lt;Semaphore&gt;,
<a href=#38 id=38 data-nosnippet>38</a> bytes: Arc&lt;Semaphore&gt;,
<a href=#39 id=39 data-nosnippet>39</a> max_writers: usize,
<a href=#40 id=40 data-nosnippet>40</a>}
<a href=#41 id=41 data-nosnippet>41</a>
<a href=#42 id=42 data-nosnippet>42</a><span class="kw">impl </span>ExportWriteBudget {
<a href=#43 id=43 data-nosnippet>43</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>new(parallelism: usize) -&gt; Arc&lt;<span class="self">Self</span>&gt; {
<a href=#44 id=44 data-nosnippet>44</a> <span class="kw">let </span>max_writers = parallelism.max(<span class="number">1</span>);
<a href=#45 id=45 data-nosnippet>45</a> Arc::new(<span class="self">Self </span>{
<a href=#46 id=46 data-nosnippet>46</a> writers: Arc::new(Semaphore::new(max_writers)),
<a href=#47 id=47 data-nosnippet>47</a> bytes: Arc::new(Semaphore::new(PAYLOAD_BYTES)),
<a href=#48 id=48 data-nosnippet>48</a> max_writers,
<a href=#49 id=49 data-nosnippet>49</a> })
<a href=#50 id=50 data-nosnippet>50</a> }
<a href=#51 id=51 data-nosnippet>51</a>
<a href=#52 id=52 data-nosnippet>52</a> <span class="attr">#[cfg(test)]
<a href=#53 id=53 data-nosnippet>53</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>available(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; (usize, usize) {
<a href=#54 id=54 data-nosnippet>54</a> (
<a href=#55 id=55 data-nosnippet>55</a> <span class="self">self</span>.writers.available_permits(),
<a href=#56 id=56 data-nosnippet>56</a> <span class="self">self</span>.bytes.available_permits(),
<a href=#57 id=57 data-nosnippet>57</a> )
<a href=#58 id=58 data-nosnippet>58</a> }
<a href=#59 id=59 data-nosnippet>59</a>
<a href=#60 id=60 data-nosnippet>60</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>writer(<span class="kw-2">&amp;</span><span class="self">self</span>, token: <span class="kw-2">&amp;</span>CancellationToken) -&gt; <span class="prelude-ty">Result</span>&lt;OwnedSemaphorePermit&gt; {
<a href=#61 id=61 data-nosnippet>61</a> <span class="macro">tokio::select!</span> {
<a href=#62 id=62 data-nosnippet>62</a> biased;
<a href=#63 id=63 data-nosnippet>63</a> <span class="kw">_ </span>= token.cancelled() =&gt; error::LogicalTableExportCancelledSnafu.fail(),
<a href=#64 id=64 data-nosnippet>64</a> permit = <span class="self">self</span>.writers.clone().acquire_owned() =&gt; permit.map_err(|<span class="kw">_</span>| error::UnexpectedSnafu { violated: <span class="string">"writer budget closed" </span>}.build()),
<a href=#65 id=65 data-nosnippet>65</a> }
<a href=#66 id=66 data-nosnippet>66</a> }
<a href=#67 id=67 data-nosnippet>67</a>
<a href=#68 id=68 data-nosnippet>68</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>reserve(
<a href=#69 id=69 data-nosnippet>69</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#70 id=70 data-nosnippet>70</a> size: usize,
<a href=#71 id=71 data-nosnippet>71</a> token: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#72 id=72 data-nosnippet>72</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;OwnedSemaphorePermit&gt; {
<a href=#73 id=73 data-nosnippet>73</a> <span class="macro">ensure!</span>(
<a href=#74 id=74 data-nosnippet>74</a> size &lt;= PAYLOAD_BYTES,
<a href=#75 id=75 data-nosnippet>75</a> error::LogicalTableExportResourceSnafu {
<a href=#76 id=76 data-nosnippet>76</a> reason: <span class="string">"batch backing buffers exceed request payload budget"
<a href=#77 id=77 data-nosnippet>77</a> </span>}
<a href=#78 id=78 data-nosnippet>78</a> );
<a href=#79 id=79 data-nosnippet>79</a> <span class="macro">tokio::select!</span> {
<a href=#80 id=80 data-nosnippet>80</a> biased;
<a href=#81 id=81 data-nosnippet>81</a> <span class="kw">_ </span>= token.cancelled() =&gt; error::LogicalTableExportCancelledSnafu.fail(),
<a href=#82 id=82 data-nosnippet>82</a> permit = <span class="self">self</span>.bytes.clone().acquire_many_owned(size <span class="kw">as </span>u32) =&gt; permit.map_err(|<span class="kw">_</span>| error::UnexpectedSnafu { violated: <span class="string">"payload budget closed" </span>}.build()),
<a href=#83 id=83 data-nosnippet>83</a> }
<a href=#84 id=84 data-nosnippet>84</a> }
<a href=#85 id=85 data-nosnippet>85</a>
<a href=#86 id=86 data-nosnippet>86</a> <span class="doccomment">/// Accounts for retained input allocations, conversion capacity growth and
<a href=#87 id=87 data-nosnippet>87</a> /// array metadata. Pass zero backing when the caller detaches oversized slices.
<a href=#88 id=88 data-nosnippet>88</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>conversion_budget(
<a href=#89 id=89 data-nosnippet>89</a> backing: usize,
<a href=#90 id=90 data-nosnippet>90</a> columns: usize,
<a href=#91 id=91 data-nosnippet>91</a> requested: usize,
<a href=#92 id=92 data-nosnippet>92</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;(usize, usize)&gt; {
<a href=#93 id=93 data-nosnippet>93</a> <span class="kw">let </span>overhead = columns.saturating_mul(<span class="number">1024</span>);
<a href=#94 id=94 data-nosnippet>94</a> <span class="kw">let </span>available = PAYLOAD_BYTES
<a href=#95 id=95 data-nosnippet>95</a> .saturating_sub(backing)
<a href=#96 id=96 data-nosnippet>96</a> .saturating_sub(overhead)
<a href=#97 id=97 data-nosnippet>97</a> / <span class="number">4</span>;
<a href=#98 id=98 data-nosnippet>98</a> <span class="kw">let </span>conversion = requested.min(available);
<a href=#99 id=99 data-nosnippet>99</a> <span class="macro">ensure!</span>(
<a href=#100 id=100 data-nosnippet>100</a> conversion &gt; <span class="number">0</span>,
<a href=#101 id=101 data-nosnippet>101</a> error::LogicalTableExportResourceSnafu {
<a href=#102 id=102 data-nosnippet>102</a> reason: <span class="string">"batch backing buffers exceed request payload budget"
<a href=#103 id=103 data-nosnippet>103</a> </span>}
<a href=#104 id=104 data-nosnippet>104</a> );
<a href=#105 id=105 data-nosnippet>105</a> <span class="prelude-val">Ok</span>((conversion, backing.saturating_add(overhead)))
<a href=#106 id=106 data-nosnippet>106</a> }
<a href=#107 id=107 data-nosnippet>107</a>}
<a href=#108 id=108 data-nosnippet>108</a>
<a href=#109 id=109 data-nosnippet>109</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>Payload {
<a href=#110 id=110 data-nosnippet>110</a> <span class="kw">pub</span>(<span class="kw">crate</span>) batch: RecordBatch,
<a href=#111 id=111 data-nosnippet>111</a> <span class="kw">pub</span>(<span class="kw">crate</span>) permit: OwnedSemaphorePermit,
<a href=#112 id=112 data-nosnippet>112</a>}
<a href=#113 id=113 data-nosnippet>113</a>
<a href=#114 id=114 data-nosnippet>114</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>TableWriters {
<a href=#115 id=115 data-nosnippet>115</a> <span class="kw">pub</span>(<span class="kw">crate</span>) packed: <span class="prelude-ty">Option</span>&lt;PackedWriterRef&gt;,
<a href=#116 id=116 data-nosnippet>116</a> current: <span class="prelude-ty">Option</span>&lt;(u32, mpsc::Sender&lt;Payload&gt;)&gt;,
<a href=#117 id=117 data-nosnippet>117</a> tasks: FuturesUnordered&lt;JoinHandle&lt;<span class="prelude-ty">Result</span>&lt;()&gt;&gt;&gt;,
<a href=#118 id=118 data-nosnippet>118</a> budget: Arc&lt;ExportWriteBudget&gt;,
<a href=#119 id=119 data-nosnippet>119</a>}
<a href=#120 id=120 data-nosnippet>120</a>
<a href=#121 id=121 data-nosnippet>121</a><span class="kw">impl </span>TableWriters {
<a href=#122 id=122 data-nosnippet>122</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>new(budget: Arc&lt;ExportWriteBudget&gt;) -&gt; <span class="self">Self </span>{
<a href=#123 id=123 data-nosnippet>123</a> <span class="self">Self </span>{
<a href=#124 id=124 data-nosnippet>124</a> current: <span class="prelude-val">None</span>,
<a href=#125 id=125 data-nosnippet>125</a> packed: <span class="prelude-val">None</span>,
<a href=#126 id=126 data-nosnippet>126</a> tasks: FuturesUnordered::new(),
<a href=#127 id=127 data-nosnippet>127</a> budget,
<a href=#128 id=128 data-nosnippet>128</a> }
<a href=#129 id=129 data-nosnippet>129</a> }
<a href=#130 id=130 data-nosnippet>130</a>
<a href=#131 id=131 data-nosnippet>131</a> <span class="attr">#[cfg(test)]
<a href=#132 id=132 data-nosnippet>132</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>pending_tasks(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; usize {
<a href=#133 id=133 data-nosnippet>133</a> <span class="self">self</span>.tasks.len()
<a href=#134 id=134 data-nosnippet>134</a> }
<a href=#135 id=135 data-nosnippet>135</a>
<a href=#136 id=136 data-nosnippet>136</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>table_id(<span class="kw-2">&amp;</span><span class="self">self</span>) -&gt; <span class="prelude-ty">Option</span>&lt;u32&gt; {
<a href=#137 id=137 data-nosnippet>137</a> <span class="self">self</span>.current.as_ref().map(|(id, <span class="kw">_</span>)| <span class="kw-2">*</span>id)
<a href=#138 id=138 data-nosnippet>138</a> }
<a href=#139 id=139 data-nosnippet>139</a>
<a href=#140 id=140 data-nosnippet>140</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>close_input(<span class="kw-2">&amp;mut </span><span class="self">self</span>) {
<a href=#141 id=141 data-nosnippet>141</a> <span class="self">self</span>.current = <span class="prelude-val">None</span>;
<a href=#142 id=142 data-nosnippet>142</a> }
<a href=#143 id=143 data-nosnippet>143</a>
<a href=#144 id=144 data-nosnippet>144</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>open(
<a href=#145 id=145 data-nosnippet>145</a> <span class="kw-2">&amp;mut </span><span class="self">self</span>,
<a href=#146 id=146 data-nosnippet>146</a> id: u32,
<a href=#147 id=147 data-nosnippet>147</a> table: <span class="kw-2">&amp;</span>LogicalTableProjection,
<a href=#148 id=148 data-nosnippet>148</a> store: <span class="kw-2">&amp;</span>ObjectStore,
<a href=#149 id=149 data-nosnippet>149</a> limits: LogicalTableExportLimits,
<a href=#150 id=150 data-nosnippet>150</a> token: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#151 id=151 data-nosnippet>151</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#152 id=152 data-nosnippet>152</a> <span class="comment">// EOF must precede acquiring the next slot, including when P is one.
<a href=#153 id=153 data-nosnippet>153</a> </span><span class="self">self</span>.close_input();
<a href=#154 id=154 data-nosnippet>154</a> <span class="kw">let </span>permit = <span class="self">self</span>.budget.writer(token).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#155 id=155 data-nosnippet>155</a> <span class="self">self</span>.reap_for_admission(token).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#156 id=156 data-nosnippet>156</a> <span class="kw">let </span>writer = <span class="kw">if let </span><span class="prelude-val">Some</span>(packed) = <span class="kw-2">&amp;</span><span class="self">self</span>.packed {
<a href=#157 id=157 data-nosnippet>157</a> ActiveWriter::open_packed(table, id, store, limits, packed.clone())<span class="question-mark">?
<a href=#158 id=158 data-nosnippet>158</a> </span>} <span class="kw">else </span>{
<a href=#159 id=159 data-nosnippet>159</a> ActiveWriter::open(table, store, limits).<span class="kw">await</span><span class="question-mark">?
<a href=#160 id=160 data-nosnippet>160</a> </span>};
<a href=#161 id=161 data-nosnippet>161</a> <span class="kw">let </span>(sender, receiver) = mpsc::channel(<span class="number">2</span>);
<a href=#162 id=162 data-nosnippet>162</a> <span class="self">self</span>.current = <span class="prelude-val">Some</span>((id, sender));
<a href=#163 id=163 data-nosnippet>163</a> <span class="kw">let </span>token = token.clone();
<a href=#164 id=164 data-nosnippet>164</a> <span class="self">self</span>.tasks.push(common_runtime::spawn_global(<span class="kw">async move </span>{
<a href=#165 id=165 data-nosnippet>165</a> <span class="kw">let </span>_permit = permit;
<a href=#166 id=166 data-nosnippet>166</a> <span class="kw">let </span>guard = token.clone().drop_guard();
<a href=#167 id=167 data-nosnippet>167</a> <span class="kw">let </span>result = run_writer(writer, receiver, <span class="kw-2">&amp;</span>token).<span class="kw">await</span>;
<a href=#168 id=168 data-nosnippet>168</a> <span class="kw">if </span>result.is_ok() {
<a href=#169 id=169 data-nosnippet>169</a> guard.disarm();
<a href=#170 id=170 data-nosnippet>170</a> }
<a href=#171 id=171 data-nosnippet>171</a> result
<a href=#172 id=172 data-nosnippet>172</a> }));
<a href=#173 id=173 data-nosnippet>173</a> <span class="prelude-val">Ok</span>(())
<a href=#174 id=174 data-nosnippet>174</a> }
<a href=#175 id=175 data-nosnippet>175</a>
<a href=#176 id=176 data-nosnippet>176</a> <span class="kw">async fn </span>reap_for_admission(<span class="kw-2">&amp;mut </span><span class="self">self</span>, token: <span class="kw-2">&amp;</span>CancellationToken) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#177 id=177 data-nosnippet>177</a> <span class="kw">while let </span><span class="prelude-val">Some</span>(<span class="prelude-val">Some</span>(result)) = <span class="self">self</span>.tasks.next().now_or_never() {
<a href=#178 id=178 data-nosnippet>178</a> result.context(error::JoinTaskSnafu)<span class="question-mark">??</span>;
<a href=#179 id=179 data-nosnippet>179</a> }
<a href=#180 id=180 data-nosnippet>180</a> <span class="comment">// A worker can release its permit before its JoinHandle becomes ready.
<a href=#181 id=181 data-nosnippet>181</a> </span><span class="kw">while </span><span class="self">self</span>.tasks.len() &gt; <span class="self">self</span>.budget.max_writers {
<a href=#182 id=182 data-nosnippet>182</a> <span class="kw">let </span>result = <span class="macro">tokio::select!</span> {
<a href=#183 id=183 data-nosnippet>183</a> biased;
<a href=#184 id=184 data-nosnippet>184</a> <span class="kw">_ </span>= token.cancelled() =&gt; <span class="kw">return </span>error::LogicalTableExportCancelledSnafu.fail(),
<a href=#185 id=185 data-nosnippet>185</a> result = <span class="self">self</span>.tasks.next() =&gt; result.context(error::UnexpectedSnafu { violated: <span class="string">"writer task queue unexpectedly empty" </span>})<span class="question-mark">?</span>,
<a href=#186 id=186 data-nosnippet>186</a> };
<a href=#187 id=187 data-nosnippet>187</a> result.context(error::JoinTaskSnafu)<span class="question-mark">??</span>;
<a href=#188 id=188 data-nosnippet>188</a> }
<a href=#189 id=189 data-nosnippet>189</a> <span class="prelude-val">Ok</span>(())
<a href=#190 id=190 data-nosnippet>190</a> }
<a href=#191 id=191 data-nosnippet>191</a>
<a href=#192 id=192 data-nosnippet>192</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>send(<span class="kw-2">&amp;</span><span class="self">self</span>, payload: Payload, token: <span class="kw-2">&amp;</span>CancellationToken) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#193 id=193 data-nosnippet>193</a> <span class="kw">let </span>(<span class="kw">_</span>, sender) = <span class="self">self</span>.current.as_ref().context(error::UnexpectedSnafu {
<a href=#194 id=194 data-nosnippet>194</a> violated: <span class="string">"missing logical writer"</span>,
<a href=#195 id=195 data-nosnippet>195</a> })<span class="question-mark">?</span>;
<a href=#196 id=196 data-nosnippet>196</a> <span class="macro">tokio::select!</span> {
<a href=#197 id=197 data-nosnippet>197</a> biased;
<a href=#198 id=198 data-nosnippet>198</a> <span class="kw">_ </span>= token.cancelled() =&gt; error::LogicalTableExportCancelledSnafu.fail(),
<a href=#199 id=199 data-nosnippet>199</a> result = sender.send(payload) =&gt; result.map_err(|<span class="kw">_</span>| error::LogicalTableExportCancelledSnafu.build()),
<a href=#200 id=200 data-nosnippet>200</a> }
<a href=#201 id=201 data-nosnippet>201</a> }
<a href=#202 id=202 data-nosnippet>202</a>
<a href=#203 id=203 data-nosnippet>203</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>drain(
<a href=#204 id=204 data-nosnippet>204</a> <span class="kw-2">&amp;mut </span><span class="self">self</span>,
<a href=#205 id=205 data-nosnippet>205</a> result: <span class="prelude-ty">Result</span>&lt;()&gt;,
<a href=#206 id=206 data-nosnippet>206</a> token: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#207 id=207 data-nosnippet>207</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#208 id=208 data-nosnippet>208</a> <span class="self">self</span>.close_input();
<a href=#209 id=209 data-nosnippet>209</a> <span class="kw">let </span><span class="kw-2">mut </span>first_error = result.err();
<a href=#210 id=210 data-nosnippet>210</a> <span class="kw">if </span>first_error.is_some() {
<a href=#211 id=211 data-nosnippet>211</a> token.cancel();
<a href=#212 id=212 data-nosnippet>212</a> }
<a href=#213 id=213 data-nosnippet>213</a> <span class="kw">while let </span><span class="prelude-val">Some</span>(result) = <span class="self">self</span>.tasks.next().<span class="kw">await </span>{
<a href=#214 id=214 data-nosnippet>214</a> <span class="kw">if let </span><span class="prelude-val">Err</span>(err) = result.context(error::JoinTaskSnafu).and_then(|r| r) {
<a href=#215 id=215 data-nosnippet>215</a> retain_error(<span class="kw-2">&amp;mut </span>first_error, err);
<a href=#216 id=216 data-nosnippet>216</a> token.cancel();
<a href=#217 id=217 data-nosnippet>217</a> }
<a href=#218 id=218 data-nosnippet>218</a> }
<a href=#219 id=219 data-nosnippet>219</a> <span class="kw">match </span>first_error {
<a href=#220 id=220 data-nosnippet>220</a> <span class="prelude-val">Some</span>(err) =&gt; <span class="prelude-val">Err</span>(err),
<a href=#221 id=221 data-nosnippet>221</a> <span class="prelude-val">None </span>=&gt; <span class="prelude-val">Ok</span>(()),
<a href=#222 id=222 data-nosnippet>222</a> }
<a href=#223 id=223 data-nosnippet>223</a> }
<a href=#224 id=224 data-nosnippet>224</a>}
<a href=#225 id=225 data-nosnippet>225</a>
<a href=#226 id=226 data-nosnippet>226</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>retain_error(first: <span class="kw-2">&amp;mut </span><span class="prelude-ty">Option</span>&lt;error::Error&gt;, error: error::Error) {
<a href=#227 id=227 data-nosnippet>227</a> <span class="kw">if </span>first.as_ref().is_none_or(|err| {
<a href=#228 id=228 data-nosnippet>228</a> <span class="macro">matches!</span>(
<a href=#229 id=229 data-nosnippet>229</a> err,
<a href=#230 id=230 data-nosnippet>230</a> error::Error::LogicalTableExportCancelled { .. }
<a href=#231 id=231 data-nosnippet>231</a> | error::Error::DatabaseExportCancelled { .. }
<a href=#232 id=232 data-nosnippet>232</a> )
<a href=#233 id=233 data-nosnippet>233</a> }) {
<a href=#234 id=234 data-nosnippet>234</a> <span class="kw-2">*</span>first = <span class="prelude-val">Some</span>(error);
<a href=#235 id=235 data-nosnippet>235</a> }
<a href=#236 id=236 data-nosnippet>236</a>}
<a href=#237 id=237 data-nosnippet>237</a>
<a href=#238 id=238 data-nosnippet>238</a><span class="kw">async fn </span>run_writer(
<a href=#239 id=239 data-nosnippet>239</a> <span class="kw-2">mut </span>writer: ActiveWriter,
<a href=#240 id=240 data-nosnippet>240</a> <span class="kw-2">mut </span>receiver: mpsc::Receiver&lt;Payload&gt;,
<a href=#241 id=241 data-nosnippet>241</a> token: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#242 id=242 data-nosnippet>242</a>) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#243 id=243 data-nosnippet>243</a> <span class="kw">let </span>result = <span class="kw">async </span>{
<a href=#244 id=244 data-nosnippet>244</a> <span class="kw">loop </span>{
<a href=#245 id=245 data-nosnippet>245</a> <span class="kw">let </span>payload = <span class="macro">tokio::select!</span> {
<a href=#246 id=246 data-nosnippet>246</a> biased;
<a href=#247 id=247 data-nosnippet>247</a> <span class="kw">_ </span>= token.cancelled() =&gt; <span class="kw">return </span>error::LogicalTableExportCancelledSnafu.fail(),
<a href=#248 id=248 data-nosnippet>248</a> payload = receiver.recv() =&gt; payload,
<a href=#249 id=249 data-nosnippet>249</a> };
<a href=#250 id=250 data-nosnippet>250</a> <span class="kw">let </span><span class="prelude-val">Some</span>(Payload { batch, permit }) = payload <span class="kw">else </span>{
<a href=#251 id=251 data-nosnippet>251</a> <span class="kw">break</span>;
<a href=#252 id=252 data-nosnippet>252</a> };
<a href=#253 id=253 data-nosnippet>253</a> <span class="kw">let </span>result = writer.writer.write(batch, <span class="prelude-val">Some</span>(token)).<span class="kw">await</span>;
<a href=#254 id=254 data-nosnippet>254</a> drop(permit);
<a href=#255 id=255 data-nosnippet>255</a> result.map_err(|error| map_writer_error(error, <span class="kw-2">&amp;</span>writer.path))<span class="question-mark">?</span>;
<a href=#256 id=256 data-nosnippet>256</a> }
<a href=#257 id=257 data-nosnippet>257</a> writer
<a href=#258 id=258 data-nosnippet>258</a> .writer
<a href=#259 id=259 data-nosnippet>259</a> .finish(<span class="prelude-val">Some</span>(token))
<a href=#260 id=260 data-nosnippet>260</a> .<span class="kw">await
<a href=#261 id=261 data-nosnippet>261</a> </span>.map_err(|error| map_writer_error(error, <span class="kw-2">&amp;</span>writer.path))<span class="question-mark">?</span>;
<a href=#262 id=262 data-nosnippet>262</a> check_cancelled(token)
<a href=#263 id=263 data-nosnippet>263</a> }
<a href=#264 id=264 data-nosnippet>264</a> .<span class="kw">await</span>;
<a href=#265 id=265 data-nosnippet>265</a> <span class="kw">if </span>result.is_err() {
<a href=#266 id=266 data-nosnippet>266</a> token.cancel();
<a href=#267 id=267 data-nosnippet>267</a> receiver.close();
<a href=#268 id=268 data-nosnippet>268</a> drop(receiver);
<a href=#269 id=269 data-nosnippet>269</a> <span class="kw">if let </span><span class="prelude-val">Err</span>(error) = writer.writer.abort().<span class="kw">await </span>{
<a href=#270 id=270 data-nosnippet>270</a> <span class="macro">common_telemetry::warn!</span>(error; <span class="string">"Failed to abort Metric export file"</span>);
<a href=#271 id=271 data-nosnippet>271</a> }
<a href=#272 id=272 data-nosnippet>272</a> }
<a href=#273 id=273 data-nosnippet>273</a> result
<a href=#274 id=274 data-nosnippet>274</a>}
<a href=#275 id=275 data-nosnippet>275</a>
<a href=#276 id=276 data-nosnippet>276</a><span class="attr">#[cfg(test)]
<a href=#277 id=277 data-nosnippet>277</a></span><span class="kw">mod </span>tests {
<a href=#278 id=278 data-nosnippet>278</a> <span class="kw">use </span>tokio::sync::oneshot;
<a href=#279 id=279 data-nosnippet>279</a>
<a href=#280 id=280 data-nosnippet>280</a> <span class="kw">use super</span>::<span class="kw-2">*</span>;
<a href=#281 id=281 data-nosnippet>281</a>
<a href=#282 id=282 data-nosnippet>282</a> <span class="attr">#[tokio::test]
<a href=#283 id=283 data-nosnippet>283</a> </span><span class="kw">async fn </span>released_permit_does_not_hide_pending_join_handles() {
<a href=#284 id=284 data-nosnippet>284</a> <span class="kw">for </span>parallelism <span class="kw">in </span>[<span class="number">1</span>, <span class="number">4</span>] {
<a href=#285 id=285 data-nosnippet>285</a> <span class="kw">let </span>budget = ExportWriteBudget::new(parallelism);
<a href=#286 id=286 data-nosnippet>286</a> <span class="kw">let </span>token = CancellationToken::new();
<a href=#287 id=287 data-nosnippet>287</a> <span class="kw">let </span><span class="kw-2">mut </span>writers = TableWriters::new(budget.clone());
<a href=#288 id=288 data-nosnippet>288</a> <span class="kw">let </span><span class="kw-2">mut </span>resume = Vec::new();
<a href=#289 id=289 data-nosnippet>289</a> <span class="kw">for _ in </span><span class="number">0</span>..=parallelism {
<a href=#290 id=290 data-nosnippet>290</a> <span class="kw">let </span>permit = budget.writer(<span class="kw-2">&amp;</span>token).<span class="kw">await</span>.unwrap();
<a href=#291 id=291 data-nosnippet>291</a> <span class="kw">let </span>(released_tx, released_rx) = oneshot::channel();
<a href=#292 id=292 data-nosnippet>292</a> <span class="kw">let </span>(resume_tx, resume_rx) = oneshot::channel();
<a href=#293 id=293 data-nosnippet>293</a> writers.tasks.push(common_runtime::spawn_global(<span class="kw">async move </span>{
<a href=#294 id=294 data-nosnippet>294</a> drop(permit);
<a href=#295 id=295 data-nosnippet>295</a> released_tx.send(()).unwrap();
<a href=#296 id=296 data-nosnippet>296</a> resume_rx.<span class="kw">await</span>.unwrap();
<a href=#297 id=297 data-nosnippet>297</a> <span class="prelude-val">Ok</span>(())
<a href=#298 id=298 data-nosnippet>298</a> }));
<a href=#299 id=299 data-nosnippet>299</a> released_rx.<span class="kw">await</span>.unwrap();
<a href=#300 id=300 data-nosnippet>300</a> resume.push(resume_tx);
<a href=#301 id=301 data-nosnippet>301</a> }
<a href=#302 id=302 data-nosnippet>302</a> <span class="macro">assert_eq!</span>(budget.available().<span class="number">0</span>, parallelism);
<a href=#303 id=303 data-nosnippet>303</a>
<a href=#304 id=304 data-nosnippet>304</a> <span class="kw">let </span><span class="kw-2">mut </span>reap = Box::pin(writers.reap_for_admission(<span class="kw-2">&amp;</span>token));
<a href=#305 id=305 data-nosnippet>305</a> <span class="macro">assert!</span>(reap.as_mut().now_or_never().is_none());
<a href=#306 id=306 data-nosnippet>306</a> resume.pop().unwrap().send(()).unwrap();
<a href=#307 id=307 data-nosnippet>307</a> reap.<span class="kw">await</span>.unwrap();
<a href=#308 id=308 data-nosnippet>308</a> <span class="macro">assert_eq!</span>(writers.pending_tasks(), parallelism);
<a href=#309 id=309 data-nosnippet>309</a>
<a href=#310 id=310 data-nosnippet>310</a> <span class="kw">for </span>tx <span class="kw">in </span>resume {
<a href=#311 id=311 data-nosnippet>311</a> tx.send(()).unwrap();
<a href=#312 id=312 data-nosnippet>312</a> }
<a href=#313 id=313 data-nosnippet>313</a> writers.drain(<span class="prelude-val">Ok</span>(()), <span class="kw-2">&amp;</span>token).<span class="kw">await</span>.unwrap();
<a href=#314 id=314 data-nosnippet>314</a> <span class="macro">assert_eq!</span>(budget.available().<span class="number">0</span>, parallelism);
<a href=#315 id=315 data-nosnippet>315</a> }
<a href=#316 id=316 data-nosnippet>316</a> }
<a href=#317 id=317 data-nosnippet>317</a>}
</code></pre></div></section></main></body></html>
@@ -0,0 +1,437 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/import_packed.rs`."><title>import_packed.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>import_packed.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::collections::HashMap;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>std::sync::Arc;
<a href=#17 id=17 data-nosnippet>17</a>
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_datasource::file_format::Format;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>common_datasource::file_format::parquet::packed_reader::{
<a href=#20 id=20 data-nosnippet>20</a> PackReadWindows, PackedParquetReader, WINDOW_SIZE,
<a href=#21 id=21 data-nosnippet>21</a>};
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>common_datasource::object_store::{BuiltBackend, build_backend_with_path};
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>common_datasource::packed_snapshot::{ObjectKind, PACK_INDEX_FILE, PackIndex};
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>common_error::ext::BoxedError;
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>common_query::Output;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>futures::stream::FuturesUnordered;
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>futures::{FutureExt, StreamExt};
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span>snafu::ResultExt;
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span>table::TableRef;
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">use </span>table::requests::CopyDatabaseRequest;
<a href=#32 id=32 data-nosnippet>32</a><span class="kw">use </span>tokio_util::compat::FuturesAsyncReadCompatExt;
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">use </span>tokio_util::sync::CancellationToken;
<a href=#34 id=34 data-nosnippet>34</a>
<a href=#35 id=35 data-nosnippet>35</a><span class="kw">use </span><span class="kw">crate</span>::error::{<span class="self">self</span>, <span class="prelude-ty">Result</span>};
<a href=#36 id=36 data-nosnippet>36</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#37 id=37 data-nosnippet>37</a><span class="kw">use </span><span class="kw">crate</span>::statement::database_copy::validate_database_directory;
<a href=#38 id=38 data-nosnippet>38</a>
<a href=#39 id=39 data-nosnippet>39</a><span class="doccomment">/// All table identities are resolved before the frontend authorizes any write.
<a href=#40 id=40 data-nosnippet>40</a></span><span class="kw">pub struct </span>PreparedPackedImport {
<a href=#41 id=41 data-nosnippet>41</a> backend: BuiltBackend,
<a href=#42 id=42 data-nosnippet>42</a> index: PackIndex,
<a href=#43 id=43 data-nosnippet>43</a> parallelism: usize,
<a href=#44 id=44 data-nosnippet>44</a> <span class="kw">pub </span>tables: Vec&lt;TableRef&gt;,
<a href=#45 id=45 data-nosnippet>45</a>}
<a href=#46 id=46 data-nosnippet>46</a>
<a href=#47 id=47 data-nosnippet>47</a><span class="kw">impl </span>StatementExecutor {
<a href=#48 id=48 data-nosnippet>48</a> <span class="kw">pub async fn </span>prepare_packed_import(
<a href=#49 id=49 data-nosnippet>49</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#50 id=50 data-nosnippet>50</a> req: CopyDatabaseRequest,
<a href=#51 id=51 data-nosnippet>51</a> ctx: <span class="kw-2">&amp;</span>QueryContextRef,
<a href=#52 id=52 data-nosnippet>52</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;PreparedPackedImport&gt; {
<a href=#53 id=53 data-nosnippet>53</a> validate_database_directory(<span class="kw-2">&amp;</span>req.location)<span class="question-mark">?</span>;
<a href=#54 id=54 data-nosnippet>54</a> <span class="kw">if </span>req.with.get(<span class="string">"metric_data_layout"</span>).map(String::as_str) != <span class="prelude-val">Some</span>(<span class="string">"packed"</span>)
<a href=#55 id=55 data-nosnippet>55</a> || !<span class="macro">matches!</span>(
<a href=#56 id=56 data-nosnippet>56</a> Format::try_from(<span class="kw-2">&amp;</span>req.with).context(error::ParseFileFormatSnafu)<span class="question-mark">?</span>,
<a href=#57 id=57 data-nosnippet>57</a> Format::Parquet(<span class="kw">_</span>)
<a href=#58 id=58 data-nosnippet>58</a> )
<a href=#59 id=59 data-nosnippet>59</a> || req
<a href=#60 id=60 data-nosnippet>60</a> .with
<a href=#61 id=61 data-nosnippet>61</a> .get(<span class="string">"continue_on_error"</span>)
<a href=#62 id=62 data-nosnippet>62</a> .is_some_and(|v| v != <span class="string">"false"</span>)
<a href=#63 id=63 data-nosnippet>63</a> {
<a href=#64 id=64 data-nosnippet>64</a> <span class="kw">return </span>error::InvalidCopyParameterSnafu {
<a href=#65 id=65 data-nosnippet>65</a> key: <span class="string">"metric_data_layout"</span>,
<a href=#66 id=66 data-nosnippet>66</a> value: <span class="string">"packed requires parquet and continue_on_error=false"</span>,
<a href=#67 id=67 data-nosnippet>67</a> }
<a href=#68 id=68 data-nosnippet>68</a> .fail();
<a href=#69 id=69 data-nosnippet>69</a> }
<a href=#70 id=70 data-nosnippet>70</a> <span class="kw">let </span>backend =
<a href=#71 id=71 data-nosnippet>71</a> build_backend_with_path(<span class="kw-2">&amp;</span>req.location, <span class="kw-2">&amp;</span>req.connection, <span class="kw-2">&amp;</span><span class="self">self</span>.local_file_access)
<a href=#72 id=72 data-nosnippet>72</a> .<span class="kw">await
<a href=#73 id=73 data-nosnippet>73</a> </span>.context(error::BuildBackendSnafu)<span class="question-mark">?</span>;
<a href=#74 id=74 data-nosnippet>74</a> <span class="kw">if </span>!backend
<a href=#75 id=75 data-nosnippet>75</a> .is_file(PACK_INDEX_FILE)
<a href=#76 id=76 data-nosnippet>76</a> .<span class="kw">await
<a href=#77 id=77 data-nosnippet>77</a> </span>.context(error::ReadObjectSnafu {
<a href=#78 id=78 data-nosnippet>78</a> path: PACK_INDEX_FILE,
<a href=#79 id=79 data-nosnippet>79</a> })<span class="question-mark">?
<a href=#80 id=80 data-nosnippet>80</a> </span>{
<a href=#81 id=81 data-nosnippet>81</a> <span class="kw">return </span>error::InvalidCopyParameterSnafu {
<a href=#82 id=82 data-nosnippet>82</a> key: <span class="string">"index"</span>,
<a href=#83 id=83 data-nosnippet>83</a> value: <span class="string">"not a regular file"</span>,
<a href=#84 id=84 data-nosnippet>84</a> }
<a href=#85 id=85 data-nosnippet>85</a> .fail();
<a href=#86 id=86 data-nosnippet>86</a> }
<a href=#87 id=87 data-nosnippet>87</a> <span class="kw">let </span>bytes = backend
<a href=#88 id=88 data-nosnippet>88</a> .object_store
<a href=#89 id=89 data-nosnippet>89</a> .read(PACK_INDEX_FILE)
<a href=#90 id=90 data-nosnippet>90</a> .<span class="kw">await
<a href=#91 id=91 data-nosnippet>91</a> </span>.context(error::ReadObjectSnafu {
<a href=#92 id=92 data-nosnippet>92</a> path: PACK_INDEX_FILE,
<a href=#93 id=93 data-nosnippet>93</a> })<span class="question-mark">?
<a href=#94 id=94 data-nosnippet>94</a> </span>.to_bytes();
<a href=#95 id=95 data-nosnippet>95</a> <span class="kw">let </span><span class="kw-2">mut </span>index: PackIndex = serde_json::from_slice(<span class="kw-2">&amp;</span>bytes).map_err(|e| {
<a href=#96 id=96 data-nosnippet>96</a> error::InvalidCopyParameterSnafu {
<a href=#97 id=97 data-nosnippet>97</a> key: <span class="string">"pack_index"</span>,
<a href=#98 id=98 data-nosnippet>98</a> value: e.to_string(),
<a href=#99 id=99 data-nosnippet>99</a> }
<a href=#100 id=100 data-nosnippet>100</a> .build()
<a href=#101 id=101 data-nosnippet>101</a> })<span class="question-mark">?</span>;
<a href=#102 id=102 data-nosnippet>102</a> index
<a href=#103 id=103 data-nosnippet>103</a> .validate()
<a href=#104 id=104 data-nosnippet>104</a> .map_err(BoxedError::new)
<a href=#105 id=105 data-nosnippet>105</a> .context(error::ExternalSnafu)<span class="question-mark">?</span>;
<a href=#106 id=106 data-nosnippet>106</a> index
<a href=#107 id=107 data-nosnippet>107</a> .tables
<a href=#108 id=108 data-nosnippet>108</a> .sort_by(|a, b| (<span class="kw-2">&amp;</span>a.object, a.offset).cmp(<span class="kw-2">&amp;</span>(<span class="kw-2">&amp;</span>b.object, b.offset)));
<a href=#109 id=109 data-nosnippet>109</a> <span class="kw">let </span>names: Vec&lt;<span class="kw">_</span>&gt; = index.tables.iter().map(|t| t.table_name.clone()).collect();
<a href=#110 id=110 data-nosnippet>110</a> <span class="kw">let </span><span class="kw-2">mut </span>tables = <span class="self">self
<a href=#111 id=111 data-nosnippet>111</a> </span>.capture_database_export_tables(<span class="kw-2">&amp;</span>req, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>names), ctx)
<a href=#112 id=112 data-nosnippet>112</a> .<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#113 id=113 data-nosnippet>113</a> <span class="kw">if </span>tables.len() != names.len() {
<a href=#114 id=114 data-nosnippet>114</a> <span class="kw">return </span>error::InvalidCopyParameterSnafu {
<a href=#115 id=115 data-nosnippet>115</a> key: <span class="string">"index"</span>,
<a href=#116 id=116 data-nosnippet>116</a> value: <span class="string">"physical tables and views are not data targets"</span>,
<a href=#117 id=117 data-nosnippet>117</a> }
<a href=#118 id=118 data-nosnippet>118</a> .fail();
<a href=#119 id=119 data-nosnippet>119</a> }
<a href=#120 id=120 data-nosnippet>120</a> <span class="kw">let </span>positions: HashMap&lt;<span class="kw">_</span>, <span class="kw">_</span>&gt; = names
<a href=#121 id=121 data-nosnippet>121</a> .iter()
<a href=#122 id=122 data-nosnippet>122</a> .enumerate()
<a href=#123 id=123 data-nosnippet>123</a> .map(|(i, n)| (n.as_str(), i))
<a href=#124 id=124 data-nosnippet>124</a> .collect();
<a href=#125 id=125 data-nosnippet>125</a> tables.sort_by_key(|table| positions.get(table.table_info().name.as_str()).copied());
<a href=#126 id=126 data-nosnippet>126</a> <span class="kw">for </span>object <span class="kw">in </span><span class="kw-2">&amp;</span>index.objects {
<a href=#127 id=127 data-nosnippet>127</a> <span class="kw">let </span>meta = backend
<a href=#128 id=128 data-nosnippet>128</a> .object_store
<a href=#129 id=129 data-nosnippet>129</a> .stat(<span class="kw-2">&amp;</span>object.path)
<a href=#130 id=130 data-nosnippet>130</a> .<span class="kw">await
<a href=#131 id=131 data-nosnippet>131</a> </span>.context(error::ReadObjectSnafu { path: <span class="kw-2">&amp;</span>object.path })<span class="question-mark">?</span>;
<a href=#132 id=132 data-nosnippet>132</a> <span class="kw">if </span>!backend
<a href=#133 id=133 data-nosnippet>133</a> .is_file_with_mode(<span class="kw-2">&amp;</span>object.path, meta.mode())
<a href=#134 id=134 data-nosnippet>134</a> .<span class="kw">await
<a href=#135 id=135 data-nosnippet>135</a> </span>.context(error::ReadObjectSnafu { path: <span class="kw-2">&amp;</span>object.path })<span class="question-mark">?
<a href=#136 id=136 data-nosnippet>136</a> </span>{
<a href=#137 id=137 data-nosnippet>137</a> <span class="kw">return </span>error::InvalidCopyParameterSnafu {
<a href=#138 id=138 data-nosnippet>138</a> key: <span class="string">"object"</span>,
<a href=#139 id=139 data-nosnippet>139</a> value: <span class="kw-2">&amp;</span>object.path,
<a href=#140 id=140 data-nosnippet>140</a> }
<a href=#141 id=141 data-nosnippet>141</a> .fail();
<a href=#142 id=142 data-nosnippet>142</a> }
<a href=#143 id=143 data-nosnippet>143</a> <span class="kw">if </span>meta.content_length() != object.length {
<a href=#144 id=144 data-nosnippet>144</a> <span class="kw">return </span>error::InvalidCopyParameterSnafu {
<a href=#145 id=145 data-nosnippet>145</a> key: <span class="string">"object_length"</span>,
<a href=#146 id=146 data-nosnippet>146</a> value: <span class="kw-2">&amp;</span>object.path,
<a href=#147 id=147 data-nosnippet>147</a> }
<a href=#148 id=148 data-nosnippet>148</a> .fail();
<a href=#149 id=149 data-nosnippet>149</a> }
<a href=#150 id=150 data-nosnippet>150</a> }
<a href=#151 id=151 data-nosnippet>151</a> <span class="prelude-val">Ok</span>(PreparedPackedImport {
<a href=#152 id=152 data-nosnippet>152</a> backend,
<a href=#153 id=153 data-nosnippet>153</a> index,
<a href=#154 id=154 data-nosnippet>154</a> tables,
<a href=#155 id=155 data-nosnippet>155</a> parallelism: <span class="kw">crate</span>::statement::database_copy::parse_parallelism_from_option_map(
<a href=#156 id=156 data-nosnippet>156</a> <span class="kw-2">&amp;</span>req.with,
<a href=#157 id=157 data-nosnippet>157</a> ),
<a href=#158 id=158 data-nosnippet>158</a> })
<a href=#159 id=159 data-nosnippet>159</a> }
<a href=#160 id=160 data-nosnippet>160</a>
<a href=#161 id=161 data-nosnippet>161</a> <span class="kw">pub async fn </span>import_packed(
<a href=#162 id=162 data-nosnippet>162</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#163 id=163 data-nosnippet>163</a> plan: PreparedPackedImport,
<a href=#164 id=164 data-nosnippet>164</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#165 id=165 data-nosnippet>165</a> ctx: QueryContextRef,
<a href=#166 id=166 data-nosnippet>166</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#167 id=167 data-nosnippet>167</a> <span class="kw">let </span>windows = PackReadWindows::new(plan.backend.object_store.clone());
<a href=#168 id=168 data-nosnippet>168</a> <span class="kw">let </span><span class="kw-2">mut </span>pending = PendingPackedInserts::new(plan.parallelism);
<a href=#169 id=169 data-nosnippet>169</a> <span class="kw">let </span>objects: HashMap&lt;<span class="kw">_</span>, <span class="kw">_</span>&gt; = plan
<a href=#170 id=170 data-nosnippet>170</a> .index
<a href=#171 id=171 data-nosnippet>171</a> .objects
<a href=#172 id=172 data-nosnippet>172</a> .iter()
<a href=#173 id=173 data-nosnippet>173</a> .map(|o| (o.path.as_str(), o))
<a href=#174 id=174 data-nosnippet>174</a> .collect();
<a href=#175 id=175 data-nosnippet>175</a> <span class="kw">let </span>result: <span class="prelude-ty">Result</span>&lt;()&gt; = <span class="kw">async </span>{
<a href=#176 id=176 data-nosnippet>176</a> <span class="kw">for </span>(entry, table) <span class="kw">in </span>plan.index.tables.iter().zip(plan.tables) {
<a href=#177 id=177 data-nosnippet>177</a> <span class="kw">if </span>cancellation.is_cancelled() {
<a href=#178 id=178 data-nosnippet>178</a> <span class="kw">return </span>error::PackedImportCancelledSnafu.fail();
<a href=#179 id=179 data-nosnippet>179</a> }
<a href=#180 id=180 data-nosnippet>180</a> <span class="kw">let </span>object = objects.get(entry.object.as_str()).ok_or_else(|| {
<a href=#181 id=181 data-nosnippet>181</a> error::InvalidCopyParameterSnafu {
<a href=#182 id=182 data-nosnippet>182</a> key: <span class="string">"object"</span>,
<a href=#183 id=183 data-nosnippet>183</a> value: <span class="kw-2">&amp;</span>entry.object,
<a href=#184 id=184 data-nosnippet>184</a> }
<a href=#185 id=185 data-nosnippet>185</a> .build()
<a href=#186 id=186 data-nosnippet>186</a> })<span class="question-mark">?</span>;
<a href=#187 id=187 data-nosnippet>187</a> <span class="kw">match </span>object.kind {
<a href=#188 id=188 data-nosnippet>188</a> ObjectKind::Pack <span class="kw">if </span>entry.length &lt;= WINDOW_SIZE <span class="kw">as </span>u64 =&gt; {
<a href=#189 id=189 data-nosnippet>189</a> <span class="kw">let </span>reader = PackedParquetReader::new(
<a href=#190 id=190 data-nosnippet>190</a> Arc::clone(<span class="kw-2">&amp;</span>windows),
<a href=#191 id=191 data-nosnippet>191</a> object.path.clone(),
<a href=#192 id=192 data-nosnippet>192</a> object.length,
<a href=#193 id=193 data-nosnippet>193</a> entry.offset,
<a href=#194 id=194 data-nosnippet>194</a> entry.length,
<a href=#195 id=195 data-nosnippet>195</a> )
<a href=#196 id=196 data-nosnippet>196</a> .context(error::ReadParquetMetadataSnafu)<span class="question-mark">?</span>;
<a href=#197 id=197 data-nosnippet>197</a> <span class="self">self</span>.copy_indexed_parquet(
<a href=#198 id=198 data-nosnippet>198</a> reader,
<a href=#199 id=199 data-nosnippet>199</a> table,
<a href=#200 id=200 data-nosnippet>200</a> entry.row_count,
<a href=#201 id=201 data-nosnippet>201</a> <span class="kw-2">&amp;mut </span>pending,
<a href=#202 id=202 data-nosnippet>202</a> cancellation,
<a href=#203 id=203 data-nosnippet>203</a> ctx.clone(),
<a href=#204 id=204 data-nosnippet>204</a> )
<a href=#205 id=205 data-nosnippet>205</a> .<span class="kw">await</span><span class="question-mark">?
<a href=#206 id=206 data-nosnippet>206</a> </span>}
<a href=#207 id=207 data-nosnippet>207</a> ObjectKind::Pack | ObjectKind::Parquet =&gt; {
<a href=#208 id=208 data-nosnippet>208</a> <span class="comment">// The virtual file keeps footer/column offsets relative to this
<a href=#209 id=209 data-nosnippet>209</a> // stream. Parquet decoder allocations are separate from the cache.
<a href=#210 id=210 data-nosnippet>210</a> </span><span class="kw">let </span>reader = plan
<a href=#211 id=211 data-nosnippet>211</a> .backend
<a href=#212 id=212 data-nosnippet>212</a> .object_store
<a href=#213 id=213 data-nosnippet>213</a> .reader_with(<span class="kw-2">&amp;</span>object.path)
<a href=#214 id=214 data-nosnippet>214</a> .chunk(<span class="number">256 </span>* <span class="number">1024</span>)
<a href=#215 id=215 data-nosnippet>215</a> .<span class="kw">await
<a href=#216 id=216 data-nosnippet>216</a> </span>.context(error::ReadObjectSnafu { path: <span class="kw-2">&amp;</span>object.path })<span class="question-mark">?
<a href=#217 id=217 data-nosnippet>217</a> </span>.into_futures_async_read(entry.offset..entry.offset + entry.length)
<a href=#218 id=218 data-nosnippet>218</a> .<span class="kw">await
<a href=#219 id=219 data-nosnippet>219</a> </span>.context(error::ReadObjectSnafu { path: <span class="kw-2">&amp;</span>object.path })<span class="question-mark">?
<a href=#220 id=220 data-nosnippet>220</a> </span>.compat();
<a href=#221 id=221 data-nosnippet>221</a> <span class="self">self</span>.copy_indexed_parquet(
<a href=#222 id=222 data-nosnippet>222</a> reader,
<a href=#223 id=223 data-nosnippet>223</a> table,
<a href=#224 id=224 data-nosnippet>224</a> entry.row_count,
<a href=#225 id=225 data-nosnippet>225</a> <span class="kw-2">&amp;mut </span>pending,
<a href=#226 id=226 data-nosnippet>226</a> cancellation,
<a href=#227 id=227 data-nosnippet>227</a> ctx.clone(),
<a href=#228 id=228 data-nosnippet>228</a> )
<a href=#229 id=229 data-nosnippet>229</a> .<span class="kw">await</span><span class="question-mark">?
<a href=#230 id=230 data-nosnippet>230</a> </span>}
<a href=#231 id=231 data-nosnippet>231</a> };
<a href=#232 id=232 data-nosnippet>232</a> }
<a href=#233 id=233 data-nosnippet>233</a> <span class="prelude-val">Ok</span>(())
<a href=#234 id=234 data-nosnippet>234</a> }
<a href=#235 id=235 data-nosnippet>235</a> .<span class="kw">await</span>;
<a href=#236 id=236 data-nosnippet>236</a> <span class="kw">let </span>drained = pending.drain().<span class="kw">await</span>;
<a href=#237 id=237 data-nosnippet>237</a> <span class="macro">common_telemetry::debug!</span>(
<a href=#238 id=238 data-nosnippet>238</a> peak_pending_bytes = pending.peak_pending_bytes,
<a href=#239 id=239 data-nosnippet>239</a> peak_batch_bytes = pending.peak_batch_bytes,
<a href=#240 id=240 data-nosnippet>240</a> <span class="string">"Packed import insertion payloads; codec and compressed windows accounted separately"
<a href=#241 id=241 data-nosnippet>241</a> </span>);
<a href=#242 id=242 data-nosnippet>242</a> result<span class="question-mark">?</span>;
<a href=#243 id=243 data-nosnippet>243</a> drained<span class="question-mark">?</span>;
<a href=#244 id=244 data-nosnippet>244</a> <span class="prelude-val">Ok</span>(Output::new(
<a href=#245 id=245 data-nosnippet>245</a> common_query::OutputData::AffectedRows(pending.rows),
<a href=#246 id=246 data-nosnippet>246</a> common_query::OutputMeta::new_with_cost(pending.cost),
<a href=#247 id=247 data-nosnippet>247</a> ))
<a href=#248 id=248 data-nosnippet>248</a> }
<a href=#249 id=249 data-nosnippet>249</a>}
<a href=#250 id=250 data-nosnippet>250</a>
<a href=#251 id=251 data-nosnippet>251</a><span class="doccomment">/// Retains all started writes until completion, including after the first error.
<a href=#252 id=252 data-nosnippet>252</a></span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">struct </span>PendingPackedInserts {
<a href=#253 id=253 data-nosnippet>253</a> tasks: FuturesUnordered&lt;tokio::task::JoinHandle&lt;(usize, <span class="prelude-ty">Result</span>&lt;Output&gt;)&gt;&gt;,
<a href=#254 id=254 data-nosnippet>254</a> bytes: usize,
<a href=#255 id=255 data-nosnippet>255</a> parallelism: usize,
<a href=#256 id=256 data-nosnippet>256</a> rows: usize,
<a href=#257 id=257 data-nosnippet>257</a> cost: usize,
<a href=#258 id=258 data-nosnippet>258</a> peak_pending_bytes: usize,
<a href=#259 id=259 data-nosnippet>259</a> peak_batch_bytes: usize,
<a href=#260 id=260 data-nosnippet>260</a>}
<a href=#261 id=261 data-nosnippet>261</a>
<a href=#262 id=262 data-nosnippet>262</a><span class="kw">impl </span>PendingPackedInserts {
<a href=#263 id=263 data-nosnippet>263</a> <span class="kw">fn </span>new(parallelism: usize) -&gt; <span class="self">Self </span>{
<a href=#264 id=264 data-nosnippet>264</a> <span class="self">Self </span>{
<a href=#265 id=265 data-nosnippet>265</a> tasks: FuturesUnordered::new(),
<a href=#266 id=266 data-nosnippet>266</a> bytes: <span class="number">0</span>,
<a href=#267 id=267 data-nosnippet>267</a> parallelism,
<a href=#268 id=268 data-nosnippet>268</a> rows: <span class="number">0</span>,
<a href=#269 id=269 data-nosnippet>269</a> cost: <span class="number">0</span>,
<a href=#270 id=270 data-nosnippet>270</a> peak_pending_bytes: <span class="number">0</span>,
<a href=#271 id=271 data-nosnippet>271</a> peak_batch_bytes: <span class="number">0</span>,
<a href=#272 id=272 data-nosnippet>272</a> }
<a href=#273 id=273 data-nosnippet>273</a> }
<a href=#274 id=274 data-nosnippet>274</a>
<a href=#275 id=275 data-nosnippet>275</a> <span class="doccomment">/// Backpressure before decoding, while allowing one current batch outside the queue.
<a href=#276 id=276 data-nosnippet>276</a> </span><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>before_decode(<span class="kw-2">&amp;mut </span><span class="self">self</span>) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#277 id=277 data-nosnippet>277</a> <span class="kw">while let </span><span class="prelude-val">Some</span>(<span class="prelude-val">Some</span>(output)) = <span class="self">self</span>.tasks.next().now_or_never() {
<a href=#278 id=278 data-nosnippet>278</a> <span class="self">self</span>.complete(output)<span class="question-mark">?</span>;
<a href=#279 id=279 data-nosnippet>279</a> }
<a href=#280 id=280 data-nosnippet>280</a> <span class="kw">while </span><span class="self">self</span>.tasks.len() &gt;= <span class="self">self</span>.parallelism || <span class="self">self</span>.bytes &gt;= <span class="number">32 </span>* <span class="number">1024 </span>* <span class="number">1024 </span>{
<a href=#281 id=281 data-nosnippet>281</a> <span class="self">self</span>.complete_one().<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#282 id=282 data-nosnippet>282</a> }
<a href=#283 id=283 data-nosnippet>283</a> <span class="prelude-val">Ok</span>(())
<a href=#284 id=284 data-nosnippet>284</a> }
<a href=#285 id=285 data-nosnippet>285</a>
<a href=#286 id=286 data-nosnippet>286</a> <span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">async fn </span>admit(
<a href=#287 id=287 data-nosnippet>287</a> <span class="kw-2">&amp;mut </span><span class="self">self</span>,
<a href=#288 id=288 data-nosnippet>288</a> bytes: usize,
<a href=#289 id=289 data-nosnippet>289</a> insert: <span class="kw">impl </span>std::future::Future&lt;Output = <span class="prelude-ty">Result</span>&lt;Output&gt;&gt; + Send + <span class="lifetime">'static</span>,
<a href=#290 id=290 data-nosnippet>290</a> cancellation: <span class="kw-2">&amp;</span>CancellationToken,
<a href=#291 id=291 data-nosnippet>291</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#292 id=292 data-nosnippet>292</a> <span class="kw">while let </span><span class="prelude-val">Some</span>(<span class="prelude-val">Some</span>(output)) = <span class="self">self</span>.tasks.next().now_or_never() {
<a href=#293 id=293 data-nosnippet>293</a> <span class="self">self</span>.complete(output)<span class="question-mark">?</span>;
<a href=#294 id=294 data-nosnippet>294</a> }
<a href=#295 id=295 data-nosnippet>295</a> <span class="self">self</span>.peak_batch_bytes = <span class="self">self</span>.peak_batch_bytes.max(bytes);
<a href=#296 id=296 data-nosnippet>296</a> <span class="kw">const </span>BUDGET: usize = <span class="number">32 </span>* <span class="number">1024 </span>* <span class="number">1024</span>;
<a href=#297 id=297 data-nosnippet>297</a> <span class="kw">if </span>bytes &gt; BUDGET {
<a href=#298 id=298 data-nosnippet>298</a> <span class="self">self</span>.drain().<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#299 id=299 data-nosnippet>299</a> <span class="kw">if </span>cancellation.is_cancelled() {
<a href=#300 id=300 data-nosnippet>300</a> <span class="kw">return </span>error::PackedImportCancelledSnafu.fail();
<a href=#301 id=301 data-nosnippet>301</a> }
<a href=#302 id=302 data-nosnippet>302</a> <span class="kw">let </span>(rows, cost) = insert.<span class="kw">await</span><span class="question-mark">?</span>.extract_rows_and_cost();
<a href=#303 id=303 data-nosnippet>303</a> <span class="self">self</span>.rows += rows;
<a href=#304 id=304 data-nosnippet>304</a> <span class="self">self</span>.cost += cost;
<a href=#305 id=305 data-nosnippet>305</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(());
<a href=#306 id=306 data-nosnippet>306</a> }
<a href=#307 id=307 data-nosnippet>307</a> <span class="kw">while </span><span class="self">self</span>.tasks.len() &gt;= <span class="self">self</span>.parallelism || <span class="self">self</span>.bytes + bytes &gt; BUDGET {
<a href=#308 id=308 data-nosnippet>308</a> <span class="self">self</span>.complete_one().<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#309 id=309 data-nosnippet>309</a> }
<a href=#310 id=310 data-nosnippet>310</a> <span class="kw">if </span>cancellation.is_cancelled() {
<a href=#311 id=311 data-nosnippet>311</a> <span class="kw">return </span>error::PackedImportCancelledSnafu.fail();
<a href=#312 id=312 data-nosnippet>312</a> }
<a href=#313 id=313 data-nosnippet>313</a> <span class="self">self</span>.bytes += bytes;
<a href=#314 id=314 data-nosnippet>314</a> <span class="self">self</span>.peak_pending_bytes = <span class="self">self</span>.peak_pending_bytes.max(<span class="self">self</span>.bytes);
<a href=#315 id=315 data-nosnippet>315</a> <span class="self">self</span>.tasks.push(common_runtime::spawn_query(
<a href=#316 id=316 data-nosnippet>316</a> <span class="kw">async move </span>{ (bytes, insert.<span class="kw">await</span>) },
<a href=#317 id=317 data-nosnippet>317</a> ));
<a href=#318 id=318 data-nosnippet>318</a> <span class="prelude-val">Ok</span>(())
<a href=#319 id=319 data-nosnippet>319</a> }
<a href=#320 id=320 data-nosnippet>320</a>
<a href=#321 id=321 data-nosnippet>321</a> <span class="kw">async fn </span>complete_one(<span class="kw-2">&amp;mut </span><span class="self">self</span>) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#322 id=322 data-nosnippet>322</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(output) = <span class="self">self</span>.tasks.next().<span class="kw">await </span>{
<a href=#323 id=323 data-nosnippet>323</a> <span class="self">self</span>.complete(output)<span class="question-mark">?</span>;
<a href=#324 id=324 data-nosnippet>324</a> }
<a href=#325 id=325 data-nosnippet>325</a> <span class="prelude-val">Ok</span>(())
<a href=#326 id=326 data-nosnippet>326</a> }
<a href=#327 id=327 data-nosnippet>327</a>
<a href=#328 id=328 data-nosnippet>328</a> <span class="kw">fn </span>complete(
<a href=#329 id=329 data-nosnippet>329</a> <span class="kw-2">&amp;mut </span><span class="self">self</span>,
<a href=#330 id=330 data-nosnippet>330</a> output: std::result::Result&lt;(usize, <span class="prelude-ty">Result</span>&lt;Output&gt;), common_runtime::JoinError&gt;,
<a href=#331 id=331 data-nosnippet>331</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#332 id=332 data-nosnippet>332</a> <span class="kw">let </span>(bytes, output) = output.context(error::JoinTaskSnafu)<span class="question-mark">?</span>;
<a href=#333 id=333 data-nosnippet>333</a> <span class="self">self</span>.bytes -= bytes;
<a href=#334 id=334 data-nosnippet>334</a> <span class="kw">let </span>(rows, cost) = output<span class="question-mark">?</span>.extract_rows_and_cost();
<a href=#335 id=335 data-nosnippet>335</a> <span class="self">self</span>.rows += rows;
<a href=#336 id=336 data-nosnippet>336</a> <span class="self">self</span>.cost += cost;
<a href=#337 id=337 data-nosnippet>337</a> <span class="prelude-val">Ok</span>(())
<a href=#338 id=338 data-nosnippet>338</a> }
<a href=#339 id=339 data-nosnippet>339</a>
<a href=#340 id=340 data-nosnippet>340</a> <span class="kw">async fn </span>drain(<span class="kw-2">&amp;mut </span><span class="self">self</span>) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#341 id=341 data-nosnippet>341</a> <span class="kw">let </span><span class="kw-2">mut </span>first = <span class="prelude-val">None</span>;
<a href=#342 id=342 data-nosnippet>342</a> <span class="kw">while </span>!<span class="self">self</span>.tasks.is_empty() {
<a href=#343 id=343 data-nosnippet>343</a> <span class="kw">if let </span><span class="prelude-val">Err</span>(error) = <span class="self">self</span>.complete_one().<span class="kw">await </span>{
<a href=#344 id=344 data-nosnippet>344</a> first.get_or_insert(error);
<a href=#345 id=345 data-nosnippet>345</a> }
<a href=#346 id=346 data-nosnippet>346</a> }
<a href=#347 id=347 data-nosnippet>347</a> first.map_or(<span class="prelude-val">Ok</span>(()), <span class="prelude-val">Err</span>)
<a href=#348 id=348 data-nosnippet>348</a> }
<a href=#349 id=349 data-nosnippet>349</a>}
<a href=#350 id=350 data-nosnippet>350</a>
<a href=#351 id=351 data-nosnippet>351</a><span class="attr">#[cfg(test)]
<a href=#352 id=352 data-nosnippet>352</a></span><span class="kw">mod </span>tests {
<a href=#353 id=353 data-nosnippet>353</a> <span class="kw">use </span>std::sync::atomic::{AtomicUsize, Ordering};
<a href=#354 id=354 data-nosnippet>354</a>
<a href=#355 id=355 data-nosnippet>355</a> <span class="kw">use super</span>::<span class="kw-2">*</span>;
<a href=#356 id=356 data-nosnippet>356</a>
<a href=#357 id=357 data-nosnippet>357</a> <span class="attr">#[tokio::test]
<a href=#358 id=358 data-nosnippet>358</a> </span><span class="kw">async fn </span>oversized_batch_drains_and_runs_exclusively() {
<a href=#359 id=359 data-nosnippet>359</a> <span class="kw">let </span>cancellation = CancellationToken::new();
<a href=#360 id=360 data-nosnippet>360</a> <span class="kw">let </span><span class="kw-2">mut </span>queue = PendingPackedInserts::new(<span class="number">4</span>);
<a href=#361 id=361 data-nosnippet>361</a> <span class="kw">let </span>completed = Arc::new(AtomicUsize::new(<span class="number">0</span>));
<a href=#362 id=362 data-nosnippet>362</a> <span class="kw">for _ in </span><span class="number">0</span>..<span class="number">2 </span>{
<a href=#363 id=363 data-nosnippet>363</a> <span class="kw">let </span>completed = completed.clone();
<a href=#364 id=364 data-nosnippet>364</a> queue
<a href=#365 id=365 data-nosnippet>365</a> .admit(
<a href=#366 id=366 data-nosnippet>366</a> <span class="number">1024</span>,
<a href=#367 id=367 data-nosnippet>367</a> <span class="kw">async move </span>{
<a href=#368 id=368 data-nosnippet>368</a> completed.fetch_add(<span class="number">1</span>, Ordering::SeqCst);
<a href=#369 id=369 data-nosnippet>369</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="number">1</span>))
<a href=#370 id=370 data-nosnippet>370</a> },
<a href=#371 id=371 data-nosnippet>371</a> <span class="kw-2">&amp;</span>cancellation,
<a href=#372 id=372 data-nosnippet>372</a> )
<a href=#373 id=373 data-nosnippet>373</a> .<span class="kw">await
<a href=#374 id=374 data-nosnippet>374</a> </span>.unwrap();
<a href=#375 id=375 data-nosnippet>375</a> }
<a href=#376 id=376 data-nosnippet>376</a> <span class="kw">let </span>observed = completed.clone();
<a href=#377 id=377 data-nosnippet>377</a> queue
<a href=#378 id=378 data-nosnippet>378</a> .admit(
<a href=#379 id=379 data-nosnippet>379</a> <span class="number">40 </span>* <span class="number">1024 </span>* <span class="number">1024</span>,
<a href=#380 id=380 data-nosnippet>380</a> <span class="kw">async move </span>{
<a href=#381 id=381 data-nosnippet>381</a> <span class="macro">assert_eq!</span>(observed.load(Ordering::SeqCst), <span class="number">2</span>);
<a href=#382 id=382 data-nosnippet>382</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="number">3</span>))
<a href=#383 id=383 data-nosnippet>383</a> },
<a href=#384 id=384 data-nosnippet>384</a> <span class="kw-2">&amp;</span>cancellation,
<a href=#385 id=385 data-nosnippet>385</a> )
<a href=#386 id=386 data-nosnippet>386</a> .<span class="kw">await
<a href=#387 id=387 data-nosnippet>387</a> </span>.unwrap();
<a href=#388 id=388 data-nosnippet>388</a> <span class="macro">assert_eq!</span>(queue.rows, <span class="number">5</span>);
<a href=#389 id=389 data-nosnippet>389</a> <span class="macro">assert!</span>(queue.tasks.is_empty());
<a href=#390 id=390 data-nosnippet>390</a> <span class="macro">assert!</span>(queue.peak_pending_bytes &lt;= <span class="number">32 </span>* <span class="number">1024 </span>* <span class="number">1024</span>);
<a href=#391 id=391 data-nosnippet>391</a> <span class="macro">assert_eq!</span>(queue.peak_batch_bytes, <span class="number">40 </span>* <span class="number">1024 </span>* <span class="number">1024</span>);
<a href=#392 id=392 data-nosnippet>392</a> }
<a href=#393 id=393 data-nosnippet>393</a>
<a href=#394 id=394 data-nosnippet>394</a> <span class="attr">#[tokio::test]
<a href=#395 id=395 data-nosnippet>395</a> </span><span class="kw">async fn </span>concurrent_insert_queue_drains_after_error_and_cancel() {
<a href=#396 id=396 data-nosnippet>396</a> <span class="kw">let </span>cancellation = CancellationToken::new();
<a href=#397 id=397 data-nosnippet>397</a> <span class="kw">let </span><span class="kw-2">mut </span>queue = PendingPackedInserts::new(<span class="number">2</span>);
<a href=#398 id=398 data-nosnippet>398</a> <span class="kw">let </span>barrier = Arc::new(tokio::sync::Barrier::new(<span class="number">2</span>));
<a href=#399 id=399 data-nosnippet>399</a> <span class="kw">let </span>completed = Arc::new(AtomicUsize::new(<span class="number">0</span>));
<a href=#400 id=400 data-nosnippet>400</a> <span class="kw">for </span>fail <span class="kw">in </span>[<span class="bool-val">true</span>, <span class="bool-val">false</span>] {
<a href=#401 id=401 data-nosnippet>401</a> <span class="kw">let </span>barrier = barrier.clone();
<a href=#402 id=402 data-nosnippet>402</a> <span class="kw">let </span>completed = completed.clone();
<a href=#403 id=403 data-nosnippet>403</a> queue
<a href=#404 id=404 data-nosnippet>404</a> .admit(
<a href=#405 id=405 data-nosnippet>405</a> <span class="number">16 </span>* <span class="number">1024 </span>* <span class="number">1024</span>,
<a href=#406 id=406 data-nosnippet>406</a> <span class="kw">async move </span>{
<a href=#407 id=407 data-nosnippet>407</a> barrier.wait().<span class="kw">await</span>;
<a href=#408 id=408 data-nosnippet>408</a> completed.fetch_add(<span class="number">1</span>, Ordering::SeqCst);
<a href=#409 id=409 data-nosnippet>409</a> <span class="kw">if </span>fail {
<a href=#410 id=410 data-nosnippet>410</a> error::InvalidCopyParameterSnafu {
<a href=#411 id=411 data-nosnippet>411</a> key: <span class="string">"test"</span>,
<a href=#412 id=412 data-nosnippet>412</a> value: <span class="string">"failure"</span>,
<a href=#413 id=413 data-nosnippet>413</a> }
<a href=#414 id=414 data-nosnippet>414</a> .fail()
<a href=#415 id=415 data-nosnippet>415</a> } <span class="kw">else </span>{
<a href=#416 id=416 data-nosnippet>416</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="number">1</span>))
<a href=#417 id=417 data-nosnippet>417</a> }
<a href=#418 id=418 data-nosnippet>418</a> },
<a href=#419 id=419 data-nosnippet>419</a> <span class="kw-2">&amp;</span>cancellation,
<a href=#420 id=420 data-nosnippet>420</a> )
<a href=#421 id=421 data-nosnippet>421</a> .<span class="kw">await
<a href=#422 id=422 data-nosnippet>422</a> </span>.unwrap();
<a href=#423 id=423 data-nosnippet>423</a> }
<a href=#424 id=424 data-nosnippet>424</a> cancellation.cancel();
<a href=#425 id=425 data-nosnippet>425</a> <span class="macro">assert!</span>(
<a href=#426 id=426 data-nosnippet>426</a> queue
<a href=#427 id=427 data-nosnippet>427</a> .admit(<span class="number">1</span>, <span class="kw">async </span>{ <span class="macro">panic!</span>(<span class="string">"must not start"</span>) }, <span class="kw-2">&amp;</span>cancellation)
<a href=#428 id=428 data-nosnippet>428</a> .<span class="kw">await
<a href=#429 id=429 data-nosnippet>429</a> </span>.is_err()
<a href=#430 id=430 data-nosnippet>430</a> );
<a href=#431 id=431 data-nosnippet>431</a> <span class="kw">let _ </span>= queue.drain().<span class="kw">await</span>;
<a href=#432 id=432 data-nosnippet>432</a> <span class="macro">assert_eq!</span>(completed.load(Ordering::SeqCst), <span class="number">2</span>);
<a href=#433 id=433 data-nosnippet>433</a> <span class="macro">assert_eq!</span>(queue.bytes, <span class="number">0</span>);
<a href=#434 id=434 data-nosnippet>434</a> <span class="macro">assert!</span>(queue.tasks.is_empty());
<a href=#435 id=435 data-nosnippet>435</a> }
<a href=#436 id=436 data-nosnippet>436</a>}
</code></pre></div></section></main></body></html>
+78
View File
@@ -0,0 +1,78 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/kill.rs`."><title>kill.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>kill.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-2"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>catalog::process_manager::ProcessManagerRef;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>common_frontend::DisplayProcessId;
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_query::Output;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_telemetry::error;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>snafu::ResultExt;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>sql::statements::kill::Kill;
<a href=#22 id=22 data-nosnippet>22</a>
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span><span class="kw">crate</span>::error;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#25 id=25 data-nosnippet>25</a>
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">impl </span>StatementExecutor {
<a href=#27 id=27 data-nosnippet>27</a> <span class="kw">pub async fn </span>execute_kill(
<a href=#28 id=28 data-nosnippet>28</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#29 id=29 data-nosnippet>29</a> query_ctx: QueryContextRef,
<a href=#30 id=30 data-nosnippet>30</a> kill: Kill,
<a href=#31 id=31 data-nosnippet>31</a> ) -&gt; error::Result&lt;Output&gt; {
<a href=#32 id=32 data-nosnippet>32</a> <span class="kw">let </span><span class="prelude-val">Some</span>(process_manager) = <span class="self">self</span>.process_manager.as_ref() <span class="kw">else </span>{
<a href=#33 id=33 data-nosnippet>33</a> <span class="macro">error!</span>(<span class="string">"Process manager is not initialized"</span>);
<a href=#34 id=34 data-nosnippet>34</a> <span class="kw">return </span>error::ProcessManagerMissingSnafu.fail();
<a href=#35 id=35 data-nosnippet>35</a> };
<a href=#36 id=36 data-nosnippet>36</a>
<a href=#37 id=37 data-nosnippet>37</a> <span class="kw">let </span>succ = <span class="kw">match </span>kill {
<a href=#38 id=38 data-nosnippet>38</a> Kill::ProcessId(process_id) =&gt; {
<a href=#39 id=39 data-nosnippet>39</a> <span class="self">self</span>.kill_process_id(process_manager, query_ctx, process_id)
<a href=#40 id=40 data-nosnippet>40</a> .<span class="kw">await</span><span class="question-mark">?
<a href=#41 id=41 data-nosnippet>41</a> </span>}
<a href=#42 id=42 data-nosnippet>42</a> Kill::ConnectionId(conn_id) =&gt; {
<a href=#43 id=43 data-nosnippet>43</a> <span class="self">self</span>.kill_connection_id(process_manager, query_ctx, conn_id)
<a href=#44 id=44 data-nosnippet>44</a> .<span class="kw">await</span><span class="question-mark">?
<a href=#45 id=45 data-nosnippet>45</a> </span>}
<a href=#46 id=46 data-nosnippet>46</a> };
<a href=#47 id=47 data-nosnippet>47</a> <span class="prelude-val">Ok</span>(Output::new_with_affected_rows(<span class="kw">if </span>succ { <span class="number">1 </span>} <span class="kw">else </span>{ <span class="number">0 </span>}))
<a href=#48 id=48 data-nosnippet>48</a> }
<a href=#49 id=49 data-nosnippet>49</a>
<a href=#50 id=50 data-nosnippet>50</a> <span class="doccomment">/// Handles `KILL &lt;PROCESS_ID&gt;` statements.
<a href=#51 id=51 data-nosnippet>51</a> </span><span class="kw">async fn </span>kill_process_id(
<a href=#52 id=52 data-nosnippet>52</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#53 id=53 data-nosnippet>53</a> pm: <span class="kw-2">&amp;</span>ProcessManagerRef,
<a href=#54 id=54 data-nosnippet>54</a> query_ctx: QueryContextRef,
<a href=#55 id=55 data-nosnippet>55</a> process_id: String,
<a href=#56 id=56 data-nosnippet>56</a> ) -&gt; error::Result&lt;bool&gt; {
<a href=#57 id=57 data-nosnippet>57</a> <span class="kw">let </span>display_id = DisplayProcessId::try_from(process_id.as_str())
<a href=#58 id=58 data-nosnippet>58</a> .map_err(|<span class="kw">_</span>| error::InvalidProcessIdSnafu { id: process_id }.build())<span class="question-mark">?</span>;
<a href=#59 id=59 data-nosnippet>59</a>
<a href=#60 id=60 data-nosnippet>60</a> <span class="kw">let </span>current_user_catalog = query_ctx.current_catalog().to_string();
<a href=#61 id=61 data-nosnippet>61</a> pm.kill_process(display_id.server_addr, current_user_catalog, display_id.id)
<a href=#62 id=62 data-nosnippet>62</a> .<span class="kw">await
<a href=#63 id=63 data-nosnippet>63</a> </span>.context(error::CatalogSnafu)
<a href=#64 id=64 data-nosnippet>64</a> }
<a href=#65 id=65 data-nosnippet>65</a>
<a href=#66 id=66 data-nosnippet>66</a> <span class="doccomment">/// Handles MySQL `KILL QUERY &lt;CONNECTION_ID&gt;` statements.
<a href=#67 id=67 data-nosnippet>67</a> </span><span class="kw">pub async fn </span>kill_connection_id(
<a href=#68 id=68 data-nosnippet>68</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#69 id=69 data-nosnippet>69</a> pm: <span class="kw-2">&amp;</span>ProcessManagerRef,
<a href=#70 id=70 data-nosnippet>70</a> query_ctx: QueryContextRef,
<a href=#71 id=71 data-nosnippet>71</a> connection_id: u32,
<a href=#72 id=72 data-nosnippet>72</a> ) -&gt; error::Result&lt;bool&gt; {
<a href=#73 id=73 data-nosnippet>73</a> pm.kill_local_process(query_ctx.current_catalog().to_string(), connection_id)
<a href=#74 id=74 data-nosnippet>74</a> .<span class="kw">await
<a href=#75 id=75 data-nosnippet>75</a> </span>.context(error::CatalogSnafu)
<a href=#76 id=76 data-nosnippet>76</a> }
<a href=#77 id=77 data-nosnippet>77</a>}
</code></pre></div></section></main></body></html>
File diff suppressed because it is too large Load Diff
@@ -0,0 +1,379 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/semantic_graph/conventions.rs`."><title>conventions.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../../" data-static-root-path="../../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../../static.files/src-script-813739b1.js"></script><script defer src="../../../../src-files.js"></script><script defer src="../../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/semantic_graph/</div>conventions.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="doccomment">//! The built-in derivation conventions of the entity graph, embedded as
<a href=#16 id=16 data-nosnippet>16</a>//! [`conventions.yaml`](./conventions.yaml): the co-declared edge vocabulary,
<a href=#17 id=17 data-nosnippet>17</a>//! the virtual-destination candidates, and the implicit declarations of
<a href=#18 id=18 data-nosnippet>18</a>//! well-known Prometheus entity-descriptor metrics. The file is data shipped
<a href=#19 id=19 data-nosnippet>19</a>//! with the binary, not an operator-editable configuration surface; explicit
<a href=#20 id=20 data-nosnippet>20</a>//! `greptime.semantic.entity.*` declarations always override it.
<a href=#21 id=21 data-nosnippet>21</a>
<a href=#22 id=22 data-nosnippet>22</a></span><span class="kw">use </span>std::collections::{BTreeMap, HashMap, HashSet};
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>std::sync::LazyLock;
<a href=#24 id=24 data-nosnippet>24</a>
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>serde::Deserialize;
<a href=#26 id=26 data-nosnippet>26</a>
<a href=#27 id=27 data-nosnippet>27</a><span class="doccomment">/// A same-row co-declaration rule: a source-table row carrying both entity
<a href=#28 id=28 data-nosnippet>28</a>/// identities witnesses `src -rel-&gt; dst`.
<a href=#29 id=29 data-nosnippet>29</a></span><span class="attr">#[derive(Debug, Deserialize)]
<a href=#30 id=30 data-nosnippet>30</a>#[serde(deny_unknown_fields)]
<a href=#31 id=31 data-nosnippet>31</a></span><span class="kw">pub struct </span>EdgeRule {
<a href=#32 id=32 data-nosnippet>32</a> <span class="kw">pub </span>src: String,
<a href=#33 id=33 data-nosnippet>33</a> <span class="kw">pub </span>dst: String,
<a href=#34 id=34 data-nosnippet>34</a> <span class="kw">pub </span>rel: String,
<a href=#35 id=35 data-nosnippet>35</a>}
<a href=#36 id=36 data-nosnippet>36</a>
<a href=#37 id=37 data-nosnippet>37</a><span class="doccomment">/// A span-attribute column that may name an uninstrumented peer, and the
<a href=#38 id=38 data-nosnippet>38</a>/// `connection_type` its match implies.
<a href=#39 id=39 data-nosnippet>39</a></span><span class="attr">#[derive(Debug, Deserialize)]
<a href=#40 id=40 data-nosnippet>40</a>#[serde(deny_unknown_fields)]
<a href=#41 id=41 data-nosnippet>41</a></span><span class="kw">pub struct </span>VirtualDstCandidate {
<a href=#42 id=42 data-nosnippet>42</a> <span class="kw">pub </span>column: String,
<a href=#43 id=43 data-nosnippet>43</a> <span class="kw">pub </span>connection_type: String,
<a href=#44 id=44 data-nosnippet>44</a>}
<a href=#45 id=45 data-nosnippet>45</a>
<a href=#46 id=46 data-nosnippet>46</a><span class="doccomment">/// One implicit entity declaration: of a whitelisted Prometheus or OTel info
<a href=#47 id=47 data-nosnippet>47</a>/// metric, or of a Trace V1/V2 table's resource attributes.
<a href=#48 id=48 data-nosnippet>48</a></span><span class="attr">#[derive(Debug, Deserialize)]
<a href=#49 id=49 data-nosnippet>49</a>#[serde(deny_unknown_fields)]
<a href=#50 id=50 data-nosnippet>50</a></span><span class="kw">pub struct </span>ImplicitEntity {
<a href=#51 id=51 data-nosnippet>51</a> <span class="kw">pub </span>entity: String,
<a href=#52 id=52 data-nosnippet>52</a> <span class="doccomment">/// Identifying label columns, ordered broad to narrow; every one must
<a href=#53 id=53 data-nosnippet>53</a> /// exist on the table for the declaration to apply.
<a href=#54 id=54 data-nosnippet>54</a> </span><span class="kw">pub </span>id: Vec&lt;String&gt;,
<a href=#55 id=55 data-nosnippet>55</a> <span class="doccomment">/// Column qualifying `id[0]` as `&lt;qualifier&gt;/&lt;id[0]&gt;`, so a source
<a href=#56 id=56 data-nosnippet>56</a> /// carrying the parts separately matches one carrying them pre-composed.
<a href=#57 id=57 data-nosnippet>57</a> /// Skipped when the column is absent or empty.
<a href=#58 id=58 data-nosnippet>58</a> </span><span class="attr">#[serde(default)]
<a href=#59 id=59 data-nosnippet>59</a> </span><span class="kw">pub </span>qualified_by: <span class="prelude-ty">Option</span>&lt;String&gt;,
<a href=#60 id=60 data-nosnippet>60</a> <span class="doccomment">/// A more specific entity type declared next to this one, which takes over
<a href=#61 id=61 data-nosnippet>61</a> /// on the rows carrying its full identity. Naming the type rather than a
<a href=#62 id=62 data-nosnippet>62</a> /// trigger column is what keeps the entity replaceable but never
<a href=#63 id=63 data-nosnippet>63</a> /// droppable: where the specific type is not derivable, this one stands.
<a href=#64 id=64 data-nosnippet>64</a> </span><span class="attr">#[serde(default)]
<a href=#65 id=65 data-nosnippet>65</a> </span><span class="kw">pub </span>superseded_by: <span class="prelude-ty">Option</span>&lt;String&gt;,
<a href=#66 id=66 data-nosnippet>66</a> <span class="doccomment">/// Descriptive label columns, filtered to those present (kube-state-metrics
<a href=#67 id=67 data-nosnippet>67</a> /// label sets vary across versions).
<a href=#68 id=68 data-nosnippet>68</a> </span><span class="attr">#[serde(default)]
<a href=#69 id=69 data-nosnippet>69</a> </span><span class="kw">pub </span>descriptive: Vec&lt;String&gt;,
<a href=#70 id=70 data-nosnippet>70</a> <span class="doccomment">/// Snapshot every tag column except the id columns as descriptive —
<a href=#71 id=71 data-nosnippet>71</a> /// `target_info`-style enrichment over an open attribute set.
<a href=#72 id=72 data-nosnippet>72</a> </span><span class="attr">#[serde(default)]
<a href=#73 id=73 data-nosnippet>73</a> </span><span class="kw">pub </span>descriptive_rest: bool,
<a href=#74 id=74 data-nosnippet>74</a>}
<a href=#75 id=75 data-nosnippet>75</a>
<a href=#76 id=76 data-nosnippet>76</a><span class="doccomment">/// The parsed, validated conventions file.
<a href=#77 id=77 data-nosnippet>77</a></span><span class="attr">#[derive(Debug, Deserialize)]
<a href=#78 id=78 data-nosnippet>78</a>#[serde(deny_unknown_fields)]
<a href=#79 id=79 data-nosnippet>79</a></span><span class="kw">pub struct </span>Conventions {
<a href=#80 id=80 data-nosnippet>80</a> <span class="kw">pub </span>co_declared_edges: Vec&lt;EdgeRule&gt;,
<a href=#81 id=81 data-nosnippet>81</a> <span class="kw">pub </span>trace_co_declared_edges: Vec&lt;EdgeRule&gt;,
<a href=#82 id=82 data-nosnippet>82</a> <span class="kw">pub </span>virtual_dst_candidates: Vec&lt;VirtualDstCandidate&gt;,
<a href=#83 id=83 data-nosnippet>83</a> <span class="kw">pub </span>otlp_trace_entities: Vec&lt;ImplicitEntity&gt;,
<a href=#84 id=84 data-nosnippet>84</a> <span class="doccomment">/// Table name -&gt; the entities that table declares, for Prometheus-sourced
<a href=#85 id=85 data-nosnippet>85</a> /// descriptor metrics (`source = prometheus`).
<a href=#86 id=86 data-nosnippet>86</a> </span><span class="kw">pub </span>prometheus_info_metrics: BTreeMap&lt;String, Vec&lt;ImplicitEntity&gt;&gt;,
<a href=#87 id=87 data-nosnippet>87</a> <span class="doccomment">/// Table name -&gt; the entities that table declares, for OTLP-sourced
<a href=#88 id=88 data-nosnippet>88</a> /// descriptor tables (`source = opentelemetry`).
<a href=#89 id=89 data-nosnippet>89</a> </span><span class="kw">pub </span>otel_info_metrics: BTreeMap&lt;String, Vec&lt;ImplicitEntity&gt;&gt;,
<a href=#90 id=90 data-nosnippet>90</a>}
<a href=#91 id=91 data-nosnippet>91</a>
<a href=#92 id=92 data-nosnippet>92</a><span class="doccomment">/// The built-in entity-type vocabulary. User-declared types are open-ended;
<a href=#93 id=93 data-nosnippet>93</a>/// the embedded conventions file must stay inside this set.
<a href=#94 id=94 data-nosnippet>94</a></span><span class="kw">pub const </span>ENTITY_TYPE_SERVICE: <span class="kw-2">&amp;</span>str = <span class="string">"service"</span>;
<a href=#95 id=95 data-nosnippet>95</a><span class="kw">pub const </span>ENTITY_TYPE_SERVICE_INSTANCE: <span class="kw-2">&amp;</span>str = <span class="string">"service.instance"</span>;
<a href=#96 id=96 data-nosnippet>96</a><span class="kw">pub const </span>ENTITY_TYPE_HOST: <span class="kw-2">&amp;</span>str = <span class="string">"host"</span>;
<a href=#97 id=97 data-nosnippet>97</a><span class="kw">pub const </span>ENTITY_TYPE_CONTAINER: <span class="kw-2">&amp;</span>str = <span class="string">"container"</span>;
<a href=#98 id=98 data-nosnippet>98</a><span class="kw">pub const </span>ENTITY_TYPE_PROCESS: <span class="kw-2">&amp;</span>str = <span class="string">"process"</span>;
<a href=#99 id=99 data-nosnippet>99</a><span class="kw">pub const </span>ENTITY_TYPE_K8S_POD: <span class="kw-2">&amp;</span>str = <span class="string">"k8s.pod"</span>;
<a href=#100 id=100 data-nosnippet>100</a><span class="kw">pub const </span>ENTITY_TYPE_K8S_NODE: <span class="kw-2">&amp;</span>str = <span class="string">"k8s.node"</span>;
<a href=#101 id=101 data-nosnippet>101</a><span class="kw">pub const </span>ENTITY_TYPE_K8S_CONTAINER: <span class="kw-2">&amp;</span>str = <span class="string">"k8s.container"</span>;
<a href=#102 id=102 data-nosnippet>102</a><span class="kw">pub const </span>ENTITY_TYPE_K8S_WORKLOAD: <span class="kw-2">&amp;</span>str = <span class="string">"k8s.workload"</span>;
<a href=#103 id=103 data-nosnippet>103</a><span class="kw">pub const </span>ENTITY_TYPE_K8S_SERVICE: <span class="kw-2">&amp;</span>str = <span class="string">"k8s.service"</span>;
<a href=#104 id=104 data-nosnippet>104</a><span class="kw">pub const </span>ENTITY_TYPE_GEN_AI_AGENT: <span class="kw-2">&amp;</span>str = <span class="string">"gen_ai.agent"</span>;
<a href=#105 id=105 data-nosnippet>105</a><span class="kw">pub const </span>ENTITY_TYPE_GEN_AI_MODEL: <span class="kw-2">&amp;</span>str = <span class="string">"gen_ai.model"</span>;
<a href=#106 id=106 data-nosnippet>106</a><span class="kw">pub const </span>ENTITY_TYPE_GEN_AI_TOOL: <span class="kw-2">&amp;</span>str = <span class="string">"gen_ai.tool"</span>;
<a href=#107 id=107 data-nosnippet>107</a>
<a href=#108 id=108 data-nosnippet>108</a><span class="kw">const </span>ENTITY_TYPES: [<span class="kw-2">&amp;</span>str; <span class="number">13</span>] = [
<a href=#109 id=109 data-nosnippet>109</a> ENTITY_TYPE_SERVICE,
<a href=#110 id=110 data-nosnippet>110</a> ENTITY_TYPE_SERVICE_INSTANCE,
<a href=#111 id=111 data-nosnippet>111</a> ENTITY_TYPE_HOST,
<a href=#112 id=112 data-nosnippet>112</a> ENTITY_TYPE_CONTAINER,
<a href=#113 id=113 data-nosnippet>113</a> ENTITY_TYPE_PROCESS,
<a href=#114 id=114 data-nosnippet>114</a> ENTITY_TYPE_K8S_POD,
<a href=#115 id=115 data-nosnippet>115</a> ENTITY_TYPE_K8S_NODE,
<a href=#116 id=116 data-nosnippet>116</a> ENTITY_TYPE_K8S_CONTAINER,
<a href=#117 id=117 data-nosnippet>117</a> ENTITY_TYPE_K8S_WORKLOAD,
<a href=#118 id=118 data-nosnippet>118</a> ENTITY_TYPE_K8S_SERVICE,
<a href=#119 id=119 data-nosnippet>119</a> ENTITY_TYPE_GEN_AI_AGENT,
<a href=#120 id=120 data-nosnippet>120</a> ENTITY_TYPE_GEN_AI_MODEL,
<a href=#121 id=121 data-nosnippet>121</a> ENTITY_TYPE_GEN_AI_TOOL,
<a href=#122 id=122 data-nosnippet>122</a>];
<a href=#123 id=123 data-nosnippet>123</a>
<a href=#124 id=124 data-nosnippet>124</a><span class="doccomment">/// The relationship vocabulary (the RFC's rel_type table) and the edge
<a href=#125 id=125 data-nosnippet>125</a>/// provenances.
<a href=#126 id=126 data-nosnippet>126</a></span><span class="kw">pub const </span>REL_TYPE_CALLS: <span class="kw-2">&amp;</span>str = <span class="string">"calls"</span>;
<a href=#127 id=127 data-nosnippet>127</a><span class="kw">pub const </span>REL_TYPE_RUNS_ON: <span class="kw-2">&amp;</span>str = <span class="string">"runs_on"</span>;
<a href=#128 id=128 data-nosnippet>128</a><span class="kw">pub const </span>REL_TYPE_CONTAINS: <span class="kw-2">&amp;</span>str = <span class="string">"contains"</span>;
<a href=#129 id=129 data-nosnippet>129</a><span class="kw">pub const </span>REL_TYPE_PART_OF: <span class="kw-2">&amp;</span>str = <span class="string">"part_of"</span>;
<a href=#130 id=130 data-nosnippet>130</a><span class="kw">pub const </span>REL_TYPE_USES: <span class="kw-2">&amp;</span>str = <span class="string">"uses"</span>;
<a href=#131 id=131 data-nosnippet>131</a><span class="kw">pub const </span>REL_TYPE_INVOKES: <span class="kw-2">&amp;</span>str = <span class="string">"invokes"</span>;
<a href=#132 id=132 data-nosnippet>132</a><span class="kw">pub const </span>REL_TYPE_DEPENDS_ON: <span class="kw-2">&amp;</span>str = <span class="string">"depends_on"</span>;
<a href=#133 id=133 data-nosnippet>133</a><span class="kw">pub const </span>REL_TYPE_OWNS: <span class="kw-2">&amp;</span>str = <span class="string">"owns"</span>;
<a href=#134 id=134 data-nosnippet>134</a><span class="kw">pub const </span>PROVENANCE_TRACE: <span class="kw-2">&amp;</span>str = <span class="string">"trace"</span>;
<a href=#135 id=135 data-nosnippet>135</a><span class="kw">pub const </span>PROVENANCE_ATTRIBUTE: <span class="kw-2">&amp;</span>str = <span class="string">"attribute"</span>;
<a href=#136 id=136 data-nosnippet>136</a><span class="kw">pub const </span>PROVENANCE_DECLARED: <span class="kw-2">&amp;</span>str = <span class="string">"declared"</span>;
<a href=#137 id=137 data-nosnippet>137</a><span class="kw">pub const </span>PROVENANCE_AGENT: <span class="kw-2">&amp;</span>str = <span class="string">"agent"</span>;
<a href=#138 id=138 data-nosnippet>138</a>
<a href=#139 id=139 data-nosnippet>139</a><span class="kw">const </span>REL_TYPES: [<span class="kw-2">&amp;</span>str; <span class="number">8</span>] = [
<a href=#140 id=140 data-nosnippet>140</a> REL_TYPE_CALLS,
<a href=#141 id=141 data-nosnippet>141</a> REL_TYPE_RUNS_ON,
<a href=#142 id=142 data-nosnippet>142</a> REL_TYPE_CONTAINS,
<a href=#143 id=143 data-nosnippet>143</a> REL_TYPE_PART_OF,
<a href=#144 id=144 data-nosnippet>144</a> REL_TYPE_USES,
<a href=#145 id=145 data-nosnippet>145</a> REL_TYPE_INVOKES,
<a href=#146 id=146 data-nosnippet>146</a> REL_TYPE_DEPENDS_ON,
<a href=#147 id=147 data-nosnippet>147</a> REL_TYPE_OWNS,
<a href=#148 id=148 data-nosnippet>148</a>];
<a href=#149 id=149 data-nosnippet>149</a>
<a href=#150 id=150 data-nosnippet>150</a><span class="doccomment">/// The `connection_type` values a virtual-destination candidate may imply;
<a href=#151 id=151 data-nosnippet>151</a>/// the calls derivation branches on them when building edge attributes.
<a href=#152 id=152 data-nosnippet>152</a></span><span class="kw">pub const </span>CONNECTION_TYPE_DATABASE: <span class="kw-2">&amp;</span>str = <span class="string">"database"</span>;
<a href=#153 id=153 data-nosnippet>153</a><span class="kw">pub const </span>CONNECTION_TYPE_VIRTUAL_NODE: <span class="kw-2">&amp;</span>str = <span class="string">"virtual_node"</span>;
<a href=#154 id=154 data-nosnippet>154</a>
<a href=#155 id=155 data-nosnippet>155</a><span class="kw">const </span>CONNECTION_TYPES: [<span class="kw-2">&amp;</span>str; <span class="number">2</span>] = [CONNECTION_TYPE_DATABASE, CONNECTION_TYPE_VIRTUAL_NODE];
<a href=#156 id=156 data-nosnippet>156</a>
<a href=#157 id=157 data-nosnippet>157</a><span class="kw">static </span>CONVENTIONS: LazyLock&lt;<span class="prelude-ty">Result</span>&lt;Conventions, String&gt;&gt; =
<a href=#158 id=158 data-nosnippet>158</a> LazyLock::new(|| parse(<span class="macro">include_str!</span>(<span class="string">"conventions.yaml"</span>)));
<a href=#159 id=159 data-nosnippet>159</a>
<a href=#160 id=160 data-nosnippet>160</a><span class="doccomment">/// The embedded conventions. `Err` means the embedded file is broken — pinned
<a href=#161 id=161 data-nosnippet>161</a>/// by unit test, and propagated by the derivation paths rather than panicking.
<a href=#162 id=162 data-nosnippet>162</a></span><span class="kw">pub fn </span>conventions() -&gt; <span class="prelude-ty">Result</span>&lt;<span class="kw-2">&amp;</span><span class="lifetime">'static </span>Conventions, String&gt; {
<a href=#163 id=163 data-nosnippet>163</a> CONVENTIONS.as_ref().map_err(Clone::clone)
<a href=#164 id=164 data-nosnippet>164</a>}
<a href=#165 id=165 data-nosnippet>165</a>
<a href=#166 id=166 data-nosnippet>166</a><span class="kw">fn </span>parse(yaml: <span class="kw-2">&amp;</span>str) -&gt; <span class="prelude-ty">Result</span>&lt;Conventions, String&gt; {
<a href=#167 id=167 data-nosnippet>167</a> <span class="kw">let </span>conventions: Conventions =
<a href=#168 id=168 data-nosnippet>168</a> serde_yaml_ng::from_str(yaml).map_err(|e| <span class="macro">format!</span>(<span class="string">"malformed conventions: {e}"</span>))<span class="question-mark">?</span>;
<a href=#169 id=169 data-nosnippet>169</a> validate(<span class="kw-2">&amp;</span>conventions)<span class="question-mark">?</span>;
<a href=#170 id=170 data-nosnippet>170</a> <span class="prelude-val">Ok</span>(conventions)
<a href=#171 id=171 data-nosnippet>171</a>}
<a href=#172 id=172 data-nosnippet>172</a>
<a href=#173 id=173 data-nosnippet>173</a><span class="kw">fn </span>validate(conventions: <span class="kw-2">&amp;</span>Conventions) -&gt; <span class="prelude-ty">Result</span>&lt;(), String&gt; {
<a href=#174 id=174 data-nosnippet>174</a> <span class="kw">let </span><span class="kw-2">mut </span>seen_edges = HashSet::new();
<a href=#175 id=175 data-nosnippet>175</a> <span class="kw">for </span>rule <span class="kw">in </span>conventions
<a href=#176 id=176 data-nosnippet>176</a> .co_declared_edges
<a href=#177 id=177 data-nosnippet>177</a> .iter()
<a href=#178 id=178 data-nosnippet>178</a> .chain(<span class="kw-2">&amp;</span>conventions.trace_co_declared_edges)
<a href=#179 id=179 data-nosnippet>179</a> {
<a href=#180 id=180 data-nosnippet>180</a> <span class="kw">for </span>ty <span class="kw">in </span>[<span class="kw-2">&amp;</span>rule.src, <span class="kw-2">&amp;</span>rule.dst] {
<a href=#181 id=181 data-nosnippet>181</a> <span class="kw">if </span>!ENTITY_TYPES.contains(<span class="kw-2">&amp;</span>ty.as_str()) {
<a href=#182 id=182 data-nosnippet>182</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(<span class="string">"unknown entity type `{ty}` in edge vocabulary"</span>));
<a href=#183 id=183 data-nosnippet>183</a> }
<a href=#184 id=184 data-nosnippet>184</a> }
<a href=#185 id=185 data-nosnippet>185</a> <span class="kw">if </span>!REL_TYPES.contains(<span class="kw-2">&amp;</span>rule.rel.as_str()) {
<a href=#186 id=186 data-nosnippet>186</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#187 id=187 data-nosnippet>187</a> <span class="string">"unknown rel_type `{}` in edge vocabulary"</span>,
<a href=#188 id=188 data-nosnippet>188</a> rule.rel
<a href=#189 id=189 data-nosnippet>189</a> ));
<a href=#190 id=190 data-nosnippet>190</a> }
<a href=#191 id=191 data-nosnippet>191</a> <span class="kw">if </span>!seen_edges.insert((<span class="kw-2">&amp;</span>rule.src, <span class="kw-2">&amp;</span>rule.dst, <span class="kw-2">&amp;</span>rule.rel)) {
<a href=#192 id=192 data-nosnippet>192</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#193 id=193 data-nosnippet>193</a> <span class="string">"duplicate edge rule `{} -{}-&gt; {}`"</span>,
<a href=#194 id=194 data-nosnippet>194</a> rule.src, rule.rel, rule.dst
<a href=#195 id=195 data-nosnippet>195</a> ));
<a href=#196 id=196 data-nosnippet>196</a> }
<a href=#197 id=197 data-nosnippet>197</a> }
<a href=#198 id=198 data-nosnippet>198</a>
<a href=#199 id=199 data-nosnippet>199</a> <span class="kw">let </span><span class="kw-2">mut </span>seen_columns = HashSet::new();
<a href=#200 id=200 data-nosnippet>200</a> <span class="kw">for </span>candidate <span class="kw">in </span><span class="kw-2">&amp;</span>conventions.virtual_dst_candidates {
<a href=#201 id=201 data-nosnippet>201</a> <span class="kw">if </span>candidate.column.is_empty() {
<a href=#202 id=202 data-nosnippet>202</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="string">"empty virtual destination column"</span>.to_string());
<a href=#203 id=203 data-nosnippet>203</a> }
<a href=#204 id=204 data-nosnippet>204</a> <span class="kw">if </span>!CONNECTION_TYPES.contains(<span class="kw-2">&amp;</span>candidate.connection_type.as_str()) {
<a href=#205 id=205 data-nosnippet>205</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#206 id=206 data-nosnippet>206</a> <span class="string">"unknown connection_type `{}` for virtual destination `{}`"</span>,
<a href=#207 id=207 data-nosnippet>207</a> candidate.connection_type, candidate.column
<a href=#208 id=208 data-nosnippet>208</a> ));
<a href=#209 id=209 data-nosnippet>209</a> }
<a href=#210 id=210 data-nosnippet>210</a> <span class="kw">if </span>!seen_columns.insert(<span class="kw-2">&amp;</span>candidate.column) {
<a href=#211 id=211 data-nosnippet>211</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#212 id=212 data-nosnippet>212</a> <span class="string">"duplicate virtual destination column `{}`"</span>,
<a href=#213 id=213 data-nosnippet>213</a> candidate.column
<a href=#214 id=214 data-nosnippet>214</a> ));
<a href=#215 id=215 data-nosnippet>215</a> }
<a href=#216 id=216 data-nosnippet>216</a> }
<a href=#217 id=217 data-nosnippet>217</a>
<a href=#218 id=218 data-nosnippet>218</a> <span class="kw">let </span>per_table = conventions
<a href=#219 id=219 data-nosnippet>219</a> .prometheus_info_metrics
<a href=#220 id=220 data-nosnippet>220</a> .iter()
<a href=#221 id=221 data-nosnippet>221</a> .chain(<span class="kw-2">&amp;</span>conventions.otel_info_metrics)
<a href=#222 id=222 data-nosnippet>222</a> .map(|(table, entities)| (table.as_str(), entities))
<a href=#223 id=223 data-nosnippet>223</a> .chain(std::iter::once((
<a href=#224 id=224 data-nosnippet>224</a> <span class="string">"otlp traces"</span>,
<a href=#225 id=225 data-nosnippet>225</a> <span class="kw-2">&amp;</span>conventions.otlp_trace_entities,
<a href=#226 id=226 data-nosnippet>226</a> )));
<a href=#227 id=227 data-nosnippet>227</a> <span class="comment">// An entity type declared with a different number of id columns by two
<a href=#228 id=228 data-nosnippet>228</a> // sources yields ids that can never match, silently splitting one entity
<a href=#229 id=229 data-nosnippet>229</a> // in two.
<a href=#230 id=230 data-nosnippet>230</a> </span><span class="kw">let </span><span class="kw-2">mut </span>arity = HashMap::new();
<a href=#231 id=231 data-nosnippet>231</a> <span class="kw">for </span>(table, entities) <span class="kw">in </span>per_table {
<a href=#232 id=232 data-nosnippet>232</a> <span class="kw">let </span><span class="kw-2">mut </span>seen_types = HashSet::new();
<a href=#233 id=233 data-nosnippet>233</a> <span class="kw">for </span>implicit <span class="kw">in </span>entities {
<a href=#234 id=234 data-nosnippet>234</a> <span class="kw">if </span>!ENTITY_TYPES.contains(<span class="kw-2">&amp;</span>implicit.entity.as_str()) {
<a href=#235 id=235 data-nosnippet>235</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#236 id=236 data-nosnippet>236</a> <span class="string">"unknown entity type `{}` for info metric `{table}`"</span>,
<a href=#237 id=237 data-nosnippet>237</a> implicit.entity
<a href=#238 id=238 data-nosnippet>238</a> ));
<a href=#239 id=239 data-nosnippet>239</a> }
<a href=#240 id=240 data-nosnippet>240</a> <span class="kw">if </span>!seen_types.insert(<span class="kw-2">&amp;</span>implicit.entity) {
<a href=#241 id=241 data-nosnippet>241</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#242 id=242 data-nosnippet>242</a> <span class="string">"duplicate entity type `{}` for info metric `{table}`"</span>,
<a href=#243 id=243 data-nosnippet>243</a> implicit.entity
<a href=#244 id=244 data-nosnippet>244</a> ));
<a href=#245 id=245 data-nosnippet>245</a> }
<a href=#246 id=246 data-nosnippet>246</a> <span class="kw">if </span>implicit.id.is_empty() || implicit.id.iter().any(String::is_empty) {
<a href=#247 id=247 data-nosnippet>247</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#248 id=248 data-nosnippet>248</a> <span class="string">"entity `{}` of info metric `{table}` needs non-empty id columns"</span>,
<a href=#249 id=249 data-nosnippet>249</a> implicit.entity
<a href=#250 id=250 data-nosnippet>250</a> ));
<a href=#251 id=251 data-nosnippet>251</a> }
<a href=#252 id=252 data-nosnippet>252</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(superseding) = <span class="kw-2">&amp;</span>implicit.superseded_by {
<a href=#253 id=253 data-nosnippet>253</a> <span class="kw">let </span>superseding = entities
<a href=#254 id=254 data-nosnippet>254</a> .iter()
<a href=#255 id=255 data-nosnippet>255</a> .find(|other| <span class="kw-2">&amp;</span>other.entity == superseding)
<a href=#256 id=256 data-nosnippet>256</a> .ok_or_else(|| {
<a href=#257 id=257 data-nosnippet>257</a> <span class="macro">format!</span>(
<a href=#258 id=258 data-nosnippet>258</a> <span class="string">"entity `{}` of info metric `{table}` is superseded by `{superseding}`, \
<a href=#259 id=259 data-nosnippet>259</a> which `{table}` does not declare"</span>,
<a href=#260 id=260 data-nosnippet>260</a> implicit.entity
<a href=#261 id=261 data-nosnippet>261</a> )
<a href=#262 id=262 data-nosnippet>262</a> })<span class="question-mark">?</span>;
<a href=#263 id=263 data-nosnippet>263</a> <span class="comment">// A chain would let the middle entity withdraw while its own
<a href=#264 id=264 data-nosnippet>264</a> // replacement is absent.
<a href=#265 id=265 data-nosnippet>265</a> </span><span class="kw">if </span>superseding.superseded_by.is_some() {
<a href=#266 id=266 data-nosnippet>266</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#267 id=267 data-nosnippet>267</a> <span class="string">"entity `{}` of info metric `{table}` is superseded by `{}`, which is \
<a href=#268 id=268 data-nosnippet>268</a> itself superseded"</span>,
<a href=#269 id=269 data-nosnippet>269</a> implicit.entity, superseding.entity
<a href=#270 id=270 data-nosnippet>270</a> ));
<a href=#271 id=271 data-nosnippet>271</a> }
<a href=#272 id=272 data-nosnippet>272</a> }
<a href=#273 id=273 data-nosnippet>273</a> <span class="kw">if </span>implicit.qualified_by.as_ref().is_some_and(String::is_empty) {
<a href=#274 id=274 data-nosnippet>274</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#275 id=275 data-nosnippet>275</a> <span class="string">"entity `{}` of info metric `{table}` has an empty qualified_by"</span>,
<a href=#276 id=276 data-nosnippet>276</a> implicit.entity
<a href=#277 id=277 data-nosnippet>277</a> ));
<a href=#278 id=278 data-nosnippet>278</a> }
<a href=#279 id=279 data-nosnippet>279</a> <span class="kw">if </span>implicit.descriptive_rest &amp;&amp; !implicit.descriptive.is_empty() {
<a href=#280 id=280 data-nosnippet>280</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#281 id=281 data-nosnippet>281</a> <span class="string">"entity `{}` of info metric `{table}` sets both descriptive and \
<a href=#282 id=282 data-nosnippet>282</a> descriptive_rest"</span>,
<a href=#283 id=283 data-nosnippet>283</a> implicit.entity
<a href=#284 id=284 data-nosnippet>284</a> ));
<a href=#285 id=285 data-nosnippet>285</a> }
<a href=#286 id=286 data-nosnippet>286</a> <span class="kw">match </span>arity.insert(implicit.entity.as_str(), implicit.id.len()) {
<a href=#287 id=287 data-nosnippet>287</a> <span class="prelude-val">Some</span>(previous) <span class="kw">if </span>previous != implicit.id.len() =&gt; {
<a href=#288 id=288 data-nosnippet>288</a> <span class="kw">return </span><span class="prelude-val">Err</span>(<span class="macro">format!</span>(
<a href=#289 id=289 data-nosnippet>289</a> <span class="string">"entity `{}` is declared with {previous} id columns elsewhere but \
<a href=#290 id=290 data-nosnippet>290</a> {} for `{table}`"</span>,
<a href=#291 id=291 data-nosnippet>291</a> implicit.entity,
<a href=#292 id=292 data-nosnippet>292</a> implicit.id.len()
<a href=#293 id=293 data-nosnippet>293</a> ));
<a href=#294 id=294 data-nosnippet>294</a> }
<a href=#295 id=295 data-nosnippet>295</a> <span class="kw">_ </span>=&gt; {}
<a href=#296 id=296 data-nosnippet>296</a> }
<a href=#297 id=297 data-nosnippet>297</a> }
<a href=#298 id=298 data-nosnippet>298</a> }
<a href=#299 id=299 data-nosnippet>299</a> <span class="prelude-val">Ok</span>(())
<a href=#300 id=300 data-nosnippet>300</a>}
<a href=#301 id=301 data-nosnippet>301</a>
<a href=#302 id=302 data-nosnippet>302</a><span class="attr">#[cfg(test)]
<a href=#303 id=303 data-nosnippet>303</a></span><span class="kw">mod </span>tests {
<a href=#304 id=304 data-nosnippet>304</a> <span class="kw">use super</span>::<span class="kw-2">*</span>;
<a href=#305 id=305 data-nosnippet>305</a>
<a href=#306 id=306 data-nosnippet>306</a> <span class="attr">#[test]
<a href=#307 id=307 data-nosnippet>307</a> </span><span class="kw">fn </span>embedded_conventions_parse_and_validate() {
<a href=#308 id=308 data-nosnippet>308</a> conventions().unwrap();
<a href=#309 id=309 data-nosnippet>309</a> }
<a href=#310 id=310 data-nosnippet>310</a>
<a href=#311 id=311 data-nosnippet>311</a> <span class="doccomment">/// Each case must fail on exactly the rule it names, so the shared
<a href=#312 id=312 data-nosnippet>312</a> /// boilerplate is valid and the mutated part is inside the vocabulary.
<a href=#313 id=313 data-nosnippet>313</a> </span><span class="kw">fn </span>broken(edges: <span class="kw-2">&amp;</span>str, info_metrics: <span class="kw-2">&amp;</span>str, otel_metrics: <span class="kw-2">&amp;</span>str) -&gt; String {
<a href=#314 id=314 data-nosnippet>314</a> <span class="macro">format!</span>(
<a href=#315 id=315 data-nosnippet>315</a> <span class="string">"co_declared_edges: [{edges}]\ntrace_co_declared_edges: []\n\
<a href=#316 id=316 data-nosnippet>316</a> virtual_dst_candidates: []\notlp_trace_entities: []\n\
<a href=#317 id=317 data-nosnippet>317</a> prometheus_info_metrics: {{{info_metrics}}}\n\
<a href=#318 id=318 data-nosnippet>318</a> otel_info_metrics: {{{otel_metrics}}}"
<a href=#319 id=319 data-nosnippet>319</a> </span>)
<a href=#320 id=320 data-nosnippet>320</a> }
<a href=#321 id=321 data-nosnippet>321</a>
<a href=#322 id=322 data-nosnippet>322</a> <span class="attr">#[test]
<a href=#323 id=323 data-nosnippet>323</a> </span><span class="kw">fn </span>validation_rejects_broken_conventions() {
<a href=#324 id=324 data-nosnippet>324</a> <span class="kw">let </span>err = |edges, info, otel| parse(<span class="kw-2">&amp;</span>broken(edges, info, otel)).unwrap_err();
<a href=#325 id=325 data-nosnippet>325</a>
<a href=#326 id=326 data-nosnippet>326</a> <span class="macro">assert!</span>(err(<span class="string">"{src: host, dst: service, rel: pets}"</span>, <span class="string">""</span>, <span class="string">""</span>).contains(<span class="string">"unknown rel_type"</span>));
<a href=#327 id=327 data-nosnippet>327</a> <span class="macro">assert!</span>(
<a href=#328 id=328 data-nosnippet>328</a> err(<span class="string">"{src: k8s.pods, dst: k8s.node, rel: runs_on}"</span>, <span class="string">""</span>, <span class="string">""</span>)
<a href=#329 id=329 data-nosnippet>329</a> .contains(<span class="string">"unknown entity type"</span>)
<a href=#330 id=330 data-nosnippet>330</a> );
<a href=#331 id=331 data-nosnippet>331</a> <span class="macro">assert!</span>(
<a href=#332 id=332 data-nosnippet>332</a> err(
<a href=#333 id=333 data-nosnippet>333</a> <span class="string">"{src: host, dst: service, rel: uses}, {src: host, dst: service, rel: uses}"</span>,
<a href=#334 id=334 data-nosnippet>334</a> <span class="string">""</span>,
<a href=#335 id=335 data-nosnippet>335</a> <span class="string">""
<a href=#336 id=336 data-nosnippet>336</a> </span>)
<a href=#337 id=337 data-nosnippet>337</a> .contains(<span class="string">"duplicate edge rule"</span>)
<a href=#338 id=338 data-nosnippet>338</a> );
<a href=#339 id=339 data-nosnippet>339</a> <span class="macro">assert!</span>(
<a href=#340 id=340 data-nosnippet>340</a> err(
<a href=#341 id=341 data-nosnippet>341</a> <span class="string">""</span>,
<a href=#342 id=342 data-nosnippet>342</a> <span class="string">"t: [{entity: host, id: [x], descriptive: [y], descriptive_rest: true}]"</span>,
<a href=#343 id=343 data-nosnippet>343</a> <span class="string">""
<a href=#344 id=344 data-nosnippet>344</a> </span>)
<a href=#345 id=345 data-nosnippet>345</a> .contains(<span class="string">"descriptive_rest"</span>)
<a href=#346 id=346 data-nosnippet>346</a> );
<a href=#347 id=347 data-nosnippet>347</a> <span class="comment">// Superseding a type the same table does not declare leaves the rule
<a href=#348 id=348 data-nosnippet>348</a> // dead and the duplicate node back in the graph.
<a href=#349 id=349 data-nosnippet>349</a> </span><span class="macro">assert!</span>(
<a href=#350 id=350 data-nosnippet>350</a> err(
<a href=#351 id=351 data-nosnippet>351</a> <span class="string">""</span>,
<a href=#352 id=352 data-nosnippet>352</a> <span class="string">"t: [{entity: container, id: [x], superseded_by: k8s.container}]"</span>,
<a href=#353 id=353 data-nosnippet>353</a> <span class="string">""
<a href=#354 id=354 data-nosnippet>354</a> </span>)
<a href=#355 id=355 data-nosnippet>355</a> .contains(<span class="string">"does not declare"</span>)
<a href=#356 id=356 data-nosnippet>356</a> );
<a href=#357 id=357 data-nosnippet>357</a> <span class="comment">// the otel map runs through the same per-table validation
<a href=#358 id=358 data-nosnippet>358</a> </span><span class="macro">assert!</span>(err(<span class="string">""</span>, <span class="string">""</span>, <span class="string">"t: [{entity: hosts, id: [x]}]"</span>).contains(<span class="string">"unknown entity type"</span>));
<a href=#359 id=359 data-nosnippet>359</a> <span class="comment">// ids of a different arity for one type can never match each other
<a href=#360 id=360 data-nosnippet>360</a> </span><span class="macro">assert!</span>(
<a href=#361 id=361 data-nosnippet>361</a> err(
<a href=#362 id=362 data-nosnippet>362</a> <span class="string">""</span>,
<a href=#363 id=363 data-nosnippet>363</a> <span class="string">"a: [{entity: host, id: [x]}]"</span>,
<a href=#364 id=364 data-nosnippet>364</a> <span class="string">"b: [{entity: host, id: [x, y]}]"
<a href=#365 id=365 data-nosnippet>365</a> </span>)
<a href=#366 id=366 data-nosnippet>366</a> .contains(<span class="string">"id columns"</span>)
<a href=#367 id=367 data-nosnippet>367</a> );
<a href=#368 id=368 data-nosnippet>368</a> <span class="comment">// Unknown YAML keys are rejected, catching typos in the embedded file.
<a href=#369 id=369 data-nosnippet>369</a> </span><span class="macro">assert!</span>(
<a href=#370 id=370 data-nosnippet>370</a> parse(<span class="kw-2">&amp;</span>broken(
<a href=#371 id=371 data-nosnippet>371</a> <span class="string">"{src: host, dst: service, rel: uses, direction: down}"</span>,
<a href=#372 id=372 data-nosnippet>372</a> <span class="string">""</span>,
<a href=#373 id=373 data-nosnippet>373</a> <span class="string">""
<a href=#374 id=374 data-nosnippet>374</a> </span>))
<a href=#375 id=375 data-nosnippet>375</a> .is_err()
<a href=#376 id=376 data-nosnippet>376</a> );
<a href=#377 id=377 data-nosnippet>377</a> }
<a href=#378 id=378 data-nosnippet>378</a>}
</code></pre></div></section></main></body></html>
File diff suppressed because it is too large Load Diff
+502
View File
@@ -0,0 +1,502 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/set.rs`."><title>set.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>set.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::str::FromStr;
<a href=#16 id=16 data-nosnippet>16</a><span class="kw">use </span>std::time::Duration;
<a href=#17 id=17 data-nosnippet>17</a>
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_time::Timezone;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>lazy_static::lazy_static;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>regex::Regex;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>session::ReadPreference;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>session::context::Channel::Postgres;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>session::session_config::{PGByteaOutputValue, PGDateOrder, PGDateTimeStyle, PGIntervalStyle};
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>snafu::{IntoError, OptionExt, ResultExt, ensure};
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>sql::ast::{Expr, Ident, Value};
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>sql::statements::set_variables::SetVariables;
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>sqlparser::ast::ValueWithSpan;
<a href=#29 id=29 data-nosnippet>29</a>
<a href=#30 id=30 data-nosnippet>30</a><span class="kw">use </span><span class="kw">crate</span>::error::{InvalidConfigValueSnafu, InvalidSqlSnafu, NotSupportedSnafu, <span class="prelude-ty">Result</span>};
<a href=#31 id=31 data-nosnippet>31</a>
<a href=#32 id=32 data-nosnippet>32</a><span class="macro">lazy_static!</span> {
<a href=#33 id=33 data-nosnippet>33</a> <span class="comment">// Regex rules:
<a href=#34 id=34 data-nosnippet>34</a> // The string must start with a number (one or more digits).
<a href=#35 id=35 data-nosnippet>35</a> // The number must be followed by one of the valid time units (ms, s, min, h, d).
<a href=#36 id=36 data-nosnippet>36</a> // The string must end immediately after the unit, meaning there can be no extra
<a href=#37 id=37 data-nosnippet>37</a> // characters or spaces after the valid time specification.
<a href=#38 id=38 data-nosnippet>38</a> </span><span class="kw">static </span><span class="kw-2">ref </span>PG_TIME_INPUT_REGEX: Regex = Regex::new(<span class="string">r"^(\d+)(ms|s|min|h|d)$"</span>).unwrap();
<a href=#39 id=39 data-nosnippet>39</a>}
<a href=#40 id=40 data-nosnippet>40</a>
<a href=#41 id=41 data-nosnippet>41</a><span class="doccomment">/// Sets the session WAL policy for ordinary inserts without changing table options.
<a href=#42 id=42 data-nosnippet>42</a></span><span class="kw">pub fn </span>set_skip_wal(exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#43 id=43 data-nosnippet>43</a> <span class="kw">let </span>[Expr::Value(value)] = exprs.as_slice() <span class="kw">else </span>{
<a href=#44 id=44 data-nosnippet>44</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#45 id=45 data-nosnippet>45</a> feat: <span class="string">"SET skip_wal requires exactly one boolean value"</span>,
<a href=#46 id=46 data-nosnippet>46</a> }
<a href=#47 id=47 data-nosnippet>47</a> .fail();
<a href=#48 id=48 data-nosnippet>48</a> };
<a href=#49 id=49 data-nosnippet>49</a> <span class="kw">let </span>skip_wal = <span class="kw">match </span><span class="kw-2">&amp;</span>value.value {
<a href=#50 id=50 data-nosnippet>50</a> Value::Boolean(value) =&gt; <span class="kw-2">*</span>value,
<a href=#51 id=51 data-nosnippet>51</a> Value::SingleQuotedString(value) | Value::DoubleQuotedString(value) =&gt; {
<a href=#52 id=52 data-nosnippet>52</a> value.parse::&lt;bool&gt;().map_err(|<span class="kw">_</span>| {
<a href=#53 id=53 data-nosnippet>53</a> NotSupportedSnafu {
<a href=#54 id=54 data-nosnippet>54</a> feat: <span class="macro">format!</span>(<span class="string">"Invalid skip_wal value {value:?}: expected true or false"</span>),
<a href=#55 id=55 data-nosnippet>55</a> }
<a href=#56 id=56 data-nosnippet>56</a> .build()
<a href=#57 id=57 data-nosnippet>57</a> })<span class="question-mark">?
<a href=#58 id=58 data-nosnippet>58</a> </span>}
<a href=#59 id=59 data-nosnippet>59</a> <span class="kw">_ </span>=&gt; {
<a href=#60 id=60 data-nosnippet>60</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#61 id=61 data-nosnippet>61</a> feat: <span class="string">"SET skip_wal requires true or false"</span>,
<a href=#62 id=62 data-nosnippet>62</a> }
<a href=#63 id=63 data-nosnippet>63</a> .fail();
<a href=#64 id=64 data-nosnippet>64</a> }
<a href=#65 id=65 data-nosnippet>65</a> };
<a href=#66 id=66 data-nosnippet>66</a> ctx.set_skip_wal(skip_wal);
<a href=#67 id=67 data-nosnippet>67</a> <span class="prelude-val">Ok</span>(())
<a href=#68 id=68 data-nosnippet>68</a>}
<a href=#69 id=69 data-nosnippet>69</a>
<a href=#70 id=70 data-nosnippet>70</a><span class="kw">pub fn </span>set_read_preference(exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#71 id=71 data-nosnippet>71</a> <span class="kw">let </span>read_preference_expr = exprs.first().context(NotSupportedSnafu {
<a href=#72 id=72 data-nosnippet>72</a> feat: <span class="string">"No read preference find in set variable statement"</span>,
<a href=#73 id=73 data-nosnippet>73</a> })<span class="question-mark">?</span>;
<a href=#74 id=74 data-nosnippet>74</a>
<a href=#75 id=75 data-nosnippet>75</a> <span class="kw">match </span>read_preference_expr {
<a href=#76 id=76 data-nosnippet>76</a> Expr::Value(ValueWithSpan {
<a href=#77 id=77 data-nosnippet>77</a> value: Value::SingleQuotedString(expr),
<a href=#78 id=78 data-nosnippet>78</a> ..
<a href=#79 id=79 data-nosnippet>79</a> })
<a href=#80 id=80 data-nosnippet>80</a> | Expr::Value(ValueWithSpan {
<a href=#81 id=81 data-nosnippet>81</a> value: Value::DoubleQuotedString(expr),
<a href=#82 id=82 data-nosnippet>82</a> ..
<a href=#83 id=83 data-nosnippet>83</a> }) =&gt; {
<a href=#84 id=84 data-nosnippet>84</a> <span class="kw">match </span>ReadPreference::from_str(expr.as_str().to_lowercase().as_str()) {
<a href=#85 id=85 data-nosnippet>85</a> <span class="prelude-val">Ok</span>(read_preference) =&gt; ctx.set_read_preference(read_preference),
<a href=#86 id=86 data-nosnippet>86</a> <span class="prelude-val">Err</span>(<span class="kw">_</span>) =&gt; {
<a href=#87 id=87 data-nosnippet>87</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#88 id=88 data-nosnippet>88</a> feat: <span class="macro">format!</span>(
<a href=#89 id=89 data-nosnippet>89</a> <span class="string">"Invalid read preference expr {} in set variable statement"</span>,
<a href=#90 id=90 data-nosnippet>90</a> expr,
<a href=#91 id=91 data-nosnippet>91</a> ),
<a href=#92 id=92 data-nosnippet>92</a> }
<a href=#93 id=93 data-nosnippet>93</a> .fail();
<a href=#94 id=94 data-nosnippet>94</a> }
<a href=#95 id=95 data-nosnippet>95</a> }
<a href=#96 id=96 data-nosnippet>96</a> <span class="prelude-val">Ok</span>(())
<a href=#97 id=97 data-nosnippet>97</a> }
<a href=#98 id=98 data-nosnippet>98</a> expr =&gt; NotSupportedSnafu {
<a href=#99 id=99 data-nosnippet>99</a> feat: <span class="macro">format!</span>(
<a href=#100 id=100 data-nosnippet>100</a> <span class="string">"Unsupported read preference expr {} in set variable statement"</span>,
<a href=#101 id=101 data-nosnippet>101</a> expr
<a href=#102 id=102 data-nosnippet>102</a> ),
<a href=#103 id=103 data-nosnippet>103</a> }
<a href=#104 id=104 data-nosnippet>104</a> .fail(),
<a href=#105 id=105 data-nosnippet>105</a> }
<a href=#106 id=106 data-nosnippet>106</a>}
<a href=#107 id=107 data-nosnippet>107</a>
<a href=#108 id=108 data-nosnippet>108</a><span class="kw">pub fn </span>set_timezone(exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#109 id=109 data-nosnippet>109</a> <span class="kw">let </span>tz_expr = exprs.first().context(NotSupportedSnafu {
<a href=#110 id=110 data-nosnippet>110</a> feat: <span class="string">"No timezone find in set variable statement"</span>,
<a href=#111 id=111 data-nosnippet>111</a> })<span class="question-mark">?</span>;
<a href=#112 id=112 data-nosnippet>112</a> <span class="kw">match </span>tz_expr {
<a href=#113 id=113 data-nosnippet>113</a> Expr::Value(ValueWithSpan {
<a href=#114 id=114 data-nosnippet>114</a> value: Value::SingleQuotedString(tz),
<a href=#115 id=115 data-nosnippet>115</a> ..
<a href=#116 id=116 data-nosnippet>116</a> })
<a href=#117 id=117 data-nosnippet>117</a> | Expr::Value(ValueWithSpan {
<a href=#118 id=118 data-nosnippet>118</a> value: Value::DoubleQuotedString(tz),
<a href=#119 id=119 data-nosnippet>119</a> ..
<a href=#120 id=120 data-nosnippet>120</a> }) =&gt; {
<a href=#121 id=121 data-nosnippet>121</a> <span class="kw">match </span>Timezone::from_tz_string(tz.as_str()) {
<a href=#122 id=122 data-nosnippet>122</a> <span class="prelude-val">Ok</span>(timezone) =&gt; ctx.set_timezone(timezone),
<a href=#123 id=123 data-nosnippet>123</a> <span class="prelude-val">Err</span>(<span class="kw">_</span>) =&gt; {
<a href=#124 id=124 data-nosnippet>124</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#125 id=125 data-nosnippet>125</a> feat: <span class="macro">format!</span>(<span class="string">"Invalid timezone expr {} in set variable statement"</span>, tz),
<a href=#126 id=126 data-nosnippet>126</a> }
<a href=#127 id=127 data-nosnippet>127</a> .fail();
<a href=#128 id=128 data-nosnippet>128</a> }
<a href=#129 id=129 data-nosnippet>129</a> }
<a href=#130 id=130 data-nosnippet>130</a> <span class="prelude-val">Ok</span>(())
<a href=#131 id=131 data-nosnippet>131</a> }
<a href=#132 id=132 data-nosnippet>132</a> expr =&gt; NotSupportedSnafu {
<a href=#133 id=133 data-nosnippet>133</a> feat: <span class="macro">format!</span>(
<a href=#134 id=134 data-nosnippet>134</a> <span class="string">"Unsupported timezone expr {} in set variable statement"</span>,
<a href=#135 id=135 data-nosnippet>135</a> expr
<a href=#136 id=136 data-nosnippet>136</a> ),
<a href=#137 id=137 data-nosnippet>137</a> }
<a href=#138 id=138 data-nosnippet>138</a> .fail(),
<a href=#139 id=139 data-nosnippet>139</a> }
<a href=#140 id=140 data-nosnippet>140</a>}
<a href=#141 id=141 data-nosnippet>141</a>
<a href=#142 id=142 data-nosnippet>142</a><span class="kw">pub fn </span>set_bytea_output(exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#143 id=143 data-nosnippet>143</a> <span class="kw">let </span><span class="prelude-val">Some</span>((var_value, [])) = exprs.split_first() <span class="kw">else </span>{
<a href=#144 id=144 data-nosnippet>144</a> <span class="kw">return </span>(NotSupportedSnafu {
<a href=#145 id=145 data-nosnippet>145</a> feat: <span class="string">"Set variable value must have one and only one value for bytea_output"</span>,
<a href=#146 id=146 data-nosnippet>146</a> })
<a href=#147 id=147 data-nosnippet>147</a> .fail();
<a href=#148 id=148 data-nosnippet>148</a> };
<a href=#149 id=149 data-nosnippet>149</a> <span class="kw">let </span>Expr::Value(value) = var_value <span class="kw">else </span>{
<a href=#150 id=150 data-nosnippet>150</a> <span class="kw">return </span>(NotSupportedSnafu {
<a href=#151 id=151 data-nosnippet>151</a> feat: <span class="string">"Set variable value must be a value"</span>,
<a href=#152 id=152 data-nosnippet>152</a> })
<a href=#153 id=153 data-nosnippet>153</a> .fail();
<a href=#154 id=154 data-nosnippet>154</a> };
<a href=#155 id=155 data-nosnippet>155</a> ctx.configuration_parameter().set_postgres_bytea_output(
<a href=#156 id=156 data-nosnippet>156</a> PGByteaOutputValue::try_from(value.value.clone()).context(InvalidConfigValueSnafu)<span class="question-mark">?</span>,
<a href=#157 id=157 data-nosnippet>157</a> );
<a href=#158 id=158 data-nosnippet>158</a> <span class="prelude-val">Ok</span>(())
<a href=#159 id=159 data-nosnippet>159</a>}
<a href=#160 id=160 data-nosnippet>160</a>
<a href=#161 id=161 data-nosnippet>161</a><span class="kw">pub fn </span>validate_client_encoding(set: SetVariables) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#162 id=162 data-nosnippet>162</a> <span class="kw">let </span><span class="prelude-val">Some</span>((encoding, [])) = set.value.split_first() <span class="kw">else </span>{
<a href=#163 id=163 data-nosnippet>163</a> <span class="kw">return </span>InvalidSqlSnafu {
<a href=#164 id=164 data-nosnippet>164</a> err_msg: <span class="string">"must provide one and only one client encoding value"</span>,
<a href=#165 id=165 data-nosnippet>165</a> }
<a href=#166 id=166 data-nosnippet>166</a> .fail();
<a href=#167 id=167 data-nosnippet>167</a> };
<a href=#168 id=168 data-nosnippet>168</a> <span class="kw">let </span>encoding = <span class="kw">match </span>encoding {
<a href=#169 id=169 data-nosnippet>169</a> Expr::Value(ValueWithSpan {
<a href=#170 id=170 data-nosnippet>170</a> value: Value::SingleQuotedString(x),
<a href=#171 id=171 data-nosnippet>171</a> ..
<a href=#172 id=172 data-nosnippet>172</a> })
<a href=#173 id=173 data-nosnippet>173</a> | Expr::Identifier(Ident {
<a href=#174 id=174 data-nosnippet>174</a> value: x,
<a href=#175 id=175 data-nosnippet>175</a> quote_style: <span class="kw">_</span>,
<a href=#176 id=176 data-nosnippet>176</a> span: <span class="kw">_</span>,
<a href=#177 id=177 data-nosnippet>177</a> }) =&gt; x.to_uppercase(),
<a href=#178 id=178 data-nosnippet>178</a> <span class="kw">_ </span>=&gt; {
<a href=#179 id=179 data-nosnippet>179</a> <span class="kw">return </span>InvalidSqlSnafu {
<a href=#180 id=180 data-nosnippet>180</a> err_msg: <span class="macro">format!</span>(<span class="string">"client encoding must be a string, actual: {:?}"</span>, encoding),
<a href=#181 id=181 data-nosnippet>181</a> }
<a href=#182 id=182 data-nosnippet>182</a> .fail();
<a href=#183 id=183 data-nosnippet>183</a> }
<a href=#184 id=184 data-nosnippet>184</a> };
<a href=#185 id=185 data-nosnippet>185</a> <span class="comment">// For the sake of simplicity, we only support "UTF8" ("UNICODE" is the alias for it,
<a href=#186 id=186 data-nosnippet>186</a> // see https://www.postgresql.org/docs/current/multibyte.html#MULTIBYTE-CHARSET-SUPPORTED).
<a href=#187 id=187 data-nosnippet>187</a> // "UTF8" is universal and sufficient for almost all cases.
<a href=#188 id=188 data-nosnippet>188</a> // GreptimeDB itself is always using "UTF8" as the internal encoding.
<a href=#189 id=189 data-nosnippet>189</a> </span><span class="macro">ensure!</span>(
<a href=#190 id=190 data-nosnippet>190</a> encoding == <span class="string">"UTF8" </span>|| encoding == <span class="string">"UNICODE"</span>,
<a href=#191 id=191 data-nosnippet>191</a> NotSupportedSnafu {
<a href=#192 id=192 data-nosnippet>192</a> feat: <span class="macro">format!</span>(<span class="string">"client encoding of '{}'"</span>, encoding)
<a href=#193 id=193 data-nosnippet>193</a> }
<a href=#194 id=194 data-nosnippet>194</a> );
<a href=#195 id=195 data-nosnippet>195</a> <span class="prelude-val">Ok</span>(())
<a href=#196 id=196 data-nosnippet>196</a>}
<a href=#197 id=197 data-nosnippet>197</a>
<a href=#198 id=198 data-nosnippet>198</a><span class="comment">// if one of original value and new value is none, return the other one
<a href=#199 id=199 data-nosnippet>199</a>// returns new values only when it equals to original one else return error.
<a href=#200 id=200 data-nosnippet>200</a>// This is only used for handling datestyle
<a href=#201 id=201 data-nosnippet>201</a></span><span class="kw">fn </span>merge_datestyle_value&lt;T&gt;(value: <span class="prelude-ty">Option</span>&lt;T&gt;, new_value: <span class="prelude-ty">Option</span>&lt;T&gt;) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="prelude-ty">Option</span>&lt;T&gt;&gt;
<a href=#202 id=202 data-nosnippet>202</a><span class="kw">where
<a href=#203 id=203 data-nosnippet>203</a> </span>T: PartialEq,
<a href=#204 id=204 data-nosnippet>204</a>{
<a href=#205 id=205 data-nosnippet>205</a> <span class="kw">match </span>(<span class="kw-2">&amp;</span>value, <span class="kw-2">&amp;</span>new_value) {
<a href=#206 id=206 data-nosnippet>206</a> (<span class="prelude-val">None</span>, <span class="kw">_</span>) =&gt; <span class="prelude-val">Ok</span>(new_value),
<a href=#207 id=207 data-nosnippet>207</a> (<span class="kw">_</span>, <span class="prelude-val">None</span>) =&gt; <span class="prelude-val">Ok</span>(value),
<a href=#208 id=208 data-nosnippet>208</a> (<span class="prelude-val">Some</span>(v1), <span class="prelude-val">Some</span>(v2)) <span class="kw">if </span>v1 == v2 =&gt; <span class="prelude-val">Ok</span>(new_value),
<a href=#209 id=209 data-nosnippet>209</a> <span class="kw">_ </span>=&gt; InvalidSqlSnafu {
<a href=#210 id=210 data-nosnippet>210</a> err_msg: <span class="string">"Conflicting \"datestyle\" specifications."</span>,
<a href=#211 id=211 data-nosnippet>211</a> }
<a href=#212 id=212 data-nosnippet>212</a> .fail(),
<a href=#213 id=213 data-nosnippet>213</a> }
<a href=#214 id=214 data-nosnippet>214</a>}
<a href=#215 id=215 data-nosnippet>215</a>
<a href=#216 id=216 data-nosnippet>216</a><span class="kw">fn </span>try_parse_datestyle(expr: <span class="kw-2">&amp;</span>Expr) -&gt; <span class="prelude-ty">Result</span>&lt;(<span class="prelude-ty">Option</span>&lt;PGDateTimeStyle&gt;, <span class="prelude-ty">Option</span>&lt;PGDateOrder&gt;)&gt; {
<a href=#217 id=217 data-nosnippet>217</a> <span class="kw">enum </span>ParsedDateStyle {
<a href=#218 id=218 data-nosnippet>218</a> Order(PGDateOrder),
<a href=#219 id=219 data-nosnippet>219</a> Style(PGDateTimeStyle),
<a href=#220 id=220 data-nosnippet>220</a> }
<a href=#221 id=221 data-nosnippet>221</a> <span class="kw">fn </span>try_parse_str(s: <span class="kw-2">&amp;</span>str) -&gt; <span class="prelude-ty">Result</span>&lt;ParsedDateStyle&gt; {
<a href=#222 id=222 data-nosnippet>222</a> PGDateTimeStyle::try_from(s)
<a href=#223 id=223 data-nosnippet>223</a> .map_or_else(
<a href=#224 id=224 data-nosnippet>224</a> |<span class="kw">_</span>| PGDateOrder::try_from(s).map(ParsedDateStyle::Order),
<a href=#225 id=225 data-nosnippet>225</a> |style| <span class="prelude-val">Ok</span>(ParsedDateStyle::Style(style)),
<a href=#226 id=226 data-nosnippet>226</a> )
<a href=#227 id=227 data-nosnippet>227</a> .context(InvalidConfigValueSnafu)
<a href=#228 id=228 data-nosnippet>228</a> }
<a href=#229 id=229 data-nosnippet>229</a> <span class="kw">match </span>expr {
<a href=#230 id=230 data-nosnippet>230</a> Expr::Identifier(Ident {
<a href=#231 id=231 data-nosnippet>231</a> value: s,
<a href=#232 id=232 data-nosnippet>232</a> quote_style: <span class="kw">_</span>,
<a href=#233 id=233 data-nosnippet>233</a> span: <span class="kw">_</span>,
<a href=#234 id=234 data-nosnippet>234</a> })
<a href=#235 id=235 data-nosnippet>235</a> | Expr::Value(ValueWithSpan {
<a href=#236 id=236 data-nosnippet>236</a> value: Value::SingleQuotedString(s),
<a href=#237 id=237 data-nosnippet>237</a> ..
<a href=#238 id=238 data-nosnippet>238</a> })
<a href=#239 id=239 data-nosnippet>239</a> | Expr::Value(ValueWithSpan {
<a href=#240 id=240 data-nosnippet>240</a> value: Value::DoubleQuotedString(s),
<a href=#241 id=241 data-nosnippet>241</a> ..
<a href=#242 id=242 data-nosnippet>242</a> }) =&gt; s
<a href=#243 id=243 data-nosnippet>243</a> .split(<span class="string">','</span>)
<a href=#244 id=244 data-nosnippet>244</a> .map(|s| s.trim())
<a href=#245 id=245 data-nosnippet>245</a> .try_fold((<span class="prelude-val">None</span>, <span class="prelude-val">None</span>), |(style, order), s| <span class="kw">match </span>try_parse_str(s)<span class="question-mark">? </span>{
<a href=#246 id=246 data-nosnippet>246</a> ParsedDateStyle::Order(o) =&gt; <span class="prelude-val">Ok</span>((style, merge_datestyle_value(order, <span class="prelude-val">Some</span>(o))<span class="question-mark">?</span>)),
<a href=#247 id=247 data-nosnippet>247</a> ParsedDateStyle::Style(s) =&gt; <span class="prelude-val">Ok</span>((merge_datestyle_value(style, <span class="prelude-val">Some</span>(s))<span class="question-mark">?</span>, order)),
<a href=#248 id=248 data-nosnippet>248</a> }),
<a href=#249 id=249 data-nosnippet>249</a> <span class="kw">_ </span>=&gt; NotSupportedSnafu {
<a href=#250 id=250 data-nosnippet>250</a> feat: <span class="string">"Not supported expression for datestyle"</span>,
<a href=#251 id=251 data-nosnippet>251</a> }
<a href=#252 id=252 data-nosnippet>252</a> .fail(),
<a href=#253 id=253 data-nosnippet>253</a> }
<a href=#254 id=254 data-nosnippet>254</a>}
<a href=#255 id=255 data-nosnippet>255</a>
<a href=#256 id=256 data-nosnippet>256</a><span class="doccomment">/// Applies supported query options and rejects unsupported owned namespaces.
<a href=#257 id=257 data-nosnippet>257</a></span><span class="kw">pub fn </span>set_query_option(variable: <span class="kw-2">&amp;</span>str, exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;bool&gt; {
<a href=#258 id=258 data-nosnippet>258</a> <span class="kw">let </span>lower = variable.to_ascii_lowercase();
<a href=#259 id=259 data-nosnippet>259</a> <span class="kw">if </span>!lower.starts_with(<span class="string">"query."</span>)
<a href=#260 id=260 data-nosnippet>260</a> &amp;&amp; !lower.starts_with(<span class="string">"datafusion."</span>)
<a href=#261 id=261 data-nosnippet>261</a> &amp;&amp; session::query_options::canonical_query_option_name(variable).is_none()
<a href=#262 id=262 data-nosnippet>262</a> {
<a href=#263 id=263 data-nosnippet>263</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(<span class="bool-val">false</span>);
<a href=#264 id=264 data-nosnippet>264</a> }
<a href=#265 id=265 data-nosnippet>265</a> <span class="kw">let </span>[Expr::Value(value)] = exprs.as_slice() <span class="kw">else </span>{
<a href=#266 id=266 data-nosnippet>266</a> <span class="kw">return </span><span class="prelude-val">Err</span>(InvalidConfigValueSnafu.into_error(
<a href=#267 id=267 data-nosnippet>267</a> session::session_config::Error::InvalidConfigValue {
<a href=#268 id=268 data-nosnippet>268</a> name: variable.to_string(),
<a href=#269 id=269 data-nosnippet>269</a> value: exprs
<a href=#270 id=270 data-nosnippet>270</a> .iter()
<a href=#271 id=271 data-nosnippet>271</a> .map(ToString::to_string)
<a href=#272 id=272 data-nosnippet>272</a> .collect::&lt;Vec&lt;<span class="kw">_</span>&gt;&gt;()
<a href=#273 id=273 data-nosnippet>273</a> .join(<span class="string">", "</span>),
<a href=#274 id=274 data-nosnippet>274</a> hint: <span class="string">"Expected exactly one literal value"</span>.to_string(),
<a href=#275 id=275 data-nosnippet>275</a> location: snafu::Location::new(<span class="macro">file!</span>(), <span class="macro">line!</span>(), <span class="macro">column!</span>()),
<a href=#276 id=276 data-nosnippet>276</a> },
<a href=#277 id=277 data-nosnippet>277</a> ));
<a href=#278 id=278 data-nosnippet>278</a> };
<a href=#279 id=279 data-nosnippet>279</a> <span class="kw">let </span>value = <span class="kw">match </span><span class="kw-2">&amp;</span>value.value {
<a href=#280 id=280 data-nosnippet>280</a> Value::Boolean(value) =&gt; value.to_string(),
<a href=#281 id=281 data-nosnippet>281</a> Value::Number(value, <span class="kw">_</span>)
<a href=#282 id=282 data-nosnippet>282</a> | Value::SingleQuotedString(value)
<a href=#283 id=283 data-nosnippet>283</a> | Value::DoubleQuotedString(value) =&gt; value.clone(),
<a href=#284 id=284 data-nosnippet>284</a> <span class="kw">_ </span>=&gt; {
<a href=#285 id=285 data-nosnippet>285</a> <span class="kw">return </span><span class="prelude-val">Err</span>(InvalidConfigValueSnafu.into_error(
<a href=#286 id=286 data-nosnippet>286</a> session::session_config::Error::InvalidConfigValue {
<a href=#287 id=287 data-nosnippet>287</a> name: variable.to_string(),
<a href=#288 id=288 data-nosnippet>288</a> value: value.value.to_string(),
<a href=#289 id=289 data-nosnippet>289</a> hint: <span class="string">"Expected a boolean, number, or quoted string literal"</span>.to_string(),
<a href=#290 id=290 data-nosnippet>290</a> location: snafu::Location::new(<span class="macro">file!</span>(), <span class="macro">line!</span>(), <span class="macro">column!</span>()),
<a href=#291 id=291 data-nosnippet>291</a> },
<a href=#292 id=292 data-nosnippet>292</a> ));
<a href=#293 id=293 data-nosnippet>293</a> }
<a href=#294 id=294 data-nosnippet>294</a> };
<a href=#295 id=295 data-nosnippet>295</a> ctx.configuration_parameter()
<a href=#296 id=296 data-nosnippet>296</a> .set_query_option(variable, <span class="kw-2">&amp;</span>value)
<a href=#297 id=297 data-nosnippet>297</a> .context(InvalidConfigValueSnafu)<span class="question-mark">?</span>;
<a href=#298 id=298 data-nosnippet>298</a> <span class="prelude-val">Ok</span>(<span class="bool-val">true</span>)
<a href=#299 id=299 data-nosnippet>299</a>}
<a href=#300 id=300 data-nosnippet>300</a>
<a href=#301 id=301 data-nosnippet>301</a><span class="kw">pub fn </span>set_intervalstyle(exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#302 id=302 data-nosnippet>302</a> <span class="kw">let </span><span class="prelude-val">Some</span>((var_value, [])) = exprs.split_first() <span class="kw">else </span>{
<a href=#303 id=303 data-nosnippet>303</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#304 id=304 data-nosnippet>304</a> feat: <span class="string">"Set variable value must have one and only one value for intervalstyle"</span>,
<a href=#305 id=305 data-nosnippet>305</a> }
<a href=#306 id=306 data-nosnippet>306</a> .fail();
<a href=#307 id=307 data-nosnippet>307</a> };
<a href=#308 id=308 data-nosnippet>308</a> <span class="kw">let </span>intervalstyle = <span class="kw">match </span>var_value {
<a href=#309 id=309 data-nosnippet>309</a> Expr::Identifier(Ident {
<a href=#310 id=310 data-nosnippet>310</a> value,
<a href=#311 id=311 data-nosnippet>311</a> quote_style: <span class="kw">_</span>,
<a href=#312 id=312 data-nosnippet>312</a> span: <span class="kw">_</span>,
<a href=#313 id=313 data-nosnippet>313</a> }) =&gt; PGIntervalStyle::try_from(value.as_str()).context(InvalidConfigValueSnafu)<span class="question-mark">?</span>,
<a href=#314 id=314 data-nosnippet>314</a> Expr::Value(value) =&gt; {
<a href=#315 id=315 data-nosnippet>315</a> PGIntervalStyle::try_from(<span class="kw-2">&amp;</span>value.value).context(InvalidConfigValueSnafu)<span class="question-mark">?
<a href=#316 id=316 data-nosnippet>316</a> </span>}
<a href=#317 id=317 data-nosnippet>317</a> <span class="kw">_ </span>=&gt; {
<a href=#318 id=318 data-nosnippet>318</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#319 id=319 data-nosnippet>319</a> feat: <span class="string">"Set variable value must be a value or identifier"</span>,
<a href=#320 id=320 data-nosnippet>320</a> }
<a href=#321 id=321 data-nosnippet>321</a> .fail();
<a href=#322 id=322 data-nosnippet>322</a> }
<a href=#323 id=323 data-nosnippet>323</a> };
<a href=#324 id=324 data-nosnippet>324</a> ctx.configuration_parameter()
<a href=#325 id=325 data-nosnippet>325</a> .set_pg_intervalstyle_format(intervalstyle);
<a href=#326 id=326 data-nosnippet>326</a> <span class="prelude-val">Ok</span>(())
<a href=#327 id=327 data-nosnippet>327</a>}
<a href=#328 id=328 data-nosnippet>328</a>
<a href=#329 id=329 data-nosnippet>329</a><span class="kw">pub fn </span>set_datestyle(exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#330 id=330 data-nosnippet>330</a> <span class="comment">// ORDER,
<a href=#331 id=331 data-nosnippet>331</a> // STYLE,
<a href=#332 id=332 data-nosnippet>332</a> // ORDER,ORDER
<a href=#333 id=333 data-nosnippet>333</a> // ORDER,STYLE
<a href=#334 id=334 data-nosnippet>334</a> // STYLE,ORDER
<a href=#335 id=335 data-nosnippet>335</a> </span><span class="kw">let </span>(style, order) = exprs
<a href=#336 id=336 data-nosnippet>336</a> .iter()
<a href=#337 id=337 data-nosnippet>337</a> .try_fold((<span class="prelude-val">None</span>, <span class="prelude-val">None</span>), |(style, order), expr| {
<a href=#338 id=338 data-nosnippet>338</a> <span class="kw">let </span>(new_style, new_order) = try_parse_datestyle(expr)<span class="question-mark">?</span>;
<a href=#339 id=339 data-nosnippet>339</a> <span class="prelude-val">Ok</span>((
<a href=#340 id=340 data-nosnippet>340</a> merge_datestyle_value(style, new_style)<span class="question-mark">?</span>,
<a href=#341 id=341 data-nosnippet>341</a> merge_datestyle_value(order, new_order)<span class="question-mark">?</span>,
<a href=#342 id=342 data-nosnippet>342</a> ))
<a href=#343 id=343 data-nosnippet>343</a> })<span class="question-mark">?</span>;
<a href=#344 id=344 data-nosnippet>344</a>
<a href=#345 id=345 data-nosnippet>345</a> <span class="kw">let </span>(old_style, older_order) = <span class="kw-2">*</span>ctx.configuration_parameter().pg_datetime_style();
<a href=#346 id=346 data-nosnippet>346</a> ctx.configuration_parameter()
<a href=#347 id=347 data-nosnippet>347</a> .set_pg_datetime_style(style.unwrap_or(old_style), order.unwrap_or(older_order));
<a href=#348 id=348 data-nosnippet>348</a> <span class="prelude-val">Ok</span>(())
<a href=#349 id=349 data-nosnippet>349</a>}
<a href=#350 id=350 data-nosnippet>350</a>
<a href=#351 id=351 data-nosnippet>351</a><span class="kw">pub fn </span>set_query_timeout(exprs: Vec&lt;Expr&gt;, ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;()&gt; {
<a href=#352 id=352 data-nosnippet>352</a> <span class="kw">let </span>timeout_expr = exprs.first().context(NotSupportedSnafu {
<a href=#353 id=353 data-nosnippet>353</a> feat: <span class="string">"No timeout value find in set query timeout statement"</span>,
<a href=#354 id=354 data-nosnippet>354</a> })<span class="question-mark">?</span>;
<a href=#355 id=355 data-nosnippet>355</a> <span class="kw">match </span>timeout_expr {
<a href=#356 id=356 data-nosnippet>356</a> Expr::Value(ValueWithSpan {
<a href=#357 id=357 data-nosnippet>357</a> value: Value::Number(timeout, <span class="kw">_</span>),
<a href=#358 id=358 data-nosnippet>358</a> ..
<a href=#359 id=359 data-nosnippet>359</a> }) =&gt; {
<a href=#360 id=360 data-nosnippet>360</a> <span class="kw">match </span>timeout.parse::&lt;u64&gt;() {
<a href=#361 id=361 data-nosnippet>361</a> <span class="prelude-val">Ok</span>(timeout) =&gt; ctx.set_query_timeout(Duration::from_millis(timeout)),
<a href=#362 id=362 data-nosnippet>362</a> <span class="prelude-val">Err</span>(<span class="kw">_</span>) =&gt; {
<a href=#363 id=363 data-nosnippet>363</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#364 id=364 data-nosnippet>364</a> feat: <span class="macro">format!</span>(<span class="string">"Invalid timeout expr {} in set variable statement"</span>, timeout),
<a href=#365 id=365 data-nosnippet>365</a> }
<a href=#366 id=366 data-nosnippet>366</a> .fail();
<a href=#367 id=367 data-nosnippet>367</a> }
<a href=#368 id=368 data-nosnippet>368</a> }
<a href=#369 id=369 data-nosnippet>369</a> <span class="prelude-val">Ok</span>(())
<a href=#370 id=370 data-nosnippet>370</a> }
<a href=#371 id=371 data-nosnippet>371</a> <span class="comment">// postgres support time units i.e. SET STATEMENT_TIMEOUT = '50ms';
<a href=#372 id=372 data-nosnippet>372</a> </span>Expr::Value(ValueWithSpan {
<a href=#373 id=373 data-nosnippet>373</a> value: Value::SingleQuotedString(timeout),
<a href=#374 id=374 data-nosnippet>374</a> ..
<a href=#375 id=375 data-nosnippet>375</a> })
<a href=#376 id=376 data-nosnippet>376</a> | Expr::Value(ValueWithSpan {
<a href=#377 id=377 data-nosnippet>377</a> value: Value::DoubleQuotedString(timeout),
<a href=#378 id=378 data-nosnippet>378</a> ..
<a href=#379 id=379 data-nosnippet>379</a> }) =&gt; {
<a href=#380 id=380 data-nosnippet>380</a> <span class="kw">if </span>ctx.channel() != Postgres {
<a href=#381 id=381 data-nosnippet>381</a> <span class="kw">return </span>NotSupportedSnafu {
<a href=#382 id=382 data-nosnippet>382</a> feat: <span class="macro">format!</span>(<span class="string">"Invalid timeout expr {} in set variable statement"</span>, timeout),
<a href=#383 id=383 data-nosnippet>383</a> }
<a href=#384 id=384 data-nosnippet>384</a> .fail();
<a href=#385 id=385 data-nosnippet>385</a> }
<a href=#386 id=386 data-nosnippet>386</a> <span class="kw">let </span>timeout = parse_pg_query_timeout_input(timeout)<span class="question-mark">?</span>;
<a href=#387 id=387 data-nosnippet>387</a> ctx.set_query_timeout(Duration::from_millis(timeout));
<a href=#388 id=388 data-nosnippet>388</a> <span class="prelude-val">Ok</span>(())
<a href=#389 id=389 data-nosnippet>389</a> }
<a href=#390 id=390 data-nosnippet>390</a> expr =&gt; NotSupportedSnafu {
<a href=#391 id=391 data-nosnippet>391</a> feat: <span class="macro">format!</span>(
<a href=#392 id=392 data-nosnippet>392</a> <span class="string">"Unsupported timeout expr {} in set variable statement"</span>,
<a href=#393 id=393 data-nosnippet>393</a> expr
<a href=#394 id=394 data-nosnippet>394</a> ),
<a href=#395 id=395 data-nosnippet>395</a> }
<a href=#396 id=396 data-nosnippet>396</a> .fail(),
<a href=#397 id=397 data-nosnippet>397</a> }
<a href=#398 id=398 data-nosnippet>398</a>}
<a href=#399 id=399 data-nosnippet>399</a>
<a href=#400 id=400 data-nosnippet>400</a><span class="comment">// support time units in ms, s, min, h, d for postgres protocol.
<a href=#401 id=401 data-nosnippet>401</a>// https://www.postgresql.org/docs/8.4/config-setting.html#:~:text=Valid%20memory%20units%20are%20kB,%2C%20and%20d%20(days).
<a href=#402 id=402 data-nosnippet>402</a></span><span class="kw">fn </span>parse_pg_query_timeout_input(input: <span class="kw-2">&amp;</span>str) -&gt; <span class="prelude-ty">Result</span>&lt;u64&gt; {
<a href=#403 id=403 data-nosnippet>403</a> <span class="kw">match </span>input.parse::&lt;u64&gt;() {
<a href=#404 id=404 data-nosnippet>404</a> <span class="prelude-val">Ok</span>(timeout) =&gt; <span class="prelude-val">Ok</span>(timeout),
<a href=#405 id=405 data-nosnippet>405</a> <span class="prelude-val">Err</span>(<span class="kw">_</span>) =&gt; {
<a href=#406 id=406 data-nosnippet>406</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(captures) = PG_TIME_INPUT_REGEX.captures(input) {
<a href=#407 id=407 data-nosnippet>407</a> <span class="kw">let </span>value = captures[<span class="number">1</span>].parse::&lt;u64&gt;().expect(<span class="string">"regex failed"</span>);
<a href=#408 id=408 data-nosnippet>408</a> <span class="kw">let </span>unit = <span class="kw-2">&amp;</span>captures[<span class="number">2</span>];
<a href=#409 id=409 data-nosnippet>409</a>
<a href=#410 id=410 data-nosnippet>410</a> <span class="kw">match </span>unit {
<a href=#411 id=411 data-nosnippet>411</a> <span class="string">"ms" </span>=&gt; <span class="prelude-val">Ok</span>(value),
<a href=#412 id=412 data-nosnippet>412</a> <span class="string">"s" </span>=&gt; <span class="prelude-val">Ok</span>(value * <span class="number">1000</span>),
<a href=#413 id=413 data-nosnippet>413</a> <span class="string">"min" </span>=&gt; <span class="prelude-val">Ok</span>(value * <span class="number">60 </span>* <span class="number">1000</span>),
<a href=#414 id=414 data-nosnippet>414</a> <span class="string">"h" </span>=&gt; <span class="prelude-val">Ok</span>(value * <span class="number">60 </span>* <span class="number">60 </span>* <span class="number">1000</span>),
<a href=#415 id=415 data-nosnippet>415</a> <span class="string">"d" </span>=&gt; <span class="prelude-val">Ok</span>(value * <span class="number">24 </span>* <span class="number">60 </span>* <span class="number">60 </span>* <span class="number">1000</span>),
<a href=#416 id=416 data-nosnippet>416</a> <span class="kw">_ </span>=&gt; <span class="macro">unreachable!</span>(<span class="string">"regex failed"</span>),
<a href=#417 id=417 data-nosnippet>417</a> }
<a href=#418 id=418 data-nosnippet>418</a> } <span class="kw">else </span>{
<a href=#419 id=419 data-nosnippet>419</a> NotSupportedSnafu {
<a href=#420 id=420 data-nosnippet>420</a> feat: <span class="macro">format!</span>(
<a href=#421 id=421 data-nosnippet>421</a> <span class="string">"Unsupported timeout expr {} in set variable statement"</span>,
<a href=#422 id=422 data-nosnippet>422</a> input
<a href=#423 id=423 data-nosnippet>423</a> ),
<a href=#424 id=424 data-nosnippet>424</a> }
<a href=#425 id=425 data-nosnippet>425</a> .fail()
<a href=#426 id=426 data-nosnippet>426</a> }
<a href=#427 id=427 data-nosnippet>427</a> }
<a href=#428 id=428 data-nosnippet>428</a> }
<a href=#429 id=429 data-nosnippet>429</a>}
<a href=#430 id=430 data-nosnippet>430</a>
<a href=#431 id=431 data-nosnippet>431</a><span class="attr">#[cfg(test)]
<a href=#432 id=432 data-nosnippet>432</a></span><span class="kw">mod </span>test {
<a href=#433 id=433 data-nosnippet>433</a> <span class="kw">use </span>session::Session;
<a href=#434 id=434 data-nosnippet>434</a> <span class="kw">use </span>session::context::{Channel, QueryContext};
<a href=#435 id=435 data-nosnippet>435</a> <span class="kw">use </span>sql::ast::{Expr, Value};
<a href=#436 id=436 data-nosnippet>436</a>
<a href=#437 id=437 data-nosnippet>437</a> <span class="kw">use </span><span class="kw">super</span>::set_skip_wal;
<a href=#438 id=438 data-nosnippet>438</a> <span class="kw">use </span><span class="kw">crate</span>::statement::set::parse_pg_query_timeout_input;
<a href=#439 id=439 data-nosnippet>439</a>
<a href=#440 id=440 data-nosnippet>440</a> <span class="attr">#[test]
<a href=#441 id=441 data-nosnippet>441</a> </span><span class="kw">fn </span>test_set_skip_wal() {
<a href=#442 id=442 data-nosnippet>442</a> <span class="kw">let </span>ctx = QueryContext::arc();
<a href=#443 id=443 data-nosnippet>443</a> <span class="kw">for </span>value <span class="kw">in </span>[<span class="bool-val">true</span>, <span class="bool-val">false</span>] {
<a href=#444 id=444 data-nosnippet>444</a> set_skip_wal(<span class="macro">vec!</span>[Expr::Value(Value::Boolean(value).into())], ctx.clone()).unwrap();
<a href=#445 id=445 data-nosnippet>445</a> <span class="macro">assert_eq!</span>(ctx.skip_wal(), value);
<a href=#446 id=446 data-nosnippet>446</a> set_skip_wal(
<a href=#447 id=447 data-nosnippet>447</a> <span class="macro">vec!</span>[Expr::Value(
<a href=#448 id=448 data-nosnippet>448</a> Value::SingleQuotedString(value.to_string()).into(),
<a href=#449 id=449 data-nosnippet>449</a> )],
<a href=#450 id=450 data-nosnippet>450</a> ctx.clone(),
<a href=#451 id=451 data-nosnippet>451</a> )
<a href=#452 id=452 data-nosnippet>452</a> .unwrap();
<a href=#453 id=453 data-nosnippet>453</a> <span class="macro">assert_eq!</span>(ctx.skip_wal(), value);
<a href=#454 id=454 data-nosnippet>454</a> }
<a href=#455 id=455 data-nosnippet>455</a> ctx.set_skip_wal(<span class="bool-val">true</span>);
<a href=#456 id=456 data-nosnippet>456</a> <span class="kw">for </span>values <span class="kw">in </span>[
<a href=#457 id=457 data-nosnippet>457</a> <span class="macro">vec!</span>[],
<a href=#458 id=458 data-nosnippet>458</a> <span class="macro">vec!</span>[Expr::Value(Value::Number(<span class="string">"1"</span>.to_string(), <span class="bool-val">false</span>).into())],
<a href=#459 id=459 data-nosnippet>459</a> <span class="macro">vec!</span>[Expr::Value(
<a href=#460 id=460 data-nosnippet>460</a> Value::SingleQuotedString(<span class="string">"invalid"</span>.to_string()).into(),
<a href=#461 id=461 data-nosnippet>461</a> )],
<a href=#462 id=462 data-nosnippet>462</a> <span class="macro">vec!</span>[Expr::Value(Value::Boolean(<span class="bool-val">false</span>).into()); <span class="number">2</span>],
<a href=#463 id=463 data-nosnippet>463</a> ] {
<a href=#464 id=464 data-nosnippet>464</a> <span class="macro">assert!</span>(set_skip_wal(values, ctx.clone()).is_err());
<a href=#465 id=465 data-nosnippet>465</a> <span class="macro">assert!</span>(ctx.skip_wal());
<a href=#466 id=466 data-nosnippet>466</a> }
<a href=#467 id=467 data-nosnippet>467</a> }
<a href=#468 id=468 data-nosnippet>468</a>
<a href=#469 id=469 data-nosnippet>469</a> <span class="attr">#[test]
<a href=#470 id=470 data-nosnippet>470</a> </span><span class="kw">fn </span>test_set_skip_wal_session_isolation() {
<a href=#471 id=471 data-nosnippet>471</a> <span class="kw">for </span>channel <span class="kw">in </span>[Channel::Mysql, Channel::Postgres] {
<a href=#472 id=472 data-nosnippet>472</a> <span class="kw">let </span>session = Session::new(<span class="prelude-val">None</span>, channel, Default::default(), <span class="number">0</span>);
<a href=#473 id=473 data-nosnippet>473</a> <span class="kw">let </span>other = Session::new(<span class="prelude-val">None</span>, channel, Default::default(), <span class="number">1</span>);
<a href=#474 id=474 data-nosnippet>474</a> <span class="macro">assert!</span>(!session.new_query_context().skip_wal());
<a href=#475 id=475 data-nosnippet>475</a> set_skip_wal(
<a href=#476 id=476 data-nosnippet>476</a> <span class="macro">vec!</span>[Expr::Value(Value::Boolean(<span class="bool-val">true</span>).into())],
<a href=#477 id=477 data-nosnippet>477</a> session.new_query_context(),
<a href=#478 id=478 data-nosnippet>478</a> )
<a href=#479 id=479 data-nosnippet>479</a> .unwrap();
<a href=#480 id=480 data-nosnippet>480</a> <span class="macro">assert!</span>(session.new_query_context().skip_wal());
<a href=#481 id=481 data-nosnippet>481</a> <span class="macro">assert!</span>(!other.new_query_context().skip_wal());
<a href=#482 id=482 data-nosnippet>482</a> }
<a href=#483 id=483 data-nosnippet>483</a> }
<a href=#484 id=484 data-nosnippet>484</a>
<a href=#485 id=485 data-nosnippet>485</a> <span class="attr">#[test]
<a href=#486 id=486 data-nosnippet>486</a> </span><span class="kw">fn </span>test_parse_pg_query_timeout_input() {
<a href=#487 id=487 data-nosnippet>487</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">""</span>).is_err());
<a href=#488 id=488 data-nosnippet>488</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">" 50 ms"</span>).is_err());
<a href=#489 id=489 data-nosnippet>489</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">"5s 1ms"</span>).is_err());
<a href=#490 id=490 data-nosnippet>490</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">"3a"</span>).is_err());
<a href=#491 id=491 data-nosnippet>491</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">"1.5min"</span>).is_err());
<a href=#492 id=492 data-nosnippet>492</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">"ms"</span>).is_err());
<a href=#493 id=493 data-nosnippet>493</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">"a"</span>).is_err());
<a href=#494 id=494 data-nosnippet>494</a> <span class="macro">assert!</span>(parse_pg_query_timeout_input(<span class="string">"-1"</span>).is_err());
<a href=#495 id=495 data-nosnippet>495</a>
<a href=#496 id=496 data-nosnippet>496</a> <span class="macro">assert_eq!</span>(<span class="number">50</span>, parse_pg_query_timeout_input(<span class="string">"50"</span>).unwrap());
<a href=#497 id=497 data-nosnippet>497</a> <span class="macro">assert_eq!</span>(<span class="number">12</span>, parse_pg_query_timeout_input(<span class="string">"12ms"</span>).unwrap());
<a href=#498 id=498 data-nosnippet>498</a> <span class="macro">assert_eq!</span>(<span class="number">2000</span>, parse_pg_query_timeout_input(<span class="string">"2s"</span>).unwrap());
<a href=#499 id=499 data-nosnippet>499</a> <span class="macro">assert_eq!</span>(<span class="number">60000</span>, parse_pg_query_timeout_input(<span class="string">"1min"</span>).unwrap());
<a href=#500 id=500 data-nosnippet>500</a> }
<a href=#501 id=501 data-nosnippet>501</a>}
</code></pre></div></section></main></body></html>
+432
View File
@@ -0,0 +1,432 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/show.rs`."><title>show.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>show.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::sync::Arc;
<a href=#16 id=16 data-nosnippet>16</a>
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_error::ext::BoxedError;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_meta::key::schema_name::SchemaNameKey;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>common_query::Output;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>common_telemetry::tracing;
<a href=#21 id=21 data-nosnippet>21</a><span class="kw">use </span>partition::manager::PartitionInfoWithVersion;
<a href=#22 id=22 data-nosnippet>22</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#23 id=23 data-nosnippet>23</a><span class="kw">use </span>session::table_name::table_idents_to_full_name;
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>snafu::{OptionExt, ResultExt};
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>sql::ast::ObjectNamePartExt;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>sql::statements::OptionMap;
<a href=#27 id=27 data-nosnippet>27</a><span class="kw">use </span>sql::statements::create::Partitions;
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span>sql::statements::show::{
<a href=#29 id=29 data-nosnippet>29</a> ShowColumns, ShowCreateFlow, ShowCreateView, ShowDatabases, ShowFlowStatus, ShowFlows,
<a href=#30 id=30 data-nosnippet>30</a> ShowIndex, ShowKind, ShowProcessList, ShowRegion, ShowTableStatus, ShowTables, ShowVariables,
<a href=#31 id=31 data-nosnippet>31</a> ShowViews,
<a href=#32 id=32 data-nosnippet>32</a>};
<a href=#33 id=33 data-nosnippet>33</a><span class="kw">use </span>table::TableRef;
<a href=#34 id=34 data-nosnippet>34</a><span class="kw">use </span>table::metadata::{TableInfo, TableType};
<a href=#35 id=35 data-nosnippet>35</a><span class="kw">use </span>table::table_name::TableName;
<a href=#36 id=36 data-nosnippet>36</a>
<a href=#37 id=37 data-nosnippet>37</a><span class="kw">use </span><span class="kw">crate</span>::error::{
<a href=#38 id=38 data-nosnippet>38</a> <span class="self">self</span>, CatalogSnafu, ExecLogicalPlanSnafu, ExecuteStatementSnafu, ExternalSnafu,
<a href=#39 id=39 data-nosnippet>39</a> FindViewInfoSnafu, InvalidSqlSnafu, <span class="prelude-ty">Result</span>, TableMetadataManagerSnafu, ViewInfoNotFoundSnafu,
<a href=#40 id=40 data-nosnippet>40</a> ViewNotFoundSnafu,
<a href=#41 id=41 data-nosnippet>41</a>};
<a href=#42 id=42 data-nosnippet>42</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#43 id=43 data-nosnippet>43</a>
<a href=#44 id=44 data-nosnippet>44</a><span class="kw">impl </span>StatementExecutor {
<a href=#45 id=45 data-nosnippet>45</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#46 id=46 data-nosnippet>46</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_databases(
<a href=#47 id=47 data-nosnippet>47</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#48 id=48 data-nosnippet>48</a> stmt: ShowDatabases,
<a href=#49 id=49 data-nosnippet>49</a> query_ctx: QueryContextRef,
<a href=#50 id=50 data-nosnippet>50</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#51 id=51 data-nosnippet>51</a> query::sql::show_databases(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#52 id=52 data-nosnippet>52</a> .<span class="kw">await
<a href=#53 id=53 data-nosnippet>53</a> </span>.context(ExecuteStatementSnafu)
<a href=#54 id=54 data-nosnippet>54</a> }
<a href=#55 id=55 data-nosnippet>55</a>
<a href=#56 id=56 data-nosnippet>56</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#57 id=57 data-nosnippet>57</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_tables(
<a href=#58 id=58 data-nosnippet>58</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#59 id=59 data-nosnippet>59</a> stmt: ShowTables,
<a href=#60 id=60 data-nosnippet>60</a> query_ctx: QueryContextRef,
<a href=#61 id=61 data-nosnippet>61</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#62 id=62 data-nosnippet>62</a> query::sql::show_tables(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#63 id=63 data-nosnippet>63</a> .<span class="kw">await
<a href=#64 id=64 data-nosnippet>64</a> </span>.context(ExecuteStatementSnafu)
<a href=#65 id=65 data-nosnippet>65</a> }
<a href=#66 id=66 data-nosnippet>66</a>
<a href=#67 id=67 data-nosnippet>67</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#68 id=68 data-nosnippet>68</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_table_status(
<a href=#69 id=69 data-nosnippet>69</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#70 id=70 data-nosnippet>70</a> stmt: ShowTableStatus,
<a href=#71 id=71 data-nosnippet>71</a> query_ctx: QueryContextRef,
<a href=#72 id=72 data-nosnippet>72</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#73 id=73 data-nosnippet>73</a> query::sql::show_table_status(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#74 id=74 data-nosnippet>74</a> .<span class="kw">await
<a href=#75 id=75 data-nosnippet>75</a> </span>.context(ExecuteStatementSnafu)
<a href=#76 id=76 data-nosnippet>76</a> }
<a href=#77 id=77 data-nosnippet>77</a>
<a href=#78 id=78 data-nosnippet>78</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#79 id=79 data-nosnippet>79</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_columns(
<a href=#80 id=80 data-nosnippet>80</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#81 id=81 data-nosnippet>81</a> stmt: ShowColumns,
<a href=#82 id=82 data-nosnippet>82</a> query_ctx: QueryContextRef,
<a href=#83 id=83 data-nosnippet>83</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#84 id=84 data-nosnippet>84</a> query::sql::show_columns(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#85 id=85 data-nosnippet>85</a> .<span class="kw">await
<a href=#86 id=86 data-nosnippet>86</a> </span>.context(ExecuteStatementSnafu)
<a href=#87 id=87 data-nosnippet>87</a> }
<a href=#88 id=88 data-nosnippet>88</a>
<a href=#89 id=89 data-nosnippet>89</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#90 id=90 data-nosnippet>90</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_index(
<a href=#91 id=91 data-nosnippet>91</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#92 id=92 data-nosnippet>92</a> stmt: ShowIndex,
<a href=#93 id=93 data-nosnippet>93</a> query_ctx: QueryContextRef,
<a href=#94 id=94 data-nosnippet>94</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#95 id=95 data-nosnippet>95</a> query::sql::show_index(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#96 id=96 data-nosnippet>96</a> .<span class="kw">await
<a href=#97 id=97 data-nosnippet>97</a> </span>.context(ExecuteStatementSnafu)
<a href=#98 id=98 data-nosnippet>98</a> }
<a href=#99 id=99 data-nosnippet>99</a>
<a href=#100 id=100 data-nosnippet>100</a> <span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_region(
<a href=#101 id=101 data-nosnippet>101</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#102 id=102 data-nosnippet>102</a> stmt: ShowRegion,
<a href=#103 id=103 data-nosnippet>103</a> query_ctx: QueryContextRef,
<a href=#104 id=104 data-nosnippet>104</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#105 id=105 data-nosnippet>105</a> query::sql::show_region(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#106 id=106 data-nosnippet>106</a> .<span class="kw">await
<a href=#107 id=107 data-nosnippet>107</a> </span>.context(ExecuteStatementSnafu)
<a href=#108 id=108 data-nosnippet>108</a> }
<a href=#109 id=109 data-nosnippet>109</a>
<a href=#110 id=110 data-nosnippet>110</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#111 id=111 data-nosnippet>111</a> </span><span class="kw">pub async fn </span>show_create_database(
<a href=#112 id=112 data-nosnippet>112</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#113 id=113 data-nosnippet>113</a> database_name: <span class="kw-2">&amp;</span>str,
<a href=#114 id=114 data-nosnippet>114</a> opts: OptionMap,
<a href=#115 id=115 data-nosnippet>115</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#116 id=116 data-nosnippet>116</a> query::sql::show_create_database(database_name, opts).context(ExecuteStatementSnafu)
<a href=#117 id=117 data-nosnippet>117</a> }
<a href=#118 id=118 data-nosnippet>118</a>
<a href=#119 id=119 data-nosnippet>119</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#120 id=120 data-nosnippet>120</a> </span><span class="kw">pub async fn </span>show_create_table(
<a href=#121 id=121 data-nosnippet>121</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#122 id=122 data-nosnippet>122</a> table_name: TableName,
<a href=#123 id=123 data-nosnippet>123</a> table: TableRef,
<a href=#124 id=124 data-nosnippet>124</a> query_ctx: QueryContextRef,
<a href=#125 id=125 data-nosnippet>125</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#126 id=126 data-nosnippet>126</a> <span class="kw">let </span><span class="kw-2">mut </span>table_info = table.table_info();
<a href=#127 id=127 data-nosnippet>127</a> <span class="kw">let </span>partition_column_names: Vec&lt;<span class="kw">_</span>&gt; =
<a href=#128 id=128 data-nosnippet>128</a> table_info.meta.partition_column_names().cloned().collect();
<a href=#129 id=129 data-nosnippet>129</a>
<a href=#130 id=130 data-nosnippet>130</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(latest) = <span class="self">self
<a href=#131 id=131 data-nosnippet>131</a> </span>.table_metadata_manager
<a href=#132 id=132 data-nosnippet>132</a> .table_info_manager()
<a href=#133 id=133 data-nosnippet>133</a> .get(table_info.table_id())
<a href=#134 id=134 data-nosnippet>134</a> .<span class="kw">await
<a href=#135 id=135 data-nosnippet>135</a> </span>.context(TableMetadataManagerSnafu)<span class="question-mark">?
<a href=#136 id=136 data-nosnippet>136</a> </span>{
<a href=#137 id=137 data-nosnippet>137</a> <span class="kw">let </span><span class="kw-2">mut </span>latest_info = latest.into_inner().table_info;
<a href=#138 id=138 data-nosnippet>138</a>
<a href=#139 id=139 data-nosnippet>139</a> <span class="kw">if </span>!partition_column_names.is_empty() {
<a href=#140 id=140 data-nosnippet>140</a> latest_info.meta.partition_key_indices = partition_column_names
<a href=#141 id=141 data-nosnippet>141</a> .iter()
<a href=#142 id=142 data-nosnippet>142</a> .filter_map(|name| latest_info.meta.schema.column_index_by_name(name.as_str()))
<a href=#143 id=143 data-nosnippet>143</a> .collect();
<a href=#144 id=144 data-nosnippet>144</a> }
<a href=#145 id=145 data-nosnippet>145</a>
<a href=#146 id=146 data-nosnippet>146</a> table_info = Arc::new(latest_info);
<a href=#147 id=147 data-nosnippet>147</a> }
<a href=#148 id=148 data-nosnippet>148</a>
<a href=#149 id=149 data-nosnippet>149</a> <span class="kw">if </span>table_info.table_type != TableType::Base {
<a href=#150 id=150 data-nosnippet>150</a> <span class="kw">return </span>error::ShowCreateTableBaseOnlySnafu {
<a href=#151 id=151 data-nosnippet>151</a> table_name: table_name.to_string(),
<a href=#152 id=152 data-nosnippet>152</a> table_type: table_info.table_type,
<a href=#153 id=153 data-nosnippet>153</a> }
<a href=#154 id=154 data-nosnippet>154</a> .fail();
<a href=#155 id=155 data-nosnippet>155</a> }
<a href=#156 id=156 data-nosnippet>156</a>
<a href=#157 id=157 data-nosnippet>157</a> <span class="kw">let </span>schema_options = <span class="self">self
<a href=#158 id=158 data-nosnippet>158</a> </span>.table_metadata_manager
<a href=#159 id=159 data-nosnippet>159</a> .schema_manager()
<a href=#160 id=160 data-nosnippet>160</a> .get(SchemaNameKey {
<a href=#161 id=161 data-nosnippet>161</a> catalog: <span class="kw-2">&amp;</span>table_name.catalog_name,
<a href=#162 id=162 data-nosnippet>162</a> schema: <span class="kw-2">&amp;</span>table_name.schema_name,
<a href=#163 id=163 data-nosnippet>163</a> })
<a href=#164 id=164 data-nosnippet>164</a> .<span class="kw">await
<a href=#165 id=165 data-nosnippet>165</a> </span>.context(TableMetadataManagerSnafu)<span class="question-mark">?
<a href=#166 id=166 data-nosnippet>166</a> </span>.map(|v| v.into_inner());
<a href=#167 id=167 data-nosnippet>167</a>
<a href=#168 id=168 data-nosnippet>168</a> <span class="kw">let </span>partition_info = <span class="self">self
<a href=#169 id=169 data-nosnippet>169</a> </span>.partition_manager
<a href=#170 id=170 data-nosnippet>170</a> .find_physical_partition_info(table_info.table_id())
<a href=#171 id=171 data-nosnippet>171</a> .<span class="kw">await
<a href=#172 id=172 data-nosnippet>172</a> </span>.context(error::FindTablePartitionRuleSnafu {
<a href=#173 id=173 data-nosnippet>173</a> table_name: <span class="kw-2">&amp;</span>table_name.table_name,
<a href=#174 id=174 data-nosnippet>174</a> })<span class="question-mark">?</span>;
<a href=#175 id=175 data-nosnippet>175</a>
<a href=#176 id=176 data-nosnippet>176</a> <span class="kw">let </span>partitions = create_partitions_stmt(<span class="kw-2">&amp;</span>table_info, <span class="kw-2">&amp;</span>partition_info.partitions)<span class="question-mark">?</span>;
<a href=#177 id=177 data-nosnippet>177</a>
<a href=#178 id=178 data-nosnippet>178</a> query::sql::show_create_table(table_info, schema_options, partitions, query_ctx)
<a href=#179 id=179 data-nosnippet>179</a> .context(ExecuteStatementSnafu)
<a href=#180 id=180 data-nosnippet>180</a> }
<a href=#181 id=181 data-nosnippet>181</a>
<a href=#182 id=182 data-nosnippet>182</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#183 id=183 data-nosnippet>183</a> </span><span class="kw">pub async fn </span>show_create_table_for_pg(
<a href=#184 id=184 data-nosnippet>184</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#185 id=185 data-nosnippet>185</a> table_name: TableName,
<a href=#186 id=186 data-nosnippet>186</a> table: TableRef,
<a href=#187 id=187 data-nosnippet>187</a> query_ctx: QueryContextRef,
<a href=#188 id=188 data-nosnippet>188</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#189 id=189 data-nosnippet>189</a> <span class="kw">let </span>table_info = table.table_info();
<a href=#190 id=190 data-nosnippet>190</a> <span class="kw">if </span>table_info.table_type != TableType::Base {
<a href=#191 id=191 data-nosnippet>191</a> <span class="kw">return </span>error::ShowCreateTableBaseOnlySnafu {
<a href=#192 id=192 data-nosnippet>192</a> table_name: table_name.to_string(),
<a href=#193 id=193 data-nosnippet>193</a> table_type: table_info.table_type,
<a href=#194 id=194 data-nosnippet>194</a> }
<a href=#195 id=195 data-nosnippet>195</a> .fail();
<a href=#196 id=196 data-nosnippet>196</a> }
<a href=#197 id=197 data-nosnippet>197</a>
<a href=#198 id=198 data-nosnippet>198</a> query::sql::show_create_foreign_table_for_pg(table, query_ctx)
<a href=#199 id=199 data-nosnippet>199</a> .context(ExecuteStatementSnafu)
<a href=#200 id=200 data-nosnippet>200</a> }
<a href=#201 id=201 data-nosnippet>201</a>
<a href=#202 id=202 data-nosnippet>202</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#203 id=203 data-nosnippet>203</a> </span><span class="kw">pub async fn </span>show_create_view(
<a href=#204 id=204 data-nosnippet>204</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#205 id=205 data-nosnippet>205</a> show: ShowCreateView,
<a href=#206 id=206 data-nosnippet>206</a> query_ctx: QueryContextRef,
<a href=#207 id=207 data-nosnippet>207</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#208 id=208 data-nosnippet>208</a> <span class="kw">let </span>(catalog, schema, view) = table_idents_to_full_name(<span class="kw-2">&amp;</span>show.view_name, <span class="kw-2">&amp;</span>query_ctx)
<a href=#209 id=209 data-nosnippet>209</a> .map_err(BoxedError::new)
<a href=#210 id=210 data-nosnippet>210</a> .context(ExternalSnafu)<span class="question-mark">?</span>;
<a href=#211 id=211 data-nosnippet>211</a>
<a href=#212 id=212 data-nosnippet>212</a> <span class="kw">let </span>table_ref = <span class="self">self
<a href=#213 id=213 data-nosnippet>213</a> </span>.catalog_manager
<a href=#214 id=214 data-nosnippet>214</a> .table(<span class="kw-2">&amp;</span>catalog, <span class="kw-2">&amp;</span>schema, <span class="kw-2">&amp;</span>view, <span class="prelude-val">Some</span>(<span class="kw-2">&amp;</span>query_ctx))
<a href=#215 id=215 data-nosnippet>215</a> .<span class="kw">await
<a href=#216 id=216 data-nosnippet>216</a> </span>.context(CatalogSnafu)<span class="question-mark">?
<a href=#217 id=217 data-nosnippet>217</a> </span>.context(ViewNotFoundSnafu { view_name: <span class="kw-2">&amp;</span>view })<span class="question-mark">?</span>;
<a href=#218 id=218 data-nosnippet>218</a>
<a href=#219 id=219 data-nosnippet>219</a> <span class="kw">let </span>view_id = table_ref.table_info().ident.table_id;
<a href=#220 id=220 data-nosnippet>220</a>
<a href=#221 id=221 data-nosnippet>221</a> <span class="kw">let </span>view_info = <span class="self">self
<a href=#222 id=222 data-nosnippet>222</a> </span>.view_info_manager
<a href=#223 id=223 data-nosnippet>223</a> .get(view_id)
<a href=#224 id=224 data-nosnippet>224</a> .<span class="kw">await
<a href=#225 id=225 data-nosnippet>225</a> </span>.context(FindViewInfoSnafu { view_name: <span class="kw-2">&amp;</span>view })<span class="question-mark">?
<a href=#226 id=226 data-nosnippet>226</a> </span>.context(ViewInfoNotFoundSnafu { view_name: <span class="kw-2">&amp;</span>view })<span class="question-mark">?</span>;
<a href=#227 id=227 data-nosnippet>227</a>
<a href=#228 id=228 data-nosnippet>228</a> query::sql::show_create_view(show.view_name, <span class="kw-2">&amp;</span>view_info.definition, query_ctx)
<a href=#229 id=229 data-nosnippet>229</a> .context(error::ExecuteStatementSnafu)
<a href=#230 id=230 data-nosnippet>230</a> }
<a href=#231 id=231 data-nosnippet>231</a>
<a href=#232 id=232 data-nosnippet>232</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#233 id=233 data-nosnippet>233</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_views(
<a href=#234 id=234 data-nosnippet>234</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#235 id=235 data-nosnippet>235</a> stmt: ShowViews,
<a href=#236 id=236 data-nosnippet>236</a> query_ctx: QueryContextRef,
<a href=#237 id=237 data-nosnippet>237</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#238 id=238 data-nosnippet>238</a> query::sql::show_views(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#239 id=239 data-nosnippet>239</a> .<span class="kw">await
<a href=#240 id=240 data-nosnippet>240</a> </span>.context(ExecuteStatementSnafu)
<a href=#241 id=241 data-nosnippet>241</a> }
<a href=#242 id=242 data-nosnippet>242</a>
<a href=#243 id=243 data-nosnippet>243</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#244 id=244 data-nosnippet>244</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_flows(
<a href=#245 id=245 data-nosnippet>245</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#246 id=246 data-nosnippet>246</a> stmt: ShowFlows,
<a href=#247 id=247 data-nosnippet>247</a> query_ctx: QueryContextRef,
<a href=#248 id=248 data-nosnippet>248</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#249 id=249 data-nosnippet>249</a> query::sql::show_flows(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#250 id=250 data-nosnippet>250</a> .<span class="kw">await
<a href=#251 id=251 data-nosnippet>251</a> </span>.context(ExecuteStatementSnafu)
<a href=#252 id=252 data-nosnippet>252</a> }
<a href=#253 id=253 data-nosnippet>253</a>
<a href=#254 id=254 data-nosnippet>254</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#255 id=255 data-nosnippet>255</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_flow_status(
<a href=#256 id=256 data-nosnippet>256</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#257 id=257 data-nosnippet>257</a> stmt: ShowFlowStatus,
<a href=#258 id=258 data-nosnippet>258</a> query_ctx: QueryContextRef,
<a href=#259 id=259 data-nosnippet>259</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#260 id=260 data-nosnippet>260</a> query::sql::show_flow_status(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#261 id=261 data-nosnippet>261</a> .<span class="kw">await
<a href=#262 id=262 data-nosnippet>262</a> </span>.context(ExecuteStatementSnafu)
<a href=#263 id=263 data-nosnippet>263</a> }
<a href=#264 id=264 data-nosnippet>264</a>
<a href=#265 id=265 data-nosnippet>265</a> <span class="attr">#[cfg(feature = <span class="string">"enterprise"</span>)]
<a href=#266 id=266 data-nosnippet>266</a> #[tracing::instrument(skip_all)]
<a href=#267 id=267 data-nosnippet>267</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>show_triggers(
<a href=#268 id=268 data-nosnippet>268</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#269 id=269 data-nosnippet>269</a> stmt: sql::statements::show::trigger::ShowTriggers,
<a href=#270 id=270 data-nosnippet>270</a> query_ctx: QueryContextRef,
<a href=#271 id=271 data-nosnippet>271</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#272 id=272 data-nosnippet>272</a> query::sql::show_triggers(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#273 id=273 data-nosnippet>273</a> .<span class="kw">await
<a href=#274 id=274 data-nosnippet>274</a> </span>.context(ExecuteStatementSnafu)
<a href=#275 id=275 data-nosnippet>275</a> }
<a href=#276 id=276 data-nosnippet>276</a>
<a href=#277 id=277 data-nosnippet>277</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#278 id=278 data-nosnippet>278</a> </span><span class="kw">pub async fn </span>show_create_flow(
<a href=#279 id=279 data-nosnippet>279</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#280 id=280 data-nosnippet>280</a> show: ShowCreateFlow,
<a href=#281 id=281 data-nosnippet>281</a> query_ctx: QueryContextRef,
<a href=#282 id=282 data-nosnippet>282</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#283 id=283 data-nosnippet>283</a> <span class="kw">let </span>obj_name = <span class="kw-2">&amp;</span>show.flow_name;
<a href=#284 id=284 data-nosnippet>284</a> <span class="kw">let </span>(catalog_name, flow_name) = <span class="kw">match </span><span class="kw-2">&amp;</span>obj_name.<span class="number">0</span>[..] {
<a href=#285 id=285 data-nosnippet>285</a> [flow] =&gt; (query_ctx.current_catalog().to_string(), flow.to_string_unquoted()),
<a href=#286 id=286 data-nosnippet>286</a> [catalog, flow] =&gt; (catalog.to_string_unquoted(), flow.to_string_unquoted()),
<a href=#287 id=287 data-nosnippet>287</a> <span class="kw">_ </span>=&gt; {
<a href=#288 id=288 data-nosnippet>288</a> <span class="kw">return </span>InvalidSqlSnafu {
<a href=#289 id=289 data-nosnippet>289</a> err_msg: <span class="macro">format!</span>(
<a href=#290 id=290 data-nosnippet>290</a> <span class="string">"expect flow name to be &lt;catalog&gt;.&lt;flow_name&gt; or &lt;flow_name&gt;, actual: {obj_name}"</span>,
<a href=#291 id=291 data-nosnippet>291</a> ),
<a href=#292 id=292 data-nosnippet>292</a> }
<a href=#293 id=293 data-nosnippet>293</a> .fail()
<a href=#294 id=294 data-nosnippet>294</a> }
<a href=#295 id=295 data-nosnippet>295</a> };
<a href=#296 id=296 data-nosnippet>296</a>
<a href=#297 id=297 data-nosnippet>297</a> <span class="kw">let </span>flow_name_val = <span class="self">self
<a href=#298 id=298 data-nosnippet>298</a> </span>.flow_metadata_manager
<a href=#299 id=299 data-nosnippet>299</a> .flow_name_manager()
<a href=#300 id=300 data-nosnippet>300</a> .get(<span class="kw-2">&amp;</span>catalog_name, <span class="kw-2">&amp;</span>flow_name)
<a href=#301 id=301 data-nosnippet>301</a> .<span class="kw">await
<a href=#302 id=302 data-nosnippet>302</a> </span>.context(error::TableMetadataManagerSnafu)<span class="question-mark">?
<a href=#303 id=303 data-nosnippet>303</a> </span>.context(error::FlowNotFoundSnafu {
<a href=#304 id=304 data-nosnippet>304</a> flow_name: <span class="kw-2">&amp;</span>flow_name,
<a href=#305 id=305 data-nosnippet>305</a> })<span class="question-mark">?</span>;
<a href=#306 id=306 data-nosnippet>306</a>
<a href=#307 id=307 data-nosnippet>307</a> <span class="kw">let </span>flow_val = <span class="self">self
<a href=#308 id=308 data-nosnippet>308</a> </span>.flow_metadata_manager
<a href=#309 id=309 data-nosnippet>309</a> .flow_info_manager()
<a href=#310 id=310 data-nosnippet>310</a> .get(flow_name_val.flow_id())
<a href=#311 id=311 data-nosnippet>311</a> .<span class="kw">await
<a href=#312 id=312 data-nosnippet>312</a> </span>.context(error::TableMetadataManagerSnafu)<span class="question-mark">?
<a href=#313 id=313 data-nosnippet>313</a> </span>.context(error::FlowNotFoundSnafu {
<a href=#314 id=314 data-nosnippet>314</a> flow_name: <span class="kw-2">&amp;</span>flow_name,
<a href=#315 id=315 data-nosnippet>315</a> })<span class="question-mark">?</span>;
<a href=#316 id=316 data-nosnippet>316</a>
<a href=#317 id=317 data-nosnippet>317</a> query::sql::show_create_flow(obj_name.clone(), flow_val, query_ctx)
<a href=#318 id=318 data-nosnippet>318</a> .context(error::ExecuteStatementSnafu)
<a href=#319 id=319 data-nosnippet>319</a> }
<a href=#320 id=320 data-nosnippet>320</a>
<a href=#321 id=321 data-nosnippet>321</a> <span class="attr">#[cfg(feature = <span class="string">"enterprise"</span>)]
<a href=#322 id=322 data-nosnippet>322</a> #[tracing::instrument(skip_all)]
<a href=#323 id=323 data-nosnippet>323</a> </span><span class="kw">pub async fn </span>show_create_trigger(
<a href=#324 id=324 data-nosnippet>324</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#325 id=325 data-nosnippet>325</a> show: sql::statements::show::trigger::ShowCreateTrigger,
<a href=#326 id=326 data-nosnippet>326</a> query_ctx: QueryContextRef,
<a href=#327 id=327 data-nosnippet>327</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#328 id=328 data-nosnippet>328</a> <span class="kw">let </span><span class="prelude-val">Some</span>(trigger_querier) = <span class="self">self</span>.trigger_querier.as_ref() <span class="kw">else </span>{
<a href=#329 id=329 data-nosnippet>329</a> <span class="kw">return </span>error::MissingTriggerQuerierSnafu.fail();
<a href=#330 id=330 data-nosnippet>330</a> };
<a href=#331 id=331 data-nosnippet>331</a>
<a href=#332 id=332 data-nosnippet>332</a> <span class="kw">let </span>obj_name = <span class="kw-2">&amp;</span>show.trigger_name;
<a href=#333 id=333 data-nosnippet>333</a> <span class="kw">let </span>(catalog_name, trigger_name) = <span class="kw">match </span><span class="kw-2">&amp;</span>obj_name.<span class="number">0</span>[..] {
<a href=#334 id=334 data-nosnippet>334</a> [trigger] =&gt; (query_ctx.current_catalog().to_string(), trigger.to_string_unquoted()),
<a href=#335 id=335 data-nosnippet>335</a> [catalog, trigger] =&gt; (catalog.to_string_unquoted(), trigger.to_string_unquoted()),
<a href=#336 id=336 data-nosnippet>336</a> <span class="kw">_ </span>=&gt; {
<a href=#337 id=337 data-nosnippet>337</a> <span class="kw">return </span>InvalidSqlSnafu {
<a href=#338 id=338 data-nosnippet>338</a> err_msg: <span class="macro">format!</span>(
<a href=#339 id=339 data-nosnippet>339</a> <span class="string">"expect trigger name to be &lt;catalog&gt;.&lt;trigger_name&gt; or &lt;trigger_name&gt;, actual: {obj_name}"</span>,
<a href=#340 id=340 data-nosnippet>340</a> ),
<a href=#341 id=341 data-nosnippet>341</a> }
<a href=#342 id=342 data-nosnippet>342</a> .fail()
<a href=#343 id=343 data-nosnippet>343</a> }
<a href=#344 id=344 data-nosnippet>344</a> };
<a href=#345 id=345 data-nosnippet>345</a> trigger_querier
<a href=#346 id=346 data-nosnippet>346</a> .show_create_trigger(<span class="kw-2">&amp;</span>catalog_name, <span class="kw-2">&amp;</span>trigger_name, <span class="kw-2">&amp;</span>query_ctx)
<a href=#347 id=347 data-nosnippet>347</a> .<span class="kw">await
<a href=#348 id=348 data-nosnippet>348</a> </span>.context(error::TriggerQuerierSnafu)
<a href=#349 id=349 data-nosnippet>349</a> }
<a href=#350 id=350 data-nosnippet>350</a>
<a href=#351 id=351 data-nosnippet>351</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#352 id=352 data-nosnippet>352</a> </span><span class="kw">pub fn </span>show_variable(<span class="kw-2">&amp;</span><span class="self">self</span>, stmt: ShowVariables, query_ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#353 id=353 data-nosnippet>353</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(output) = query::sql::show_query_option(
<a href=#354 id=354 data-nosnippet>354</a> stmt.clone(),
<a href=#355 id=355 data-nosnippet>355</a> query_ctx.clone(),
<a href=#356 id=356 data-nosnippet>356</a> <span class="self">self</span>.query_engine.as_ref(),
<a href=#357 id=357 data-nosnippet>357</a> )
<a href=#358 id=358 data-nosnippet>358</a> .context(error::ExecuteStatementSnafu)<span class="question-mark">?
<a href=#359 id=359 data-nosnippet>359</a> </span>{
<a href=#360 id=360 data-nosnippet>360</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(output);
<a href=#361 id=361 data-nosnippet>361</a> }
<a href=#362 id=362 data-nosnippet>362</a> query::sql::show_variable(stmt, query_ctx).context(error::ExecuteStatementSnafu)
<a href=#363 id=363 data-nosnippet>363</a> }
<a href=#364 id=364 data-nosnippet>364</a>
<a href=#365 id=365 data-nosnippet>365</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#366 id=366 data-nosnippet>366</a> </span><span class="kw">pub async fn </span>show_collation(
<a href=#367 id=367 data-nosnippet>367</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#368 id=368 data-nosnippet>368</a> kind: ShowKind,
<a href=#369 id=369 data-nosnippet>369</a> query_ctx: QueryContextRef,
<a href=#370 id=370 data-nosnippet>370</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#371 id=371 data-nosnippet>371</a> query::sql::show_collations(kind, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#372 id=372 data-nosnippet>372</a> .<span class="kw">await
<a href=#373 id=373 data-nosnippet>373</a> </span>.context(error::ExecuteStatementSnafu)
<a href=#374 id=374 data-nosnippet>374</a> }
<a href=#375 id=375 data-nosnippet>375</a>
<a href=#376 id=376 data-nosnippet>376</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#377 id=377 data-nosnippet>377</a> </span><span class="kw">pub async fn </span>show_charset(<span class="kw-2">&amp;</span><span class="self">self</span>, kind: ShowKind, query_ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#378 id=378 data-nosnippet>378</a> query::sql::show_charsets(kind, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#379 id=379 data-nosnippet>379</a> .<span class="kw">await
<a href=#380 id=380 data-nosnippet>380</a> </span>.context(error::ExecuteStatementSnafu)
<a href=#381 id=381 data-nosnippet>381</a> }
<a href=#382 id=382 data-nosnippet>382</a>
<a href=#383 id=383 data-nosnippet>383</a> <span class="attr">#[tracing::instrument(skip_all)]
<a href=#384 id=384 data-nosnippet>384</a> </span><span class="kw">pub async fn </span>show_status(<span class="kw-2">&amp;</span><span class="self">self</span>, query_ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#385 id=385 data-nosnippet>385</a> query::sql::show_status(query_ctx)
<a href=#386 id=386 data-nosnippet>386</a> .<span class="kw">await
<a href=#387 id=387 data-nosnippet>387</a> </span>.context(error::ExecuteStatementSnafu)
<a href=#388 id=388 data-nosnippet>388</a> }
<a href=#389 id=389 data-nosnippet>389</a> <span class="kw">pub async fn </span>show_search_path(<span class="kw-2">&amp;</span><span class="self">self</span>, query_ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#390 id=390 data-nosnippet>390</a> query::sql::show_search_path(query_ctx)
<a href=#391 id=391 data-nosnippet>391</a> .<span class="kw">await
<a href=#392 id=392 data-nosnippet>392</a> </span>.context(error::ExecuteStatementSnafu)
<a href=#393 id=393 data-nosnippet>393</a> }
<a href=#394 id=394 data-nosnippet>394</a>
<a href=#395 id=395 data-nosnippet>395</a> <span class="kw">pub async fn </span>show_processlist(
<a href=#396 id=396 data-nosnippet>396</a> <span class="kw-2">&amp;</span><span class="self">self</span>,
<a href=#397 id=397 data-nosnippet>397</a> stmt: ShowProcessList,
<a href=#398 id=398 data-nosnippet>398</a> query_ctx: QueryContextRef,
<a href=#399 id=399 data-nosnippet>399</a> ) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#400 id=400 data-nosnippet>400</a> query::sql::show_processlist(stmt, <span class="kw-2">&amp;</span><span class="self">self</span>.query_engine, <span class="kw-2">&amp;</span><span class="self">self</span>.catalog_manager, query_ctx)
<a href=#401 id=401 data-nosnippet>401</a> .<span class="kw">await
<a href=#402 id=402 data-nosnippet>402</a> </span>.context(ExecLogicalPlanSnafu)
<a href=#403 id=403 data-nosnippet>403</a> }
<a href=#404 id=404 data-nosnippet>404</a>}
<a href=#405 id=405 data-nosnippet>405</a>
<a href=#406 id=406 data-nosnippet>406</a><span class="kw">pub</span>(<span class="kw">crate</span>) <span class="kw">fn </span>create_partitions_stmt(
<a href=#407 id=407 data-nosnippet>407</a> table_info: <span class="kw-2">&amp;</span>TableInfo,
<a href=#408 id=408 data-nosnippet>408</a> partitions: <span class="kw-2">&amp;</span>[PartitionInfoWithVersion],
<a href=#409 id=409 data-nosnippet>409</a>) -&gt; <span class="prelude-ty">Result</span>&lt;<span class="prelude-ty">Option</span>&lt;Partitions&gt;&gt; {
<a href=#410 id=410 data-nosnippet>410</a> <span class="kw">if </span>partitions.is_empty() {
<a href=#411 id=411 data-nosnippet>411</a> <span class="kw">return </span><span class="prelude-val">Ok</span>(<span class="prelude-val">None</span>);
<a href=#412 id=412 data-nosnippet>412</a> }
<a href=#413 id=413 data-nosnippet>413</a>
<a href=#414 id=414 data-nosnippet>414</a> <span class="kw">let </span>column_list = table_info
<a href=#415 id=415 data-nosnippet>415</a> .meta
<a href=#416 id=416 data-nosnippet>416</a> .partition_column_names()
<a href=#417 id=417 data-nosnippet>417</a> .map(|name| name[..].into())
<a href=#418 id=418 data-nosnippet>418</a> .collect();
<a href=#419 id=419 data-nosnippet>419</a>
<a href=#420 id=420 data-nosnippet>420</a> <span class="kw">let </span>exprs = partitions
<a href=#421 id=421 data-nosnippet>421</a> .iter()
<a href=#422 id=422 data-nosnippet>422</a> .filter_map(|partition| {
<a href=#423 id=423 data-nosnippet>423</a> partition
<a href=#424 id=424 data-nosnippet>424</a> .partition_expr
<a href=#425 id=425 data-nosnippet>425</a> .as_ref()
<a href=#426 id=426 data-nosnippet>426</a> .map(|expr| expr.to_parser_expr())
<a href=#427 id=427 data-nosnippet>427</a> })
<a href=#428 id=428 data-nosnippet>428</a> .collect();
<a href=#429 id=429 data-nosnippet>429</a>
<a href=#430 id=430 data-nosnippet>430</a> <span class="prelude-val">Ok</span>(<span class="prelude-val">Some</span>(Partitions { column_list, exprs }))
<a href=#431 id=431 data-nosnippet>431</a>}
</code></pre></div></section></main></body></html>
+121
View File
@@ -0,0 +1,121 @@
<!DOCTYPE html><html lang="en"><head><meta charset="utf-8"><meta name="viewport" content="width=device-width, initial-scale=1.0"><meta name="generator" content="rustdoc"><meta name="description" content="Source of the Rust file `src/operator/src/statement/tql.rs`."><title>tql.rs - source</title><script>if(window.location.protocol!=="file:")document.head.insertAdjacentHTML("beforeend","SourceSerif4-Regular-6b053e98.ttf.woff2,FiraSans-Italic-81dc35de.woff2,FiraSans-Regular-0fe48ade.woff2,FiraSans-MediumItalic-ccf7e434.woff2,FiraSans-Medium-e1aa3f0a.woff2,SourceCodePro-Regular-8badfe75.ttf.woff2,SourceCodePro-Semibold-aa29a496.ttf.woff2".split(",").map(f=>`<link rel="preload" as="font" type="font/woff2"href="../../../static.files/${f}">`).join(""))</script><link rel="stylesheet" href="../../../static.files/normalize-9960930a.css"><link rel="stylesheet" href="../../../static.files/rustdoc-1ed95124.css"><meta name="rustdoc-vars" data-root-path="../../../" data-static-root-path="../../../static.files/" data-current-crate="operator" data-themes="" data-resource-suffix="" data-rustdoc-version="1.98.1 (48a229cea 2026-09-01)" data-channel="1.98.1" data-search-js="search-9944d736.js" data-stringdex-js="stringdex-2da4960a.js" data-settings-js="settings-170eb4bf.js" ><script src="../../../static.files/storage-41dd4d93.js"></script><script defer src="../../../static.files/src-script-813739b1.js"></script><script defer src="../../../src-files.js"></script><script defer src="../../../static.files/main-417ebf0a.js"></script><noscript><link rel="stylesheet" href="../../../static.files/noscript-f7c3ffd8.css"></noscript><link rel="alternate icon" type="image/png" href="../../../static.files/favicon-32x32-eab170b8.png"><link rel="icon" type="image/svg+xml" href="../../../static.files/favicon-044be391.svg"></head><body class="rustdoc src"><a class="skip-main-content" href="#main-content">Skip to main content</a><!--[if lte IE 11]><div class="warning">This old browser is unsupported and will most likely display funky things.</div><![endif]--><nav class="sidebar"><div class="src-sidebar-title"><h2>Files</h2></div></nav><div class="sidebar-resizer" title="Drag to resize sidebar"></div><main><section id="main-content" class="content" tabindex="-1"><div class="main-heading"><h1><div class="sub-heading">operator/statement/</div>tql.rs</h1><rustdoc-toolbar></rustdoc-toolbar></div><div class="example-wrap digits-3"><pre class="rust"><code><a href=#1 id=1 data-nosnippet>1</a><span class="comment">// Copyright 2023 Greptime Team
<a href=#2 id=2 data-nosnippet>2</a>//
<a href=#3 id=3 data-nosnippet>3</a>// Licensed under the Apache License, Version 2.0 (the "License");
<a href=#4 id=4 data-nosnippet>4</a>// you may not use this file except in compliance with the License.
<a href=#5 id=5 data-nosnippet>5</a>// You may obtain a copy of the License at
<a href=#6 id=6 data-nosnippet>6</a>//
<a href=#7 id=7 data-nosnippet>7</a>// http://www.apache.org/licenses/LICENSE-2.0
<a href=#8 id=8 data-nosnippet>8</a>//
<a href=#9 id=9 data-nosnippet>9</a>// Unless required by applicable law or agreed to in writing, software
<a href=#10 id=10 data-nosnippet>10</a>// distributed under the License is distributed on an "AS IS" BASIS,
<a href=#11 id=11 data-nosnippet>11</a>// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
<a href=#12 id=12 data-nosnippet>12</a>// See the License for the specific language governing permissions and
<a href=#13 id=13 data-nosnippet>13</a>// limitations under the License.
<a href=#14 id=14 data-nosnippet>14</a>
<a href=#15 id=15 data-nosnippet>15</a></span><span class="kw">use </span>std::collections::HashMap;
<a href=#16 id=16 data-nosnippet>16</a>
<a href=#17 id=17 data-nosnippet>17</a><span class="kw">use </span>common_query::Output;
<a href=#18 id=18 data-nosnippet>18</a><span class="kw">use </span>common_telemetry::tracing;
<a href=#19 id=19 data-nosnippet>19</a><span class="kw">use </span>datafusion_expr::LogicalPlan;
<a href=#20 id=20 data-nosnippet>20</a><span class="kw">use </span>query::parser::{
<a href=#21 id=21 data-nosnippet>21</a> ANALYZE_NODE_NAME, ANALYZE_VERBOSE_NODE_NAME, DEFAULT_LOOKBACK_STRING, EXPLAIN_NODE_NAME,
<a href=#22 id=22 data-nosnippet>22</a> EXPLAIN_VERBOSE_NODE_NAME, PromQuery, QueryLanguageParser,
<a href=#23 id=23 data-nosnippet>23</a>};
<a href=#24 id=24 data-nosnippet>24</a><span class="kw">use </span>session::context::QueryContextRef;
<a href=#25 id=25 data-nosnippet>25</a><span class="kw">use </span>snafu::ResultExt;
<a href=#26 id=26 data-nosnippet>26</a><span class="kw">use </span>sql::statements::tql::Tql;
<a href=#27 id=27 data-nosnippet>27</a>
<a href=#28 id=28 data-nosnippet>28</a><span class="kw">use </span><span class="kw">crate</span>::error::{ExecLogicalPlanSnafu, ParseQuerySnafu, PlanStatementSnafu, <span class="prelude-ty">Result</span>};
<a href=#29 id=29 data-nosnippet>29</a><span class="kw">use </span><span class="kw">crate</span>::statement::StatementExecutor;
<a href=#30 id=30 data-nosnippet>30</a>
<a href=#31 id=31 data-nosnippet>31</a><span class="kw">impl </span>StatementExecutor {
<a href=#32 id=32 data-nosnippet>32</a> <span class="doccomment">/// Plan the given [Tql] query and return the [LogicalPlan].
<a href=#33 id=33 data-nosnippet>33</a> </span><span class="attr">#[tracing::instrument(skip_all)]
<a href=#34 id=34 data-nosnippet>34</a> </span><span class="kw">pub async fn </span>plan_tql(<span class="kw-2">&amp;</span><span class="self">self</span>, tql: Tql, query_ctx: <span class="kw-2">&amp;</span>QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;LogicalPlan&gt; {
<a href=#35 id=35 data-nosnippet>35</a> <span class="kw">let </span>stmt = <span class="kw">match </span>tql {
<a href=#36 id=36 data-nosnippet>36</a> Tql::Eval(eval) =&gt; {
<a href=#37 id=37 data-nosnippet>37</a> <span class="kw">let </span>promql = PromQuery {
<a href=#38 id=38 data-nosnippet>38</a> start: eval.start,
<a href=#39 id=39 data-nosnippet>39</a> end: eval.end,
<a href=#40 id=40 data-nosnippet>40</a> step: eval.step,
<a href=#41 id=41 data-nosnippet>41</a> query: eval.query,
<a href=#42 id=42 data-nosnippet>42</a> lookback: eval
<a href=#43 id=43 data-nosnippet>43</a> .lookback
<a href=#44 id=44 data-nosnippet>44</a> .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
<a href=#45 id=45 data-nosnippet>45</a> alias: eval.alias,
<a href=#46 id=46 data-nosnippet>46</a> };
<a href=#47 id=47 data-nosnippet>47</a> QueryLanguageParser::parse_promql(<span class="kw-2">&amp;</span>promql, query_ctx).context(ParseQuerySnafu)<span class="question-mark">?
<a href=#48 id=48 data-nosnippet>48</a> </span>}
<a href=#49 id=49 data-nosnippet>49</a> Tql::Explain(explain) =&gt; {
<a href=#50 id=50 data-nosnippet>50</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(format) = <span class="kw-2">&amp;</span>explain.format {
<a href=#51 id=51 data-nosnippet>51</a> query_ctx.set_explain_format(format.to_string());
<a href=#52 id=52 data-nosnippet>52</a> }
<a href=#53 id=53 data-nosnippet>53</a>
<a href=#54 id=54 data-nosnippet>54</a> <span class="kw">let </span>promql = PromQuery {
<a href=#55 id=55 data-nosnippet>55</a> query: explain.query,
<a href=#56 id=56 data-nosnippet>56</a> start: explain.start,
<a href=#57 id=57 data-nosnippet>57</a> end: explain.end,
<a href=#58 id=58 data-nosnippet>58</a> step: explain.step,
<a href=#59 id=59 data-nosnippet>59</a> lookback: explain
<a href=#60 id=60 data-nosnippet>60</a> .lookback
<a href=#61 id=61 data-nosnippet>61</a> .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
<a href=#62 id=62 data-nosnippet>62</a> alias: explain.alias,
<a href=#63 id=63 data-nosnippet>63</a> };
<a href=#64 id=64 data-nosnippet>64</a> <span class="kw">let </span>explain_node_name = <span class="kw">if </span>explain.is_verbose {
<a href=#65 id=65 data-nosnippet>65</a> EXPLAIN_VERBOSE_NODE_NAME
<a href=#66 id=66 data-nosnippet>66</a> } <span class="kw">else </span>{
<a href=#67 id=67 data-nosnippet>67</a> EXPLAIN_NODE_NAME
<a href=#68 id=68 data-nosnippet>68</a> }
<a href=#69 id=69 data-nosnippet>69</a> .to_string();
<a href=#70 id=70 data-nosnippet>70</a> <span class="kw">let </span>params = HashMap::from([(<span class="string">"name"</span>.to_string(), explain_node_name)]);
<a href=#71 id=71 data-nosnippet>71</a> QueryLanguageParser::parse_promql(<span class="kw-2">&amp;</span>promql, query_ctx)
<a href=#72 id=72 data-nosnippet>72</a> .context(ParseQuerySnafu)<span class="question-mark">?
<a href=#73 id=73 data-nosnippet>73</a> </span>.post_process(params)
<a href=#74 id=74 data-nosnippet>74</a> .context(ParseQuerySnafu)<span class="question-mark">?
<a href=#75 id=75 data-nosnippet>75</a> </span>}
<a href=#76 id=76 data-nosnippet>76</a> Tql::Analyze(analyze) =&gt; {
<a href=#77 id=77 data-nosnippet>77</a> <span class="kw">if let </span><span class="prelude-val">Some</span>(format) = <span class="kw-2">&amp;</span>analyze.format {
<a href=#78 id=78 data-nosnippet>78</a> query_ctx.set_explain_format(format.to_string());
<a href=#79 id=79 data-nosnippet>79</a> }
<a href=#80 id=80 data-nosnippet>80</a>
<a href=#81 id=81 data-nosnippet>81</a> <span class="kw">let </span>promql = PromQuery {
<a href=#82 id=82 data-nosnippet>82</a> start: analyze.start,
<a href=#83 id=83 data-nosnippet>83</a> end: analyze.end,
<a href=#84 id=84 data-nosnippet>84</a> step: analyze.step,
<a href=#85 id=85 data-nosnippet>85</a> query: analyze.query,
<a href=#86 id=86 data-nosnippet>86</a> lookback: analyze
<a href=#87 id=87 data-nosnippet>87</a> .lookback
<a href=#88 id=88 data-nosnippet>88</a> .unwrap_or_else(|| DEFAULT_LOOKBACK_STRING.to_string()),
<a href=#89 id=89 data-nosnippet>89</a> alias: analyze.alias,
<a href=#90 id=90 data-nosnippet>90</a> };
<a href=#91 id=91 data-nosnippet>91</a> <span class="kw">let </span>analyze_node_name = <span class="kw">if </span>analyze.is_verbose {
<a href=#92 id=92 data-nosnippet>92</a> ANALYZE_VERBOSE_NODE_NAME
<a href=#93 id=93 data-nosnippet>93</a> } <span class="kw">else </span>{
<a href=#94 id=94 data-nosnippet>94</a> ANALYZE_NODE_NAME
<a href=#95 id=95 data-nosnippet>95</a> }
<a href=#96 id=96 data-nosnippet>96</a> .to_string();
<a href=#97 id=97 data-nosnippet>97</a> <span class="kw">let </span>params = HashMap::from([(<span class="string">"name"</span>.to_string(), analyze_node_name)]);
<a href=#98 id=98 data-nosnippet>98</a> QueryLanguageParser::parse_promql(<span class="kw-2">&amp;</span>promql, query_ctx)
<a href=#99 id=99 data-nosnippet>99</a> .context(ParseQuerySnafu)<span class="question-mark">?
<a href=#100 id=100 data-nosnippet>100</a> </span>.post_process(params)
<a href=#101 id=101 data-nosnippet>101</a> .context(ParseQuerySnafu)<span class="question-mark">?
<a href=#102 id=102 data-nosnippet>102</a> </span>}
<a href=#103 id=103 data-nosnippet>103</a> };
<a href=#104 id=104 data-nosnippet>104</a> <span class="self">self</span>.query_engine
<a href=#105 id=105 data-nosnippet>105</a> .planner()
<a href=#106 id=106 data-nosnippet>106</a> .plan(<span class="kw-2">&amp;</span>stmt, query_ctx.clone())
<a href=#107 id=107 data-nosnippet>107</a> .<span class="kw">await
<a href=#108 id=108 data-nosnippet>108</a> </span>.context(PlanStatementSnafu)
<a href=#109 id=109 data-nosnippet>109</a> }
<a href=#110 id=110 data-nosnippet>110</a>
<a href=#111 id=111 data-nosnippet>111</a> <span class="doccomment">/// Execute the given [Tql] query and return the result.
<a href=#112 id=112 data-nosnippet>112</a> </span><span class="attr">#[tracing::instrument(skip_all)]
<a href=#113 id=113 data-nosnippet>113</a> </span><span class="kw">pub</span>(<span class="kw">super</span>) <span class="kw">async fn </span>execute_tql(<span class="kw-2">&amp;</span><span class="self">self</span>, tql: Tql, query_ctx: QueryContextRef) -&gt; <span class="prelude-ty">Result</span>&lt;Output&gt; {
<a href=#114 id=114 data-nosnippet>114</a> <span class="kw">let </span>plan = <span class="self">self</span>.plan_tql(tql, <span class="kw-2">&amp;</span>query_ctx).<span class="kw">await</span><span class="question-mark">?</span>;
<a href=#115 id=115 data-nosnippet>115</a> <span class="self">self</span>.query_engine
<a href=#116 id=116 data-nosnippet>116</a> .execute(plan, query_ctx)
<a href=#117 id=117 data-nosnippet>117</a> .<span class="kw">await
<a href=#118 id=118 data-nosnippet>118</a> </span>.context(ExecLogicalPlanSnafu)
<a href=#119 id=119 data-nosnippet>119</a> }
<a href=#120 id=120 data-nosnippet>120</a>}
</code></pre></div></section></main></body></html>