feat: add get_tranco_rank_category function

Author: hatamiarash7

README.md

@@ -255,6 +255,18 @@ D SELECT get_tranco_rank('cloudflare.com') as rank;
 └─────────┘
 ```
 
+You can use the `get_tranco_rank_category` function to retrieve the category utility column that gives you the rank category of the domain. The `category` value is on a log10 scale with half steps (e.g. top 1k, top 5k, top 10k, top 50k, top 100k, top 500k, top 1M, top 5m, etc.) with each rank excluding the previous (e.g. top 5k is actually 4k domains, excluding top 1k).
+
+```sql
+D SELECT get_tranco_rank_category('microsoft.com') as category;
+┌──────────┐
+│ category │
+│ varchar  │
+├──────────┤
+│ top1k    │
+└──────────┘
+```
+
 ### Get Extension Version
 
 You can use the `netquack_version` function to get the version of the extension.
@@ -273,7 +285,6 @@ D select * from netquack_version();
 
 - [ ] Create a `TableFunction` for `extract_query_parameters` that return each key-value pair as a row.
 - [ ] Save Tranco data as Parquet
-- [ ] Create Rank category for Tranco ( `top1k` , `top5k`, `top10k`, `top100k`, `top500k`, `top1m` )
 - [ ] Implement GeoIP functionality
 - [ ] Add new functions to work with IPs
 - [ ] Return default value for `get_tranco_rank`

src/functions/get_tranco.cpp

@@ -124,6 +124,7 @@ namespace duckdb
                            " WHEN rank <= 1000 THEN 'top1k'"
                            " WHEN rank <= 5000 THEN 'top5k'"
                            " WHEN rank <= 10000 THEN 'top10k'"
+                           " WHEN rank <= 50000 THEN 'top50k'"
                            " WHEN rank <= 100000 THEN 'top100k'"
                            " WHEN rank <= 500000 THEN 'top500k'"
                            " WHEN rank <= 1000000 THEN 'top1m'"
@@ -192,5 +193,44 @@ namespace duckdb
                 }
             }
         }
+
+        // Function to get the Tranco rank category of a domain
+        void GetTrancoRankCategoryFunction (DataChunk &args, ExpressionState &state, Vector &result)
+        {
+            auto &db = *state.GetContext ().db;
+            netquack::LoadPublicSuffixList (db, false);
+            Connection con (db);
+
+            auto table_exists = con.Query ("SELECT 1 FROM information_schema.tables WHERE table_name = 'tranco_list'");
+
+            if (table_exists->RowCount () == 0)
+            {
+                throw std::runtime_error ("Tranco table not found. Download it first using `SELECT update_tranco(true);`");
+            }
+
+            // Extract the input from the arguments
+            auto &input_vector = args.data[0];
+            auto result_data   = FlatVector::GetData<string_t> (result);
+
+            for (idx_t i = 0; i < args.size (); i++)
+            {
+                auto input = input_vector.GetValue (i).ToString ();
+                std::transform (input.begin (), input.end (), input.begin (), ::tolower);
+
+                try
+                {
+                    auto query = "SELECT category FROM tranco_list WHERE domain = '" + input + "'";
+
+                    auto query_result = con.Query (query);
+                    auto category     = query_result->RowCount () > 0 ? query_result->GetValue (0, 0) : Value ();
+
+                    result_data[i] = StringVector::AddString (result, category.ToString ());
+                }
+                catch (const std::exception &e)
+                {
+                    result_data[i] = "Error extracting tranco category: " + std::string (e.what ());
+                }
+            }
+        }
     } // namespace netquack
 } // namespace duckdb

src/functions/get_tranco.hpp

@@ -11,5 +11,8 @@ namespace duckdb
 
         // Function to get the Tranco rank of a domain
         void GetTrancoRankFunction (DataChunk &args, ExpressionState &state, Vector &result);
+
+        // Function to get the Tranco rank category of a domain
+        void GetTrancoRankCategoryFunction (DataChunk &args, ExpressionState &state, Vector &result);
     } // namespace netquack
 } // namespace duckdb

src/netquack_extension.cpp

@@ -99,6 +99,13 @@ namespace duckdb
             netquack::GetTrancoRankFunction);
         ExtensionUtil::RegisterFunction (instance, get_tranco_rank_function);
 
+        auto get_tranco_rank_category_function = ScalarFunction (
+            "get_tranco_rank_category",
+            { LogicalType::VARCHAR },
+            LogicalType::VARCHAR,
+            netquack::GetTrancoRankCategoryFunction);
+        ExtensionUtil::RegisterFunction (instance, get_tranco_rank_category_function);
+
         auto version_function = TableFunction (
             "netquack_version",
             {},