tinybird 0.0.1.dev26__py3-none-any.whl → 0.0.1.dev28__py3-none-any.whl

tinybird/tb/modules/create.py

@@ -4,9 +4,9 @@ from pathlib import Path
 from typing import Optional
 
 import click
-import requests
 
 from tinybird.client import TinyB
+from tinybird.prompts import create_prompt, mock_prompt
 from tinybird.tb.modules.cicd import init_cicd
 from tinybird.tb.modules.cli import cli
 from tinybird.tb.modules.common import _generate_datafile, check_user_token_with_client, coro, generate_datafile
@@ -15,15 +15,11 @@ from tinybird.tb.modules.datafile.fixture import build_fixture_name, persist_fix
 from tinybird.tb.modules.exceptions import CLIException
 from tinybird.tb.modules.feedback_manager import FeedbackManager
 from tinybird.tb.modules.llm import LLM
+from tinybird.tb.modules.llm_utils import extract_xml, parse_xml
 from tinybird.tb.modules.local_common import get_tinybird_local_client
 
 
 @cli.command()
-@click.option(
-    "--demo",
-    is_flag=True,
-    help="Demo data and files to get started",
-)
 @click.option(
     "--data",
     type=click.Path(exists=True),
@@ -39,13 +35,12 @@ from tinybird.tb.modules.local_common import get_tinybird_local_client
 @click.option(
     "--folder",
     default=".",
-    type=click.Path(exists=True, file_okay=False),
+    type=click.Path(exists=False, file_okay=False),
     help="Folder where datafiles will be placed",
 )
 @click.option("--rows", type=int, default=10, help="Number of events to send")
 @coro
 async def create(
-    demo: bool,
     data: Optional[str],
     prompt: Optional[str],
     folder: Optional[str],
@@ -53,6 +48,10 @@ async def create(
 ) -> None:
     """Initialize a new project."""
     folder = folder or getcwd()
+    folder_path = Path(folder)
+    if not folder_path.exists():
+        folder_path.mkdir()
+
     try:
         config = CLIConfig.get_project_config(folder)
         tb_client = config.get_client()
@@ -71,132 +70,103 @@ async def create(
                 )
                 return
         local_client = await get_tinybird_local_client(folder)
-        click.echo(FeedbackManager.highlight(message="\n» Creating new project structure..."))
-        await project_create(local_client, tb_client, user_token, data, prompt, folder)
-        click.echo(FeedbackManager.success(message="✓ Scaffolding completed!\n"))
 
-        click.echo(FeedbackManager.highlight(message="\n» Creating CI/CD files for GitHub and GitLab..."))
-        init_git(folder)
-        await init_cicd(data_project_dir=os.path.relpath(folder))
-        click.echo(FeedbackManager.success(message="✓ Done!\n"))
+        if not validate_project_structure(folder):
+            click.echo(FeedbackManager.highlight(message="\n» Creating new project structure..."))
+            create_project_structure(folder)
+            click.echo(FeedbackManager.success(message="✓ Scaffolding completed!\n"))
 
-        click.echo(FeedbackManager.highlight(message="\n» Building fixtures..."))
+        click.echo(FeedbackManager.highlight(message="\n» Creating resources..."))
+        datasources_created = await create_resources(local_client, tb_client, user_token, data, prompt, folder)
+        click.echo(FeedbackManager.success(message="✓ Done!\n"))
 
-        if demo:
-            # Users datasource
-            ds_name = "users"
-            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
-            datasource_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/b48fb9c92825ed27c04e3104b9e871e1/raw/1f33c20eefbabc4903f38e234329e028d8ef9def/users.datasource"
-            )
-            datasource_path.write_text(datasource_content)
-            click.echo(FeedbackManager.info(message=f"✓ /datasources/{ds_name}.datasource"))
+        if not already_has_cicd(folder):
+            click.echo(FeedbackManager.highlight(message="\n» Creating CI/CD files for GitHub and GitLab..."))
+            init_git(folder)
+            await init_cicd(data_project_dir=os.path.relpath(folder))
+            click.echo(FeedbackManager.success(message="✓ Done!\n"))
 
-            # Users fixtures
-            fixture_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/8e8f66a39d7576ce3a2529bf773334a8/raw/9cab636767990e97d44a141867e5f226e992de8c/users.ndjson"
-            )
-            fixture_name = build_fixture_name(
-                datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
-            )
-            persist_fixture(fixture_name, fixture_content)
-            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
-
-            # Events datasource
-            ds_name = "events"
-            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
-            datasource_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/f8ca37b5b1f6707c75206b618de26bc9/raw/cd625da0dcd1ba8de29f12bc1c8600b9ff7c809c/events.datasource"
-            )
-            datasource_path.write_text(datasource_content)
-            click.echo(FeedbackManager.info(message=f"✓ /datasources/{ds_name}.datasource"))
+        if validate_fixtures(folder) and datasources_created:
+            click.echo(FeedbackManager.highlight(message="\n» Generating fixtures..."))
 
