Implement support for Nested Subqueries and Column Projection

Author: colloceo

minidb/database.py

@@ -210,26 +210,44 @@ def execute_query(self, query_string):
 
             if cmd_type == 'SELECT':
                 limit = parsed.get('limit')
-                # SELECT always reads from current state
+                condition = parsed.get('condition')
+                columns = parsed.get('columns', '*')
+                
+                # Handle nested subquery in WHERE clause
+                if condition and isinstance(condition['value'], str) and condition['value'].startswith('(') and 'SELECT' in condition['value'].upper():
+                    sub_sql = condition['value'][1:-1].strip()
+                    sub_res = self.execute_query(sub_sql)
+                    
+                    if isinstance(sub_res, list):
+                        # Flatten subquery result to a simple list of values
+                        if sub_res and isinstance(sub_res[0], dict):
+                            flattened = []
+                            for row in sub_res:
+                                flattened.extend(row.values())
+                            condition['value'] = flattened
+                        else:
+                            condition['value'] = sub_res
+
+                # Execution Logic
                 if self.transaction.in_transaction and table_name in self.transaction.staging_area:
                     # Read from staging area if modified in transaction
                     staged_data = self.transaction.staging_area[table_name]['data']
-                    condition = parsed.get('condition')
                     res = staged_data
                     if condition:
-                        # Apply WHERE filter to staged data
                         res = [row for row in staged_data 
                                 if table._matches_condition(row, condition['column'], 
                                                            condition['operator'], condition['value'])]
                     if limit:
-                        return res[:limit]
-                    return res
+                        res = res[:limit]
                 else:
                     # Read from disk
-                    condition = parsed.get('condition')
                     if condition:
-                        return table.select_where(condition['column'], condition['operator'], condition['value'], limit=limit)
-                    return table.select_all(limit=limit)
+                        res = table.select_where(condition['column'], condition['operator'], condition['value'], limit=limit)
+                    else:
+                        res = table.select_all(limit=limit)
+                
+                # Apply column projection
+                return table.project_columns(res, columns)
 
             if cmd_type == 'DELETE':
                 condition = parsed['condition']

minidb/parser.py

@@ -8,7 +8,7 @@ def __init__(self):
             'CREATE': re.compile(r"CREATE\s+TABLE\s+(\w+)\s*\((.*)\)", re.IGNORECASE),
             'INSERT': re.compile(r"INSERT\s+INTO\s+(\w+)\s+VALUES\s*\((.*)\)", re.IGNORECASE),
             'SELECT_JOIN': re.compile(r"SELECT\s+\*\s+FROM\s+(\w+)\s+JOIN\s+(\w+)\s+ON\s+(\w+)\.(\w+)\s*=\s*(\w+)\.(\w+)", re.IGNORECASE),
-            'SELECT': re.compile(r"SELECT\s+\*\s+FROM\s+(\w+)(?:\s+WHERE\s+(\w+)\s*(>=|<=|!=|>|<|=)\s*(.*?))?(?:\s+LIMIT\s+(\d+))?$", re.IGNORECASE),
+            'SELECT': re.compile(r"SELECT\s+(\*|[\w,\s]+)\s+FROM\s+(\w+)(?:\s+WHERE\s+(\w+)\s*(>=|<=|!=|>|<|=|\s+IN\s+)\s*(.*?))?(?:\s+LIMIT\s+(\d+))?$", re.IGNORECASE),
             'DELETE': re.compile(r"DELETE\s+FROM\s+(\w+)\s+WHERE\s+(\w+)\s*(>=|<=|!=|>|<|=)\s*(.*)", re.IGNORECASE),
             'UPDATE': re.compile(r"UPDATE\s+(\w+)\s+SET\s+(\w+)\s*=\s*(.*)\s+WHERE\s+(\w+)\s*(>=|<=|!=|>|<|=)\s*(.*)", re.IGNORECASE),
             'ALTER_TABLE': re.compile(r"ALTER\s+TABLE\s+(\w+)\s+ADD\s+(\w+)\s+(\w+)", re.IGNORECASE),
@@ -90,19 +90,21 @@ def _process_match(self, cmd_type, match):
             }
 
         elif cmd_type == 'SELECT':
-            table_name = match.group(1)
+            columns = match.group(1).strip()
+            table_name = match.group(2)
             condition = None
-            if match.group(2) and match.group(3) and match.group(4):
+            if match.group(3) and match.group(4) and match.group(5):
                 condition = {
-                    'column': match.group(2),
-                    'operator': match.group(3),
-                    'value': self._infer_type(match.group(4).strip())
+                    'column': match.group(3),
+                    'operator': match.group(4).strip().upper(),
+                    'value': self._infer_type(match.group(5).strip())
                 }
             return {
                 'type': 'SELECT',
+                'columns': columns,
                 'table': table_name,
                 'condition': condition,
-                'limit': int(match.group(5)) if len(match.groups()) >= 5 and match.group(5) else None
+                'limit': int(match.group(6)) if len(match.groups()) >= 6 and match.group(6) else None
             }
 
         elif cmd_type == 'SELECT_JOIN':

minidb/table.py

@@ -226,9 +226,25 @@ def _evaluate_condition(self, row_value, operator, target_value):
         if operator == '<': return row_value < target_value
         if operator == '>=': return row_value >= target_value
         if operator == '<=': return row_value <= target_value
+        if operator == 'IN':
+            if isinstance(target_value, list):
+                return row_value in target_value
+            return row_value == target_value
 
         return False
 
+    def project_columns(self, rows, columns_str):
+        """Helper to return only specific columns from a list of rows."""
+        if columns_str == '*':
+            return rows
+        
+        cols = [c.strip() for c in columns_str.split(',')]
+        result = []
+        for row in rows:
+            new_row = {k: row[k] for k in cols if k in row}
+            result.append(new_row)
+        return result
+
     def add_column(self, column_name, column_type=None):
         """Adds a new column to the table schema and updates all existing rows."""
         # Check if column already exists