enhancement(vrl): Add geoip enrichment table (#13338)

* Inital impl * Add tests * Add docs * Rename to path Signed-off-by: Kruno Tomola Fabro <krunotf@gmail.com> * Remove examples Signed-off-by: Kruno Tomola Fabro <krunotf@gmail.com> * Small fixes Signed-off-by: Kruno Tomola Fabro <krunotf@gmail.com> * Apply review comments Signed-off-by: Kruno Tomola Fabro <krunotf@gmail.com> * Add benches Signed-off-by: Kruno Tomola Fabro <krunotf@gmail.com> * Bump Signed-off-by: Kruno Tomola Fabro <krunotf@gmail.com> * Bump Signed-off-by: Kruno Tomola Fabro <krunotf@gmail.com>
vectordotdev · Jul 7, 2022 · aee1984 · aee1984
1 parent ab2c073
commit aee1984
Show file tree

Hide file tree

Showing 6 changed files with 609 additions and 11 deletions.
diff --git a/Cargo.toml b/Cargo.toml
@@ -428,8 +428,9 @@ protobuf-build = ["dep:tonic-build", "dep:prost-build"]
 gcp = ["dep:base64", "dep:goauth", "dep:smpl_jwt"]
 
 # Enrichment Tables
-enrichment-tables = ["enrichment-tables-file"]
+enrichment-tables = ["enrichment-tables-file","enrichment-tables-geoip"]
 enrichment-tables-file = [ "dep:csv", "dep:seahash", "dep:hash_hasher" ]
+enrichment-tables-geoip = ["dep:maxminddb"]
 
 # Sources
 sources = ["sources-logs", "sources-metrics"]
@@ -809,7 +810,7 @@ remap-benches = ["transforms-remap"]
 transform-benches = ["transforms-filter", "transforms-dedupe", "transforms-reduce", "transforms-route"]
 codecs-benches = []
 loki-benches = ["sinks-loki"]
-enrichment-tables-benches = ["enrichment-tables-file"]
+enrichment-tables-benches = ["enrichment-tables-file","enrichment-tables-geoip"]
 
 [[bench]]
 name = "default"
@@ -828,7 +829,7 @@ harness = false
 required-features = ["remap-benches"]
 
 [[bench]]
-name = "enrichment_tables_file"
+name = "enrichment_tables"
 harness = false
 required-features = ["enrichment-tables-benches"]
 

diff --git a/benches/enrichment_tables_file.rs → benches/enrichment_tables.rs b/benches/enrichment_tables_file.rs → benches/enrichment_tables.rs
@@ -4,12 +4,16 @@ use chrono::prelude::*;
 use criterion::{criterion_group, criterion_main, BatchSize, Criterion};
 use enrichment::Case;
 use value::Value;
-use vector::enrichment_tables::{file::File, Condition, Table};
+use vector::enrichment_tables::{
+    file::File,
+    geoip::{Geoip, GeoipConfig},
+    Condition, Table,
+};
 
 criterion_group!(
     name = benches;
     config = Criterion::default().noise_threshold(0.02).sample_size(10);
-    targets = benchmark_enrichment_tables_file
+    targets = benchmark_enrichment_tables_file, benchmark_enrichment_tables_geoip
 );
 criterion_main!(benches);
 
@@ -221,3 +225,86 @@ fn benchmark_enrichment_tables_file(c: &mut Criterion) {
         },
     );
 }
+
+fn benchmark_enrichment_tables_geoip(c: &mut Criterion) {
+    let mut group = c.benchmark_group("enrichment_tables_geoip");
+    let build = |path: &str| {
+        Geoip::new(GeoipConfig {
+            path: path.to_string(),
+            locale: "en".to_string(),
+        })
+        .unwrap()
+    };
+
+    group.bench_function("enrichment_tables/geoip_isp", |b| {
+        let table = build("tests/data/GeoIP2-ISP-Test.mmdb");
+        let ip = "208.192.1.2";
+        let mut expected = BTreeMap::<String, Value>::new();
+        expected.insert("autonomous_system_number".to_string(), 701i64.into());
+        expected.insert(
+            "autonomous_system_organization".to_string(),
+            "MCI Communications Services, Inc. d/b/a Verizon Business".into(),
+        );
+        expected.insert("isp".to_string(), "Verizon Business".into());
+        expected.insert("organization".to_string(), "Verizon Business".into());
+
+        b.iter_batched(
+            || (&table, ip, &expected),
+            |(table, ip, expected)| {
+                assert_eq!(
+                    Ok(expected),
+                    table
+                        .find_table_row(
+                            Case::Insensitive,
+                            &[Condition::Equals {
+                                field: "ip",
+                                value: ip.into(),
+                            }],
+                            None,
+                            None,
+                        )
+                        .as_ref()
+                )
+            },
+            BatchSize::SmallInput,
+        );
+    });
+
+    group.bench_function("enrichment_tables/geoip_city", |b| {
+        let table = build("tests/data/GeoIP2-City-Test.mmdb");
+        let ip = "67.43.156.9";
+        let mut expected = BTreeMap::<String, Value>::new();
+        expected.insert("city_name".to_string(), Value::Null);
+        expected.insert("country_code".to_string(), "BT".into());
+        expected.insert("country_name".to_string(), "Bhutan".into());
+        expected.insert("continent_code".to_string(), "AS".into());
+        expected.insert("region_code".to_string(), Value::Null);
+        expected.insert("region_name".to_string(), Value::Null);
+        expected.insert("timezone".to_string(), "Asia/Thimphu".into());
+        expected.insert("latitude".to_string(), Value::from(27.5));
+        expected.insert("longitude".to_string(), Value::from(90.5));
+        expected.insert("postal_code".to_string(), Value::Null);
+        expected.insert("metro_code".to_string(), Value::Null);
+
+        b.iter_batched(
+            || (&table, ip, &expected),
+            |(table, ip, expected)| {
+                assert_eq!(
+                    Ok(expected),
+                    table
+                        .find_table_row(
+                            Case::Insensitive,
+                            &[Condition::Equals {
+                                field: "ip",
+                                value: ip.into(),
+                            }],
+                            None,
+                            None,
+                        )
+                        .as_ref()
+                )
+            },
+            BatchSize::SmallInput,
+        );
+    });
+}