SQL INSERT INTO in python (#88)

data/sql_input_2.sql

@@ -0,0 +1,21 @@
+-- https://cratedb.com/docs/sql-99/en/latest/chapters/01.html
+-- https://www.postgresql.org/docs/16/sql-createtable.html
+-- https://www.postgresql.org/docs/16/sql-insert.html
+-- https://www.postgresql.org/docs/16/sql-select.html
+CREATE TABLE city (
+ name VARCHAR,
+ population INT,
+ timezone INT
+);
+
+INSERT INTO city (name, population, timezone)
+VALUES ('San Francisco', 852469, -8);
+
+INSERT INTO city (name, population, timezone)
+VALUES ('New York', 8405837, -5);
+
+SELECT
+ name,
+ population,
+ timezone
+FROM city;

data/sql_input_3.sql

@@ -0,0 +1,21 @@
+-- https://cratedb.com/docs/sql-99/en/latest/chapters/01.html
+-- https://www.postgresql.org/docs/16/sql-createtable.html
+-- https://www.postgresql.org/docs/16/sql-insert.html
+-- https://www.postgresql.org/docs/16/sql-select.html
+CREATE TABLE city (
+ name VARCHAR,
+ population INT,
+ timezone INT
+);
+
+INSERT INTO city (name, timezone)
+VALUES ('San Francisco', -8);
+
+INSERT INTO city (name, population)
+VALUES ('New York', 8405837);
+
+SELECT
+ name,
+ population,
+ timezone
+FROM city;

data/sql_output_0.json

@@ -1,3 +1 @@
-{
- "table_name": ["city"]
-}
+[{ "table_name": "city" }]

data/sql_output_1.json

@@ -1,3 +1 @@
-{
- "table_name": ["city", "town"]
-}
+[{ "table_name": "city" }, { "table_name": "town" }]

data/sql_output_2.json

@@ -0,0 +1,4 @@
+[
+ { "name": "San Francisco", "population": 852469, "timezone": -8 },
+ { "name": "New York", "population": 8405837, "timezone": -5 }
+]

data/sql_output_3.json

@@ -0,0 +1,4 @@
+[
+ { "name": "San Francisco", "population": null, "timezone": -8 },
+ { "name": "New York", "population": 8405837, "timezone": null }
+]

snippets/python/sql_test.py

@@ -1,31 +1,119 @@
 
+import dataclasses
 import json
+import typing
+
+
+@dataclasses.dataclass(frozen=True)
+class SQLState:
+ state: dict
+
+ def read_table_meta(self, table_name: str) -> dict:
+ return self.state.get(table_name, {}).get("metadata", {})
+
+ def read_table_rows(self, table_name: str) -> list[dict]:
+ return self.state.get(table_name, {}).get("rows", [])
+
+ def read_information_schema(self) -> list[dict]:
+ return [data["metadata"] for data in self.state.values()]
+
+ def write_table_meta(self, table_name: str, data: dict):
+ state = self.state
+ table = state.get(table_name, {})
+ metadata = table.get("metadata", {})
+ metadata.update(data)
+ table["metadata"] = metadata
+ state[table_name] = table
+ return self.__class__(state)
+
+ def write_table_rows(self, table_name: str, data: dict):
+ state = self.state
+ table = state.get(table_name, {})
+ rows = table.get("rows", [])
+ rows.append(data)
+ table["rows"] = rows
+ state[table_name] = table
+ return self.__class__(state)
+
+
+class SQLType:
+ @staticmethod
+ def varchar(data) -> str:
+ data_str = str(data).strip()
+ if data_str.startswith("'") or data_str.startswith('"'):
+ data_str = data_str[1:]
+ if data_str.endswith("'") or data_str.endswith('"'):
+ data_str = data_str[:-1]
+ return data_str
+
+ @staticmethod
+ def int(data) -> int:
+ return int(data.strip())
+
+
+sql_type_map = {
+ "VARCHAR": SQLType.varchar,
+ "INT": SQLType.int,
+}
+
+
+class SQLFunctions:
+ @staticmethod
+ def create_table(state: SQLState, *args, table_schema="public") -> typing.Tuple[list, SQLState]:
+ output: list[dict] = []
+ table_name = args[2]
 
+ # get columns
+ columns = {}
+ columns_str = " ".join(args[3:]).replace("(", "").replace(")", "").strip()
+ if columns_str:
+ # fmt: off
+ columns = {
+ column.strip().split(" ")[0]: column.strip().split(" ")[1]
+ for column in columns_str.split(",")
+ }
+ # fmt: on
 
-class SQL:
- data: dict = {}
-
- def __init__(self) -> None:
- self.data = {}
-
- def information_schema_tables(self) -> list[dict]:
- return [data["metadata"] for data in self.data.values()]
-
- def create_table(self, *args, table_schema="public") -> dict:
- table_name = args[2]
- if not self.data.get(table_name):
- self.data[table_name] = {
- "metadata": {
+ if not state.read_table_meta(table_name):
+ state = state.write_table_meta(
+ table_name,
+ {
  "table_name": table_name,
  "table_schema": table_schema,
+ "colums": columns,
  },
- }
- return {}
+ )
+ return (output, state)
+
+ @staticmethod
+ def insert_into(state: SQLState, *args) -> typing.Tuple[list, SQLState]:
+ output: list[dict] = []
+ table_name = args[2]
+
+ values_index = None
+ for i, arg in enumerate(args):
+ if arg == "VALUES":
+ values_index = i
+ if values_index is None:
+ raise ValueError("VALUES not found")
+
+ keys = " ".join(args[3:values_index]).replace("(", "").replace(")", "").split(",")
+ keys = [key.strip() for key in keys]
+ values = " ".join(args[values_index + 1 :]).replace("(", "").replace(")", "").split(",")
+ values = [value.strip() for value in values]
+ key_value_map = dict(zip(keys, values))
 
- create_table.sql = "CREATE TABLE"
+ data = {}
+ if metadata := state.read_table_meta(table_name):
+ for key, value in key_value_map.items():
+ data[key] = sql_type_map[metadata["colums"][key]](value)
+ state = state.write_table_rows(table_name, data)
 
- def select(self, *args) -> dict:
- output = {}
+ return (output, state)
+
+ @staticmethod
+ def select(state: SQLState, *args) -> typing.Tuple[list, SQLState]:
+ output: list[dict] = []
 
  from_index = None
  where_index = None
@@ -34,49 +122,59 @@ def select(self, *args) -> dict:
  from_index = i
  if arg == "WHERE":
  where_index = i
+ if from_index is None:
+ raise ValueError("FROM not found")
 
  # get select keys by getting the slice of args before FROM
  select_keys = " ".join(args[1:from_index]).split(",")
+ select_keys = [key.strip() for key in select_keys]
 
  # get where keys by getting the slice of args after WHERE
  from_value = args[from_index + 1]
 
- # consider "information_schema.tables" a special case until
- # we figure out why its so different from the others
+ # `information_schema.tables` is a special case
  if from_value == "information_schema.tables":
- target = self.information_schema_tables()
-
- # fmt: off
- output = {
- key: [
- value for data in target
- for key, value in data.items()
- if key in select_keys
- ]
- for key in select_keys
- }
- # fmt: on
-
- return output
-
- select.sql = "SELECT"
-
- sql_map = {
- create_table.sql: create_table,
- select.sql: select,
- }
-
- def run(self, input_sql: list[str]) -> list[str]:
- output = {}
-
- for line in input_sql:
- if not line.startswith("--"):
- words = line.split(" ")
- for i in reversed(range(len(words))):
- key = " ".join(words[:i])
- if func := self.sql_map.get(key):
- output = func(self, *words)
- break
-
- return [json.dumps(output)]
+ data = state.read_information_schema()
+ else:
+ data = state.read_table_rows(from_value)
+
+ output = []
+ for datum in data:
+ # fmt: off
+ output.append({
+ key: datum.get(key)
+ for key in select_keys
+ })
+ # fmt: on
+
+ return (output, state)
+
+
+sql_function_map: dict[str, typing.Callable] = {
+ "CREATE TABLE": SQLFunctions.create_table,
+ "SELECT": SQLFunctions.select,
+ "INSERT INTO": SQLFunctions.insert_into,
+}
+
+
+def run_sql(input_sql: list[str]) -> list[str]:
+ output = []
+ state = SQLState(state={})
+
+ # remove comments
+ input_sql = [line.strip() for line in input_sql if not line.startswith("--")]
+
+ # re-split on semi-colons
+ input_sql = " ".join(input_sql).split(";")
+
+ # iterate over each line of sql
+ for line in input_sql:
+ words = line.split(" ")
+ for i in reversed(range(len(words) + 1)):
+ key = " ".join(words[:i]).strip()
+ if func := sql_function_map.get(key):
+ output, state = func(state, *[word for word in words if word])
+ break
+
+ return [json.dumps(output)]