Repository: tobymao/sqlglot

Issue:
Support catalog for Trino SQL execution
Related to this issue: https://github.com/tobymao/sqlglot/issues/93

Is there an easy way, to support Trino SQL execution with Catalog and Schema?
Example code: 
```
from sqlglot.executor import execute
tables = {
    "some_catalog.some_schema.some_table": [
        {"id": 1, "price": 1.0},
        {"id": 2, "price": 2.0},
        {"id": 3, "price": 3.0},
    ]
}

result = execute(
    """
    SELECT
      *
    FROM some_catalog.some_schema.some_table s
    """,
    tables=tables
)

print(result)
```


```
sqlglot.errors.ExecuteError: Step 'Scan: s (4324464000)' failed: 'NoneType' object has no attribute 'range_reader'
```


Discussion:

Patch:
diff --git a/sqlglot/executor/__init__.py b/sqlglot/executor/__init__.py
index a67c155c..017d5bcc 100644
--- a/sqlglot/executor/__init__.py
+++ b/sqlglot/executor/__init__.py
@@ -14,9 +14,10 @@ from sqlglot import maybe_parse
 from sqlglot.errors import ExecuteError
 from sqlglot.executor.python import PythonExecutor
 from sqlglot.executor.table import Table, ensure_tables
+from sqlglot.helper import dict_depth
 from sqlglot.optimizer import optimize
 from sqlglot.planner import Plan
-from sqlglot.schema import ensure_schema
+from sqlglot.schema import ensure_schema, flatten_schema, nested_get, nested_set
 
 logger = logging.getLogger("sqlglot")
 
@@ -52,10 +53,15 @@ def execute(
     tables_ = ensure_tables(tables)
 
     if not schema:
-        schema = {
-            name: {column: type(table[0][column]).__name__ for column in table.columns}
-            for name, table in tables_.mapping.items()
-        }
+        schema = {}
+        flattened_tables = flatten_schema(tables_.mapping, depth=dict_depth(tables_.mapping))
+
+        for keys in flattened_tables:
+            table = nested_get(tables_.mapping, *zip(keys, keys))
+            assert table is not None
+
+            for column in table.columns:
+                nested_set(schema, [*keys, column], type(table[0][column]).__name__)
 
     schema = ensure_schema(schema, dialect=read)
 
diff --git a/sqlglot/schema.py b/sqlglot/schema.py
index a308ae4d..502d691f 100644
--- a/sqlglot/schema.py
+++ b/sqlglot/schema.py
@@ -128,12 +128,12 @@ class AbstractMappingSchema(t.Generic[T]):
                 if raise_on_missing:
                     raise SchemaError(f"Ambiguous mapping for {table}: {message}.")
                 return None
-        return self._nested_get(parts, raise_on_missing=raise_on_missing)
+        return self.nested_get(parts, raise_on_missing=raise_on_missing)
 
-    def _nested_get(
+    def nested_get(
         self, parts: t.Sequence[str], d: t.Optional[t.Dict] = None, raise_on_missing=True
     ) -> t.Optional[t.Any]:
-        return _nested_get(
+        return nested_get(
             d or self.mapping,
             *zip(self.supported_table_args, reversed(parts)),
             raise_on_missing=raise_on_missing,
@@ -209,7 +209,7 @@ class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
 
         parts = self.table_parts(normalized_table)
 
-        _nested_set(
+        nested_set(
             self.mapping,
             tuple(reversed(parts)),
             normalized_column_mapping,
@@ -226,7 +226,7 @@ class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
         if not only_visible or not self.visible:
             return list(schema)
 
-        visible = self._nested_get(self.table_parts(table_), self.visible)
+        visible = self.nested_get(self.table_parts(table_), self.visible)
         return [col for col in schema if col in visible]  # type: ignore
 
     def get_column_type(self, table: exp.Table | str, column: exp.Column | str) -> exp.DataType:
@@ -259,12 +259,12 @@ class MappingSchema(AbstractMappingSchema[t.Dict[str, str]], Schema):
 
         normalized_mapping: t.Dict = {}
         for keys in flattened_schema:
-            columns = _nested_get(schema, *zip(keys, keys))
+            columns = nested_get(schema, *zip(keys, keys))
             assert columns is not None
 
             normalized_keys = [self._normalize_name(key) for key in keys]
             for column_name, column_type in columns.items():
-                _nested_set(
+                nested_set(
                     normalized_mapping,
                     normalized_keys + [self._normalize_name(column_name)],
                     column_type,
@@ -365,7 +365,7 @@ def flatten_schema(
     return tables
 
 
-def _nested_get(
+def nested_get(
     d: t.Dict, *path: t.Tuple[str, str], raise_on_missing: bool = True
 ) -> t.Optional[t.Any]:
     """
@@ -390,15 +390,15 @@ def _nested_get(
     return d
 
 
-def _nested_set(d: t.Dict, keys: t.Sequence[str], value: t.Any) -> t.Dict:
+def nested_set(d: t.Dict, keys: t.Sequence[str], value: t.Any) -> t.Dict:
     """
     In-place set a value for a nested dictionary
 
     Example:
-        >>> _nested_set({}, ["top_key", "second_key"], "value")
+        >>> nested_set({}, ["top_key", "second_key"], "value")
         {'top_key': {'second_key': 'value'}}
 
-        >>> _nested_set({"top_key": {"third_key": "third_value"}}, ["top_key", "second_key"], "value")
+        >>> nested_set({"top_key": {"third_key": "third_value"}}, ["top_key", "second_key"], "value")
         {'top_key': {'third_key': 'third_value', 'second_key': 'value'}}
 
     Args:


Test patch:
diff --git a/tests/test_executor.py b/tests/test_executor.py
index e61d64ab..8f3e2e6f 100644
--- a/tests/test_executor.py
+++ b/tests/test_executor.py
@@ -604,3 +604,21 @@ class TestExecutor(unittest.TestCase):
         result = execute("SELECT A AS A FROM X", tables={"x": [{"a": 1}]})
         self.assertEqual(result.columns, ("A",))
         self.assertEqual(result.rows, [(1,)])
+
+    def test_nested_table_reference(self):
+        tables = {
+            "some_catalog": {
+                "some_schema": {
+                    "some_table": [
+                        {"id": 1, "price": 1.0},
+                        {"id": 2, "price": 2.0},
+                        {"id": 3, "price": 3.0},
+                    ]
+                }
+            }
+        }
+
+        result = execute("SELECT * FROM some_catalog.some_schema.some_table s", tables=tables)
+
+        self.assertEqual(result.columns, ("id", "price"))
+        self.assertEqual(result.rows, [(1, 1.0), (2, 2.0), (3, 3.0)])