Split in-memory with fs tables

Author: brunovcosta

json_sql/apply.py

@@ -1,5 +1,5 @@
 from typing import List, Dict, Optional
-from .tables import TablesSnapshot
+from .tables import ITablesSnapshot
 from .ast import (
     Expression,
     StringExpression,
@@ -203,7 +203,7 @@ def apply_select_fields(fields: List[SelectField], data: List[dict], ctx: dict):
 
 
 def apply_from(
-    from_part: Optional[From], tables: TablesSnapshot, ctx: dict
+    from_part: Optional[From], tables: ITablesSnapshot, ctx: dict
 ) -> List[dict]:
     if from_part is None:
         return [{}]
@@ -225,7 +225,7 @@ def apply_from(
     return data
 
 
-def apply_select(select: Select, tables: TablesSnapshot, ctx: dict):
+def apply_select(select: Select, tables: ITablesSnapshot, ctx: dict):
     data = apply_from(select.from_part, tables, ctx)
     if select.where_part:
         data = apply_where(select.where_part, data, ctx)
@@ -240,7 +240,7 @@ def apply_select(select: Select, tables: TablesSnapshot, ctx: dict):
     return data
 
 
-def apply_command(command: Command, tables: TablesSnapshot, ctx: dict):
+def apply_command(command: Command, tables: ITablesSnapshot, ctx: dict):
     if isinstance(command, Select):
         return apply_select(command, tables, ctx)
     else:

json_sql/cli.py

@@ -1,18 +1,14 @@
 from typing import List
 from pathlib import Path
-from json_sql.tables import TablesSnapshot, Table
+from json_sql.tables import FileSystemTables, Table
 from json import loads
 from eval import eval_sql
 
 
-def query(code: str, json_paths: List[Path], ctx: dict) -> List[dict]:
-    tables = TablesSnapshot(
-        tables=[
-            Table(name=path.stem, columns=[], data=loads(path.read_text()))
-            for path in json_paths
-        ]
-    )
-    return eval_sql(code=code, tables=tables, ctx=ctx)
+def query(code: str, workdir: Path, ctx: dict):
+    tables = FileSystemTables(workdir=workdir)
+    result = eval_sql(code=code, tables=tables, ctx=ctx)
+    return result
 
 
 def help(): ...

json_sql/eval.py

@@ -1,10 +1,10 @@
 from .lexer import scan
 from .parser import parse
 from .apply import apply_command
-from .tables import TablesSnapshot
+from .tables import ITablesSnapshot
 
 
-def eval_sql(code: str, tables: TablesSnapshot, ctx: dict):
+def eval_sql(code: str, tables: ITablesSnapshot, ctx: dict):
     tokens = scan(code)
     ast = parse(tokens)
     result = apply_command(command=ast, tables=tables, ctx=ctx)

json_sql/eval_test.py

@@ -1,12 +1,12 @@
 from unittest import TestCase
 from .eval import eval_sql
-from .tables import TablesSnapshot, Table, Column
+from .tables import InMemoryTables
 
 
 class TestEvalSQL(TestCase):
     def test_eval_sql(self):
         code = "select 1+1"
-        tables = TablesSnapshot(
+        tables = InMemoryTables(
             tables=[],
         )
         ctx = {}
@@ -15,7 +15,7 @@ def test_eval_sql(self):
 
     def test_eval_select_alias(self):
         code = "select 1+1 as a"
-        tables = TablesSnapshot(
+        tables = InMemoryTables(
             tables=[],
         )
         ctx = {}

json_sql/tables.py

@@ -1,5 +1,8 @@
 from typing import List, Optional
 from pydantic import BaseModel
+from pathlib import Path
+from abc import ABC, abstractmethod
+import json
 
 
 class ForeignKey(BaseModel):
@@ -20,11 +23,38 @@ class Table(BaseModel):
     data: List[dict] = []
 
 
-class TablesSnapshot(BaseModel):
+class ITablesSnapshot(ABC):
+    @abstractmethod
+    def get_table(self, name: str) -> Optional[Table]:
+        raise NotImplementedError("get_table method must be implemented")
+
+
+class InMemoryTables(BaseModel, ITablesSnapshot):
     tables: List[Table]
 
     def get_table(self, name: str) -> Optional[Table]:
         for table in self.tables:
             if table.name == name:
                 return table
         return None
+
+
+class FileSystemTables(ITablesSnapshot):
+    workdir: Path
+
+    def __init__(self, workdir: Path):
+        self.workdir = workdir
+
+    def get_table(self, name: str):
+        table_path = self.workdir / f"{name}.json"
+        if not table_path.exists():
+            raise FileNotFoundError(f"File {table_path} does not exist")
+        rows = json.loads(table_path.read_text())
+        columns = set()
+        for row in rows:
+            assert isinstance(row, dict), f"Row {row} is not a dictionary"
+            for key, value in row.items():
+                if key not in columns:
+                    col = Column(name=key, type=type(value).__name__)
+                    columns.add(col)
+        return Table(name=name, columns=list(columns), data=rows)