77
88package org .elasticsearch .xpack .ml .datafeed .extractor .esql ;
99
10+ import org .apache .logging .log4j .LogManager ;
11+ import org .apache .logging .log4j .Logger ;
12+ import org .elasticsearch .client .internal .Client ;
13+ import org .elasticsearch .common .Strings ;
14+ import org .elasticsearch .common .io .stream .BytesStreamOutput ;
15+ import org .elasticsearch .xcontent .XContentBuilder ;
16+ import org .elasticsearch .xcontent .json .JsonXContent ;
17+ import org .elasticsearch .xpack .core .ClientHelper ;
18+ import org .elasticsearch .xpack .core .ml .datafeed .DatafeedConfig ;
19+ import org .elasticsearch .xpack .core .ml .datafeed .SearchInterval ;
20+ import org .elasticsearch .xpack .esql .action .ColumnInfo ;
21+ import org .elasticsearch .xpack .esql .action .EsqlQueryAction ;
22+ import org .elasticsearch .xpack .esql .action .EsqlQueryRequest ;
23+ import org .elasticsearch .xpack .esql .action .EsqlQueryResponse ;
1024import org .elasticsearch .xpack .ml .datafeed .extractor .DataExtractor ;
25+ import org .elasticsearch .xpack .ql .util .DateUtils ;
1126
1227import java .io .IOException ;
28+ import java .util .Iterator ;
29+ import java .util .List ;
30+ import java .util .Objects ;
31+ import java .util .Optional ;
1332
1433public class EsqlDataExtractor implements DataExtractor {
1534
35+ private static final Logger logger = LogManager .getLogger (EsqlDataExtractor .class );
36+
37+ private final Client client ;
38+ private final DatafeedConfig datafeed ;
39+ private final String timeField ;
40+ private final SearchInterval interval ;
41+ private boolean isCancelled ;
42+
43+ EsqlDataExtractor (Client client , DatafeedConfig datafeed , String timeField , long start , long end ) {
44+ this .client = Objects .requireNonNull (client );
45+ this .datafeed = Objects .requireNonNull (datafeed );
46+ this .timeField = timeField ;
47+ this .interval = new SearchInterval (start , end );
48+ this .isCancelled = false ;
49+ }
50+
51+ // TODO: check whether these expressions facilitate injection attacks!
52+ private String esqlTimeFilter () {
53+ return Strings .format (
54+ " | WHERE %s >= \" %s\" AND %s < \" %s\" " ,
55+ timeField ,
56+ DateUtils .UTC_DATE_TIME_FORMATTER .formatMillis (
57+ Math .min (interval .startMs (), org .elasticsearch .common .time .DateUtils .MAX_MILLIS_BEFORE_9999 )
58+ ),
59+ timeField ,
60+ DateUtils .UTC_DATE_TIME_FORMATTER .formatMillis (
61+ Math .min (interval .endMs (), org .elasticsearch .common .time .DateUtils .MAX_MILLIS_BEFORE_9999 )
62+ )
63+ );
64+ }
65+
66+ private String esqlSortByTime () {
67+ return Strings .format (" | SORT %s" , timeField );
68+ }
69+
70+ private String esqlSummaryStats () {
71+ return Strings .format (" | STATS earliest_time=MIN(%s), latest_time=MAX(%s), total_hits=COUNT(*)" , timeField , timeField );
72+ }
73+
1674 @ Override
1775 public DataSummary getSummary () {
18- return null ;
76+ EsqlQueryRequest request = new EsqlQueryRequest ();
77+ request .query (datafeed .getEsqlQuery () + esqlTimeFilter () + esqlSummaryStats ());
78+
79+ try (EsqlQueryResponse response = execute (request )) {
80+ Iterator <Object > values = response .values ().next ();
81+ String earliestTime = (String ) values .next ();
82+ String latestTime = (String ) values .next ();
83+ long totalHits = (long ) values .next ();
84+ return new DataSummary (
85+ earliestTime == null ? null : DateUtils .UTC_DATE_TIME_FORMATTER .parseMillis (earliestTime ),
86+ latestTime == null ? null : DateUtils .UTC_DATE_TIME_FORMATTER .parseMillis (latestTime ),
87+ totalHits
88+ );
89+ }
1990 }
2091
2192 @ Override
@@ -25,26 +96,71 @@ public boolean hasNext() {
2596
2697 @ Override
2798 public Result next () throws IOException {
28- return null ;
99+ EsqlQueryRequest request = new EsqlQueryRequest ();
100+ request .query (datafeed .getEsqlQuery () + esqlTimeFilter () + esqlSortByTime ());
101+
102+ EsqlQueryResponse response = execute (request );
103+
104+ try (BytesStreamOutput outputStream = new BytesStreamOutput ()) {
105+ XContentBuilder jsonBuilder = new XContentBuilder (JsonXContent .jsonXContent , outputStream );
106+
107+ List <ColumnInfo > columns = response .columns ();
108+ int valueCount = 0 ;
109+ for (Iterator <Iterator <Object >> itValues = response .values (); itValues .hasNext ();) {
110+ jsonBuilder .startObject ();
111+ int index = 0 ;
112+ for (Iterator <Object > itValue = itValues .next (); itValue .hasNext ();) {
113+ Object value = itValue .next ();
114+ if ("date" .equals (columns .get (index ).type ())) {
115+ if (value instanceof String && Strings .isNullOrEmpty ((String ) value ) == false ) {
116+ value = DateUtils .UTC_DATE_TIME_FORMATTER .parseMillis ((String ) value );
117+ }
118+ // TODO: something with arrays of dates? (e.g. kibana_sample_data_ecommerce -> products.created_on)
119+ }
120+ jsonBuilder .field (columns .get (index ).name (), value );
121+ index ++;
122+ }
123+ jsonBuilder .endObject ();
124+ valueCount ++;
125+ }
126+ jsonBuilder .close ();
127+
128+ logger .info (
129+ "query interval: {} - {}, valueCount: {}" ,
130+ DateUtils .UTC_DATE_TIME_FORMATTER .formatMillis (interval .startMs ()),
131+ DateUtils .UTC_DATE_TIME_FORMATTER .formatMillis (interval .endMs ()),
132+ valueCount
133+ );
134+
135+ return new Result (interval , Optional .of (outputStream .bytes ().streamInput ()));
136+ }
29137 }
30138
31139 @ Override
32140 public boolean isCancelled () {
33- return false ;
141+ return isCancelled ;
34142 }
35143
36144 @ Override
37145 public void cancel () {
38-
146+ logger .trace ("Data extractor received cancel request" );
147+ isCancelled = true ;
39148 }
40149
41150 @ Override
42- public void destroy () {
43-
44- }
151+ public void destroy () {}
45152
46153 @ Override
47154 public long getEndTime () {
48- return 0 ;
155+ return interval .endMs ();
156+ }
157+
158+ private EsqlQueryResponse execute (EsqlQueryRequest request ) {
159+ return ClientHelper .executeWithHeaders (
160+ datafeed .getHeaders (),
161+ ClientHelper .ML_ORIGIN ,
162+ client ,
163+ () -> client .execute (EsqlQueryAction .INSTANCE , request ).actionGet ()
164+ );
49165 }
50166}
0 commit comments