From f0441ee16ac0c168e6c3d1124c8b2ba5c75413c1 Mon Sep 17 00:00:00 2001
From: Timothy Flynn <trflynn89@pm.me>
Date: Fri, 3 Feb 2023 09:58:38 -0500
Subject: [PATCH] LibSQL: Store selected column names in the results for SELECT
 statements

---
 Userland/Libraries/LibSQL/AST/Select.cpp | 45 +++++++++++++++++++++---
 Userland/Libraries/LibSQL/ResultSet.h    |  8 +++++
 2 files changed, 49 insertions(+), 4 deletions(-)

diff --git a/Userland/Libraries/LibSQL/AST/Select.cpp b/Userland/Libraries/LibSQL/AST/Select.cpp
index df06aed99e..5b421dcd11 100644
--- a/Userland/Libraries/LibSQL/AST/Select.cpp
+++ b/Userland/Libraries/LibSQL/AST/Select.cpp
@@ -12,9 +12,35 @@
 
 namespace SQL::AST {
 
+static DeprecatedString result_column_name(ResultColumn const& column, size_t column_index)
+{
+    auto fallback_column_name = [column_index]() {
+        return DeprecatedString::formatted("Column{}", column_index);
+    };
+
+    if (auto const& alias = column.column_alias(); !alias.is_empty())
+        return alias;
+
+    if (column.select_from_expression()) {
+        if (is<ColumnNameExpression>(*column.expression())) {
+            auto const& column_name_expression = verify_cast<ColumnNameExpression>(*column.expression());
+            return column_name_expression.column_name();
+        }
+
+        // FIXME: Generate column names from other result column expressions.
+        return fallback_column_name();
+    }
+
+    VERIFY(column.select_from_table());
+
+    // FIXME: Generate column names from select-from-table result columns.
+    return fallback_column_name();
+}
+
 ResultOr<ResultSet> Select::execute(ExecutionContext& context) const
 {
     NonnullRefPtrVector<ResultColumn> columns;
+    Vector<DeprecatedString> column_names;
 
     auto const& result_column_list = this->result_column_list();
     VERIFY(!result_column_list.is_empty());
@@ -26,27 +52,38 @@ ResultOr<ResultSet> Select::execute(ExecutionContext& context) const
         auto table_def = TRY(context.database->get_table(table_descriptor.schema_name(), table_descriptor.table_name()));
 
         if (result_column_list.size() == 1 && result_column_list[0].type() == ResultType::All) {
+            TRY(columns.try_ensure_capacity(columns.size() + table_def->columns().size()));
+            TRY(column_names.try_ensure_capacity(column_names.size() + table_def->columns().size()));
+
             for (auto& col : table_def->columns()) {
-                columns.append(
+                columns.unchecked_append(
                     create_ast_node<ResultColumn>(
                         create_ast_node<ColumnNameExpression>(table_def->parent()->name(), table_def->name(), col.name()),
                         ""));
+
+                column_names.unchecked_append(col.name());
             }
         }
     }
 
     if (result_column_list.size() != 1 || result_column_list[0].type() != ResultType::All) {
-        for (auto& col : result_column_list) {
+        TRY(columns.try_ensure_capacity(result_column_list.size()));
+        TRY(column_names.try_ensure_capacity(result_column_list.size()));
+
+        for (size_t i = 0; i < result_column_list.size(); ++i) {
+            auto const& col = result_column_list[i];
+
             if (col.type() == ResultType::All) {
                 // FIXME can have '*' for example in conjunction with computed columns
                 return Result { SQLCommand::Select, SQLErrorCode::SyntaxError, "*"sv };
             }
 
-            columns.append(col);
+            columns.unchecked_append(col);
+            column_names.unchecked_append(result_column_name(col, i));
         }
     }
 
-    ResultSet result { SQLCommand::Select };
+    ResultSet result { SQLCommand::Select, move(column_names) };
 
     auto descriptor = adopt_ref(*new TupleDescriptor);
     Tuple tuple(descriptor);
diff --git a/Userland/Libraries/LibSQL/ResultSet.h b/Userland/Libraries/LibSQL/ResultSet.h
index 8dcde2cac2..765fc52e73 100644
--- a/Userland/Libraries/LibSQL/ResultSet.h
+++ b/Userland/Libraries/LibSQL/ResultSet.h
@@ -25,7 +25,14 @@ public:
     {
     }
 
+    ALWAYS_INLINE ResultSet(SQLCommand command, Vector<DeprecatedString> column_names)
+        : m_command(command)
+        , m_column_names(move(column_names))
+    {
+    }
+
     SQLCommand command() const { return m_command; }
+    Vector<DeprecatedString> const& column_names() const { return m_column_names; }
 
     void insert_row(Tuple const& row, Tuple const& sort_key);
     void limit(size_t offset, size_t limit);
@@ -34,6 +41,7 @@ private:
     size_t binary_search(Tuple const& sort_key, size_t low, size_t high);
 
     SQLCommand m_command { SQLCommand::Unknown };
+    Vector<DeprecatedString> m_column_names;
 };
 
 }