snips

snippets/python/sql_test.py

@@ -2,14 +2,81 @@
 import json
 
 
-def run_sql(input_sql: list[str]) -> list[str]:
- output = {}
- lines = []
+class SQL:
+ data: dict = {}
 
- for line in input_sql:
- if line.startswith("--"):
- continue
- lines.append(line)
+ def __init__(self) -> None:
+ self.data = {}
 
- return [json.dumps(output)]
+ 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": {
+ "table_name": table_name,
+ "table_schema": table_schema,
+ },
+ }
+ return {}
+
+ create_table.sql = "CREATE TABLE"
+
+ def select(self, *args) -> dict:
+ output = {}
+
+ from_index = None
+ where_index = None
+ for i, arg in enumerate(args):
+ if arg == "FROM":
+ from_index = i
+ if arg == "WHERE":
+ where_index = i
+
+ # get select keys by getting the slice of args before FROM
+ select_keys = " ".join(args[1:from_index]).split(",")
+
+ # 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
+ 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)]