-            # Events fixtures
-            fixture_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/859ab9439c17e77241d0c14a5a532809/raw/251f2f3f00a968f8759ec4068cebde915256b054/events.ndjson"
-            )
-            fixture_name = build_fixture_name(
-                datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
-            )
-            persist_fixture(fixture_name, fixture_content)
-            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
-
-            # Create sample endpoint
-            pipe_name = "api_token_usage"
-            pipe_path = Path(folder) / "endpoints" / f"{pipe_name}.pipe"
-            pipe_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/68ecc47472c2b754b0ae0c1187022963/raw/52cc3aa3afdf939e58d43355bfe4ddc739989ddd/api_token_usage.pipe"
-            )
-            pipe_path.write_text(pipe_content)
-            click.echo(FeedbackManager.info(message=f"✓ /endpoints/{pipe_name}.pipe"))
-
-            # Create sample test
-            test_name = "api_token_usage"
-            test_path = Path(folder) / "tests" / f"{test_name}.yaml"
-            test_content = fetch_gist_content(
-                "https://gist.githubusercontent.com/gnzjgo/e58620bbb977d6f42f1d0c2a7b46ac8f/raw/a3a1cd0ce3a90bcd2f6dfce00da51e6051443612/api_token_usage.yaml"
-            )
-            test_path.write_text(test_content)
-            click.echo(FeedbackManager.info(message=f"✓ /tests/{test_name}.yaml"))
-
-        elif data:
-            ds_name = os.path.basename(data.split(".")[0])
-            data_content = Path(data).read_text()
-            datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
-            fixture_name = build_fixture_name(
-                datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
-            )
-            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
-            persist_fixture(fixture_name, data_content)
-        elif prompt and user_token:
-            datasource_files = [f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")]
-            for datasource_file in datasource_files:
-                datasource_path = Path(folder) / "datasources" / datasource_file
-                llm = LLM(user_token=user_token, client=tb_client)
-                datasource_name = datasource_path.stem
-                datasource_content = datasource_path.read_text()
-                has_json_path = "`json:" in datasource_content
-                if has_json_path:
-                    sql = await llm.generate_sql_sample_data(schema=datasource_content, rows=rows, prompt=prompt)
-                    result = await local_client.query(f"{sql} FORMAT JSON")
-                    data = result.get("data", [])
-                    fixture_name = build_fixture_name(
-                        datasource_path.absolute().as_posix(), datasource_name, datasource_content
-                    )
-                    if data:
-                        persist_fixture(fixture_name, data)
-                        click.echo(FeedbackManager.info(message=f"✓ /fixtures/{datasource_name}"))
+            if data:
+                ds_name = os.path.basename(data.split(".")[0])
+                data_content = Path(data).read_text()
+                datasource_path = Path(folder) / "datasources" / f"{ds_name}.datasource"
+                fixture_name = build_fixture_name(
+                    datasource_path.absolute().as_posix(), ds_name, datasource_path.read_text()
+                )
+                click.echo(FeedbackManager.info(message=f"✓ /fixtures/{ds_name}"))
+                persist_fixture(fixture_name, data_content, folder)
+            elif prompt and user_token:
+                datasource_files = [f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")]
+                for datasource_file in datasource_files:
+                    datasource_path = Path(folder) / "datasources" / datasource_file
+                    llm = LLM(user_token=user_token, client=tb_client)
+                    datasource_name = datasource_path.stem
+                    datasource_content = datasource_path.read_text()
+                    has_json_path = "`json:" in datasource_content
+                    if has_json_path:
+                        prompt = f"<datasource_schema>{datasource_content}</datasource_schema>\n<user_input>{prompt}</user_input>"
+                        response = await llm.ask(prompt, system_prompt=mock_prompt(rows))
+                        sql = extract_xml(response, "sql")
+                        sql = sql.split("FORMAT")[0]
+                        result = await local_client.query(f"{sql} FORMAT JSON")
+                        data = result.get("data", [])
+                        fixture_name = build_fixture_name(
+                            datasource_path.absolute().as_posix(), datasource_name, datasource_content
+                        )
+                        if data:
+                            persist_fixture(fixture_name, data, folder)
+                            click.echo(FeedbackManager.info(message=f"✓ /fixtures/{datasource_name}"))
 
-        click.echo(FeedbackManager.success(message="✓ Done!\n"))
+            click.echo(FeedbackManager.success(message="✓ Done!\n"))
     except Exception as e:
         click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
 
 
-async def project_create(
+PROJECT_PATHS = ("datasources", "endpoints", "materializations", "copies", "sinks", "fixtures", "tests")
+
+
+def validate_project_structure(folder: str) -> bool:
+    return all((Path(folder) / path).exists() for path in PROJECT_PATHS)
+
+
+def validate_fixtures(folder: str) -> bool:
+    datasource_files = [f for f in os.listdir(Path(folder) / "datasources") if f.endswith(".datasource")]
+    return len(datasource_files) > 0
+
+
+def already_has_cicd(folder: str) -> bool:
+    ci_cd_paths = (".gitlab", ".github")
+    return any((Path(folder) / path).exists() for path in ci_cd_paths)
+
+
+def create_project_structure(folder: str):
+    folder_path = Path(folder)
+    for x in PROJECT_PATHS:
+        try:
+            f = folder_path / x
+            f.mkdir()
+        except FileExistsError:
+            pass
+        click.echo(FeedbackManager.info_path_created(path=x))
+
+
+async def create_resources(
     local_client: TinyB,
-    client: TinyB,
+    tb_client: TinyB,
     user_token: Optional[str],
     data: Optional[str],
     prompt: Optional[str],
     folder: str,
 ):
-    project_paths = ["datasources", "endpoints", "materializations", "copies", "sinks", "fixtures", "tests"]
     force = True
-    for x in project_paths:
-        try:
-            f = Path(folder) / x
-            f.mkdir()
-        except FileExistsError:
-            pass
-        click.echo(FeedbackManager.info_path_created(path=x))
-
+    folder_path = Path(folder)
     if data:
-        path = Path(folder) / data
+        path = folder_path / data
         format = path.suffix.lstrip(".")
-        try:
-            await _generate_datafile(str(path), local_client, format=format, force=force)
-        except Exception as e:
-            click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
+        await _generate_datafile(str(path), local_client, format=format, force=force)
         name = data.split(".")[0]
         generate_pipe_file(
             f"{name}_endpoint",
@@ -208,21 +178,68 @@ TYPE ENDPOINT
             """,
             folder,
         )
+        return True
     elif prompt and user_token:
-        try:
-            llm = LLM(user_token=user_token, client=client)
-            result = await llm.create_project(prompt)
-            for ds in result.datasources:
-                content = ds.content.replace("```", "")
-                generate_datafile(
-                    content, filename=f"{ds.name}.datasource", data=None, _format="ndjson", force=force, folder=folder
-                )
+        datasource_paths = [
+            Path(folder) / "datasources" / f
+            for f in os.listdir(Path(folder) / "datasources")
+            if f.endswith(".datasource")
+        ]
+        pipes_paths = [
+            Path(folder) / "endpoints" / f for f in os.listdir(Path(folder) / "endpoints") if f.endswith(".pipe")
+        ]
+        resources_xml = "\n".join(
+            [
+                f"<resource><type>{resource_type}</type><name>{resource_name}</name><content>{resource_content}</content></resource>"
+                for resource_type, resource_name, resource_content in [
+                    ("datasource", ds.stem, ds.read_text()) for ds in datasource_paths
+                ]
+                + [
+                    (
+                        "pipe",
+                        pipe.stem,
+                        pipe.read_text(),
+                    )
+                    for pipe in pipes_paths
+                ]
+            ]
+        )
+        llm = LLM(user_token=user_token, client=tb_client)
+        result = await llm.ask(prompt, system_prompt=create_prompt(resources_xml))
+        result = extract_xml(result, "response")
+        resources = parse_xml(result, "resource")
+        datasources = []
+        pipes = []
+        for resource_xml in resources:
+            resource_type = extract_xml(resource_xml, "type")
+            name = extract_xml(resource_xml, "name")
+            content = extract_xml(resource_xml, "content")
+            resource = {
+                "name": name,
+                "content": content,
+            }
+            if resource_type.lower() == "datasource":
+                datasources.append(resource)
+            elif resource_type.lower() == "pipe":
+                pipes.append(resource)
+
+        for ds in datasources:
+            content = ds["content"].replace("```", "")
+            filename = f"{ds['name']}.datasource"
+            generate_datafile(
+                content,
+                filename=filename,
+                data=None,
+                _format="ndjson",
+                force=force,
+                folder=folder,
+            )
 
-            for pipe in result.pipes:
-                content = pipe.content.replace("```", "")
-                generate_pipe_file(pipe.name, content, folder)
-        except Exception as e:
-            click.echo(FeedbackManager.error(message=f"Error: {str(e)}"))
+        for pipe in pipes:
+            content = pipe["content"].replace("```", "")
+            generate_pipe_file(pipe["name"], content, folder)
+
+        return len(datasources) > 0
 
 
 def init_git(folder: str):
@@ -250,9 +267,3 @@ def generate_pipe_file(name: str, content: str, folder: str):
     with open(f"{f}", "w") as file:
         file.write(content)
     click.echo(FeedbackManager.info_file_created(file=f.relative_to(folder)))
-
-
-def fetch_gist_content(url: str) -> str:  # TODO: replace this with a function that fetches the content from a repo
-    response = requests.get(url)
-    response.raise_for_status()
-    return response.text

tinybird/tb/modules/datafile/build.py

@@ -722,6 +722,34 @@ async def process(
     except Exception as e:
         raise click.ClickException(str(e))
 
+    # datasource
+    # {
+    #     "resource": "datasources",
+    #     "resource_name": name,
+    #     "version": doc.version,
+    #     "params": params,
+    #     "filename": filename,
+    #     "deps": deps,
+    #     "tokens": doc.tokens,
+    #     "shared_with": doc.shared_with,
+    #     "filtering_tags": doc.filtering_tags,
+    # }
+    # pipe
+    # {
+    #     "resource": "pipes",
+    #     "resource_name": name,
+    #     "version": doc.version,
+    #     "filename": filename,
+    #     "name": name + version,
+    #     "nodes": nodes,
+    #     "deps": [x for x in set(deps)],
+    #     "tokens": doc.tokens,
+    #     "description": description,
+    #     "warnings": doc.warnings,
+    #     "filtering_tags": doc.filtering_tags,
+    # }
+
+    # r is essentially a Datasource or a Pipe in dict shape, like in the comment above
     for r in res:
         resource_name = r["resource_name"]
         warnings = r.get("warnings", [])

tinybird/tb/modules/datafile/common.py

@@ -1340,6 +1340,7 @@ def parse(
         "export_compression": assign_var("export_compression"),
         "export_write_strategy": assign_var("export_write_strategy"),
         "export_kafka_topic": assign_var("export_kafka_topic"),
+        "forward_query": sql("forward_query"),
     }
 
     engine_vars = set()

tinybird/tb/modules/datafile/fixture.py

@@ -31,22 +31,26 @@ def build_fixture_name(filename: str, datasource_name: str, datasource_content:
     return f"{datasource_name}_{hash_str}"
 
 
-def get_fixture_dir() -> Path:
-    fixture_dir = Path("fixtures")
+def get_fixture_dir(folder: str) -> Path:
+    fixture_dir = Path(folder) / "fixtures"
     if not fixture_dir.exists():
         fixture_dir.mkdir()
     return fixture_dir
 
 
-def persist_fixture(fixture_name: str, data: Union[List[Dict[str, Any]], str], format="ndjson") -> Path:
-    fixture_dir = get_fixture_dir()
+def persist_fixture(fixture_name: str, data: Union[List[Dict[str, Any]], str], folder: str, format="ndjson") -> Path:
+    fixture_dir = get_fixture_dir(folder)
     fixture_file = fixture_dir / f"{fixture_name}.{format}"
     fixture_file.write_text(data if isinstance(data, str) else format_data_to_ndjson(data))
     return fixture_file
 
 
-def load_fixture(fixture_name: str, format="ndjson") -> Union[Path, None]:
-    fixture_dir = get_fixture_dir()
+def load_fixture(
+    fixture_name: str,
+    folder: str,
+    format="ndjson",
+) -> Union[Path, None]:
+    fixture_dir = get_fixture_dir(folder)
     fixture_file = fixture_dir / f"{fixture_name}.{format}"
     if not fixture_file.exists():
         return None

tinybird/tb/modules/datafile/parse_pipe.py

@@ -45,6 +45,8 @@ def parse_pipe(
         for node in doc.nodes:
             sql = node.get("sql", "")
             if sql.strip()[0] == "%":
+                # Note(eclbg): not sure what test_mode is for. I think it does something like using placeholder values
+                # for the variables in the template.
                 sql, _, variable_warnings = render_sql_template(sql[1:], test_mode=True, name=node["name"])
                 doc.warnings = variable_warnings
             # it'll fail with a ModuleNotFoundError when the toolset is not available but it returns the parsed doc

tinybird/tb/modules/datasource.py

@@ -453,7 +453,7 @@ async def generate_datasource(ctx: Context, connector: str, filenames, force: bo
     """Generate a data source file based on a sample CSV file from local disk or url"""
     client: TinyB = ctx.ensure_object(dict)["client"]
 
-    _connector: Optional["Connector"] = None
+    _connector: Optional[Connector] = None
     if connector:
         load_connector_config(ctx, connector, False, check_uninstalled=False)
         if connector not in ctx.ensure_object(dict):