improve migration class

Query.php

@@ -1,6 +1,7 @@
 <?php
 namespace DatabaseHelper;
 
+use DatabaseHelper\enums\Aggregation;
 use DatabaseHelper\enums\Join;
 use DatabaseHelper\enums\Order;
 use http\Exception\InvalidArgumentException;
@@ -12,6 +13,7 @@ class Query
 
     protected Schema $schema;
     protected array $columns = ['*'];
+    protected array $aggregations = [];
     protected array $joins = [];
     public array $orderBy;
 
@@ -36,7 +38,7 @@ class Query
         return $this;
     }
 
-    public function join(Schema $other, Join $join): Query {
+    public function join(Join $join, Schema $other): Query {
         $foreignKey = null;
         if($this->schema->existsReference($other))
             $foreignKey = $this->schema->foreignKeys[$other->name];
@@ -65,17 +67,20 @@ class Query
     }
 
     public function toSql(): string {
+        // Merge any aggregations with the standard columns.
+        $selectColumns = $this->columns;
+        if ($this->hasAggregations())
+            $selectColumns = array_merge($selectColumns, $this->aggregations);
+
         // Build the SELECT clause.
-        $columns = implode(", ", $this->columns);
+        $columns = implode(", ", $selectColumns);
         $primaryTable = $this->schema->name;
         $sqlStatement = "SELECT $columns FROM $primaryTable";
 
         // Append join clauses, if any.
-        if ($this->isJoined()) {
-            foreach ($this->joins as $join) {
+        if ($this->isJoined())
+            foreach ($this->joins as $join)
                 $sqlStatement .= " " . $join['type']->toString() . " NATURAL JOIN " . $join['table'];
-            }
-        }
 
         // Append the WHERE clause if conditions exist.
         if ($this->isConditioned()) {
@@ -92,18 +97,46 @@ class Query
         return esc_sql($sqlStatement);
     }
 
-    public function query(): array {
+    public function aggregate(string $col, string $alias, Aggregation $func): Query {
+        if ($col != '*')
+            $this->schema->requireColumn($col);
+        $this->aggregations[] = strtoupper($func->toString()) . "($col) AS $alias";
+        return $this;
+    }
+
+    public function hasAggregations(): bool {
+        return !empty($this->aggregations);
+    }
+
+    public function query(): mixed {
         global $wpdb;
         $query = $this->toSql();
         $results = $wpdb->get_results($query, ARRAY_A);
-        return $this->castResults($results);
+        return $this->formatResults($results);
     }
 
-    protected function castResults(array $results): array {
-        foreach ($results as &$row)
+    protected function formatResults(array $results) {
+        $formatted = [];
+
+        foreach ($results as $row) {
+            // Apply type casting to each column in each row.
             foreach ($row as $column => &$value)
                 if (isset($this->columnTypes[$column]))
                     $value = $this->columnTypes[$column]->valCast($value);
-        return $results;
+            // Use the primary key for row indexing
+            $primaryKey = $this->schema->primaryKey();
+            $formatted[$row[$primaryKey]] = $row;
+        }
+
+        if (count($formatted) === 1) {
+            // Unpack single row results
+            $row = array_shift($formatted);
+            if (count($row) === 1)
+                // Unpack single column results
+                return array_shift($row);
+            return $row;
+        }
+
+        return $formatted;
     }
 }