Discover schema from __table__ data classes

labs.yml

@@ -34,3 +34,6 @@ commands:
         description: Publish the dashboard after creating by setting to `yes` or `y`.
       - name: open-browser
         description: Open the dashboard in the browser after creating by setting to `yes` or `y`.
+
+  - name: deploy-schema
+    description: Create schema from dataclasses in the given folder

pyproject.toml

@@ -25,6 +25,9 @@ dependencies = [
   "sqlglot>=22.3.1"
 ]
 
+[project.optional-dependencies]
+df = ["sqlframe~=3.3.1"]
+
 [project.urls]
 Documentation = "https://github.com/databrickslabs/lsql#readme"
 Issues = "https://github.com/databrickslabs/lsql/issues"
@@ -43,6 +46,7 @@ path = "src/databricks/labs/lsql/__about__.py"
 
 [tool.hatch.envs.default]
 dependencies = [
+    "databricks-labs-lsql[df]",
     "coverage[toml]>=6.5",
     "pytest",
     "pylint",

src/databricks/labs/lsql/discovery.py

@@ -0,0 +1,45 @@
+import ast
+from pathlib import Path
+
+
+class DataclassTableFinder(ast.NodeVisitor):
+    def __init__(self):
+        self.tables = []
+
+    def visit_ClassDef(self, node):
+        # Check if the class is a dataclass
+        is_dataclass = any(isinstance(decorator, ast.Name) and decorator.id == 'dataclass'
+                           for decorator in node.decorator_list)
+
+        # Look for __table__ assignment in class body
+        has_table_field = any(isinstance(n, ast.Assign) and
+                              any(isinstance(t, ast.Name) and t.id == '__table__' for t in n.targets)
+                              for n in node.body)
+
+        # If both conditions are met, store the class name
+        if is_dataclass and has_table_field:
+            self.tables.append(node.name)
+
+        # Continue visiting the rest of the AST
+        self.generic_visit(node)
+
+
+class Scanner:
+    def __init__(self, start: Path):
+        self._start = start
+
+    def find_all(self):
+        for f in self._start.glob('**/*.py'): # TODO: skip virtual environments
+            yield from self._find_dataclasses_with_table(f)
+
+    def _find_dataclasses_with_table(self, path: Path):
+        # Parse the source code into an AST
+        tree = ast.parse(path.read_text())
+
+        # Create a finder instance and visit the parsed tree
+        finder = DataclassTableFinder()
+        finder.visit(tree)
+
+        # Return the list of dataclasses with __table__ field
+        return finder.tables
+

tests/unit/test_discovery.py

@@ -0,0 +1,9 @@
+from pathlib import Path
+
+from databricks.labs.lsql.discovery import Scanner
+
+
+def test_finds():
+    s = Scanner(Path('/Users/serge.smertin/git/labs/lsql/tests'))
+    x = list(s.find_all())
+    assert x == ['Nested']

tests/unit/test_structs.py

@@ -15,6 +15,8 @@ class Foo:
 
 @dataclass
 class Nested:
+    __table__ = 'x'
+
     foo: Foo
     mapping: dict[str, int]
     array: list[int]