lea-cli 0.7.2__tar.gz → 0.7.3__tar.gz

{lea_cli-0.7.2 → lea_cli-0.7.3}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: lea-cli
-Version: 0.7.2
+Version: 0.7.3
 Summary: A minimalist alternative to dbt
 Author: Max Halford
 Author-email: maxhalford25@gmail.com

{lea_cli-0.7.2 → lea_cli-0.7.3}/lea/conductor.py

@@ -115,7 +115,7 @@ class Session:
 
         if self.incremental_field_name is not None:
             self.filterable_table_refs = {
-                table_ref
+                table_ref.replace_dataset(self.write_dataset)
                 for table_ref in scripts
                 if any(
                     field.name == incremental_field_name
@@ -123,7 +123,7 @@ class Session:
                 )
             }
             self.incremental_table_refs = {
-                table_ref
+                table_ref.replace_dataset(self.write_dataset)
                 for table_ref in selected_table_refs | set(existing_audit_tables)
                 if any(
                     field.name == incremental_field_name and FieldTag.INCREMENTAL in field.tags
@@ -140,36 +140,18 @@ class Session:
     def add_write_context_to_table_ref(self, table_ref: TableRef) -> TableRef:
         table_ref = table_ref.replace_dataset(self.write_dataset)
         table_ref = table_ref.add_audit_suffix()
-        if isinstance(self.database_client, databases.BigQueryClient):
-            table_ref = table_ref.replace_project(self.database_client.write_project_id)
         return table_ref
 
     def remove_write_context_from_table_ref(self, table_ref: TableRef) -> TableRef:
         table_ref = table_ref.replace_dataset(self.base_dataset)
         table_ref = table_ref.remove_audit_suffix()
-        if isinstance(self.database_client, databases.BigQueryClient):
-            table_ref = table_ref.replace_project(None)
         return table_ref
 
     def add_context_to_script(self, script: Script) -> Script:
-        # If a script is marked as incremental, it implies that it can be run incrementally. This
-        # means that we have to filter the script's dependencies, as well as filter the output.
-        # This logic is implemented by the script's SQL dialect.
-        if script.table_ref in self.incremental_table_refs:
-            script = dataclasses.replace(
-                script,
-                code=script.sql_dialect.add_dependency_filters(
-                    code=script.code,
-                    incremental_field_name=self.incremental_field_name,
-                    incremental_field_values=self.incremental_field_values,
-                    # One caveat is the dependencies which are not incremental do not have to be
-                    # filtered. Indeed, they are already filtered by the fact that they are
-                    # incremental.
-                    dependencies_to_filter=self.filterable_table_refs - self.incremental_table_refs,
-                ),
-            )
+        def add_context_to_dependency(dependency: TableRef) -> TableRef | None:
+            if dependency.project != script.table_ref.project:
+                return None
 
-        def add_context_to_dependency(dependency: TableRef) -> TableRef:
             if (
                 dependency.replace_dataset(self.base_dataset)
                 in self.selected_table_refs
@@ -180,27 +162,36 @@ class Session:
                 and dependency.replace_dataset(self.base_dataset) in self.scripts
             ):
                 dependency = dependency.add_audit_suffix()
+
             dependency = dependency.replace_dataset(self.write_dataset)
 
-            # Replace occurences of the dataset in script queries clauses to ensure it points
-            # to the right project. If we don't do that and a separate compute project is used,
-            # queries will default to looking for tables in the compute project, which is not
-            # what we want.
-            if (
-                isinstance(self.database_client, databases.BigQueryClient)
-                and dependency.project is None
-            ):
-                dependency = dependency.replace_project(self.database_client.write_project_id)
             return dependency
 
         script = replace_script_dependencies(script=script, replace_func=add_context_to_dependency)
 
+        # If a script is marked as incremental, it implies that it can be run incrementally. This
+        # means that we have to filter the script's dependencies, as well as filter the output.
+        # This logic is implemented by the script's SQL dialect.
+        if script.table_ref.replace_dataset(self.write_dataset) in self.incremental_table_refs:
+            script = dataclasses.replace(
+                script,
+                code=script.sql_dialect.add_dependency_filters(
+                    code=script.code,
+                    incremental_field_name=self.incremental_field_name,
+                    incremental_field_values=self.incremental_field_values,
+                    # One caveat is the dependencies which are not incremental do not have to be
+                    # filtered. Indeed, they are already filtered by the fact that they are
+                    # incremental.
+                    dependencies_to_filter=self.filterable_table_refs - self.incremental_table_refs,
+                ),
+            )
+
         # If the script is not incremental, we're not out of the woods! All scripts are
         # materialized into side-tables which we call "audit" tables. This is the WAP pattern.
         # Therefore, if a script is not incremental, but it depends on an incremental script, we
         # have to modify the script to use both the incremental and non-incremental versions of
         # the dependency. This is handled by the script's SQL dialect.
-        if script.table_ref not in self.incremental_table_refs and self.incremental_table_refs:
+        elif self.incremental_table_refs:
             script = dataclasses.replace(
                 script,
                 code=script.sql_dialect.handle_incremental_dependencies(
@@ -208,10 +199,8 @@ class Session:
                     incremental_field_name=self.incremental_field_name,
                     incremental_field_values=self.incremental_field_values,
                     incremental_dependencies={
-                        self.add_write_context_to_table_ref(
-                            table_ref
-                        ).remove_audit_suffix(): self.add_write_context_to_table_ref(table_ref)
-                        for table_ref in self.incremental_table_refs
+                        incremental_table_ref: incremental_table_ref.add_audit_suffix()
+                        for incremental_table_ref in self.incremental_table_refs
                     },
                 ),
             )
@@ -234,7 +223,8 @@ class Session:
         self.jobs.append(job)
 
         msg = f"{job.status} {script.table_ref}"
-        if script.table_ref in self.incremental_table_refs:
+
+        if script.table_ref.remove_audit_suffix() in self.incremental_table_refs:
             msg += " (incremental)"
         log.info(msg)
 
@@ -291,9 +281,7 @@ class Session:
         to_table_ref = table_ref.remove_audit_suffix()
 
         is_incremental = (
-            self.incremental_field_name is not None
-            and to_table_ref.replace_dataset(self.base_dataset).replace_project(None)
-            in self.incremental_table_refs
+            self.incremental_field_name is not None and to_table_ref in self.incremental_table_refs
         )
         if is_incremental:
             database_job = self.database_client.delete_and_insert(
@@ -345,11 +333,21 @@ def replace_script_dependencies(
 
     """
     code = script.code
+
     for dependency_to_edit in script.dependencies:
-        dependency_to_edit_str = script.sql_dialect.format_table_ref(dependency_to_edit)
         new_dependency = replace_func(dependency_to_edit)
+        if new_dependency is None:
+            continue
+
+        dependency_to_edit_without_project_str = script.sql_dialect.format_table_ref(
+            dependency_to_edit.replace_project(None)
+        )
         new_dependency_str = script.sql_dialect.format_table_ref(new_dependency)
-        code = re.sub(rf"\b{dependency_to_edit_str}\b", new_dependency_str, code)
+        code = re.sub(
+            rf"\b{dependency_to_edit_without_project_str}\b",
+            new_dependency_str,
+            code,
+        )
 
         # We also have to handle the case where the table is referenced to access a field.
         # TODO: refactor this with the above
@@ -394,32 +392,45 @@ def delete_table_refs(
 
 
 class Conductor:
-    def __init__(self, scripts_dir: str, dataset_name: str | None = None):
+    def __init__(
+        self, scripts_dir: str, dataset_name: str | None = None, project_name: str | None = None
+    ):
         # Load environment variables from .env file
         # TODO: is is Pythonic to do this here?
         dotenv.load_dotenv(".env", verbose=True)
 
+        self.warehouse = os.environ["LEA_WAREHOUSE"].lower()
+
         self.scripts_dir = pathlib.Path(scripts_dir)
         if not self.scripts_dir.is_dir():
             raise ValueError(f"Directory {self.scripts_dir} not found")
-        dataset_name = dataset_name or os.environ.get("LEA_BQ_DATASET_NAME")
+
+        if dataset_name is None:
+            if self.warehouse == "bigquery":
+                dataset_name = os.environ.get("LEA_BQ_DATASET_NAME")
         if dataset_name is None:
             raise ValueError("Dataset name could not be inferred")
         self.dataset_name = dataset_name
 
+        if project_name is None:
+            if self.warehouse == "bigquery":
+                project_name = os.environ.get("LEA_BQ_PROJECT_ID")
+        if project_name is None:
+            raise ValueError("Project name could not be inferred")
+        self.project_name = project_name
+
         log.info("📝 Reading scripts")
 
         self.dag = DAGOfScripts.from_directory(
             scripts_dir=self.scripts_dir,
             sql_dialect=BigQueryDialect(),
             dataset_name=self.dataset_name,
+            project_name=self.project_name,
         )
         log.info(f"{len(self.dag.scripts):,d} scripts found")
 
     def make_client(self, dry_run: bool = False, print_mode: bool = False) -> DatabaseClient:
-        warehouse = os.environ["LEA_WAREHOUSE"]
-
-        if warehouse.lower() == "bigquery":
+        if self.warehouse.lower() == "bigquery":
             # Do imports here to avoid loading them all the time
             from google.oauth2 import service_account
 
@@ -448,7 +459,7 @@ class Conductor:
                 print_mode=print_mode,
             )
 
-        raise ValueError(f"Unsupported warehouse {warehouse!r}")
+        raise ValueError(f"Unsupported warehouse {self.warehouse!r}")
 
     def name_user_dataset(self) -> str:
         username = os.environ.get("LEA_USERNAME", getpass.getuser())
@@ -611,10 +622,14 @@ def promote_audit_tables(session: Session):
 
 
 def delete_audit_tables(session: Session):
-    if session.existing_audit_tables:
+    table_refs_to_delete = set(session.existing_audit_tables) | {
+        session.add_write_context_to_table_ref(table_ref)
+        for table_ref in session.selected_table_refs
+    }
+    if table_refs_to_delete:
         log.info("🧹 Deleting audit tables")
         delete_table_refs(
-            table_refs=session.existing_audit_tables,
+            table_refs=table_refs_to_delete,
             database_client=session.database_client,
             executor=concurrent.futures.ThreadPoolExecutor(max_workers=None),
             verbose=False,

{lea_cli-0.7.2 → lea_cli-0.7.3}/lea/dag.py

@@ -19,19 +19,28 @@ class DAGOfScripts(graphlib.TopologicalSorter):
         scripts: list[Script],
         scripts_dir: pathlib.Path,
         dataset_name: str,
+        project_name: str,
     ):
         graphlib.TopologicalSorter.__init__(self, dependency_graph)
         self.dependency_graph = dependency_graph
         self.scripts = {script.table_ref: script for script in scripts}
         self.scripts_dir = scripts_dir
         self.dataset_name = dataset_name
+        self.project_name = project_name
 
     @classmethod
     def from_directory(
-        cls, scripts_dir: pathlib.Path, sql_dialect: SQLDialect, dataset_name: str
+        cls,
+        scripts_dir: pathlib.Path,
+        sql_dialect: SQLDialect,
+        dataset_name: str,
+        project_name: str,
     ) -> DAGOfScripts:
         scripts = read_scripts(
-            scripts_dir=scripts_dir, sql_dialect=sql_dialect, dataset_name=dataset_name
+            scripts_dir=scripts_dir,
+            sql_dialect=sql_dialect,
+            dataset_name=dataset_name,
+            project_name=project_name,
         )
 
         # Fields in the script's code may contain tags. These tags induce assertion tests, which
@@ -53,6 +62,7 @@ class DAGOfScripts(graphlib.TopologicalSorter):
             scripts=scripts,
             scripts_dir=scripts_dir,
             dataset_name=dataset_name,
+            project_name=project_name,
         )
 
     def select(self, *queries: str) -> set[TableRef]:
@@ -111,7 +121,12 @@ class DAGOfScripts(graphlib.TopologicalSorter):
                 return
 
             *schema, name = query.split(".")
-            table_ref = TableRef(dataset=self.dataset_name, schema=tuple(schema), name=name)
+            table_ref = TableRef(
+                dataset=self.dataset_name,
+                schema=tuple(schema),
+                name=name,
+                project=self.project_name,
+            )
             yield table_ref
             if include_ancestors:
                 yield from self.iter_ancestors(node=table_ref)

{lea_cli-0.7.2 → lea_cli-0.7.3}/lea/databases.py

@@ -285,7 +285,9 @@ class BigQueryClient:
         """
         job = self.client.query(query, location=self.location)
         return {
-            BigQueryDialect.parse_table_ref(f"{dataset_name}.{row['table_id']}"): TableStats(
+            BigQueryDialect.parse_table_ref(
+                f"{self.write_project_id}.{dataset_name}.{row['table_id']}"
+            ): TableStats(
                 n_rows=row["row_count"],
                 n_bytes=row["size_bytes"],
                 updated_at=(
@@ -302,9 +304,9 @@ class BigQueryClient:
         """
         job = self.client.query(query, location=self.location)
         return {
-            BigQueryDialect.parse_table_ref(f"{dataset_name}.{table_name}"): [
-                scripts.Field(name=row["column_name"]) for _, row in rows.iterrows()
-            ]
+            BigQueryDialect.parse_table_ref(
+                f"{self.write_project_id}.{dataset_name}.{table_name}"
+            ): [scripts.Field(name=row["column_name"]) for _, row in rows.iterrows()]
             for table_name, rows in job.result()
             .to_dataframe()
             .sort_values(["table_name", "column_name"])
@@ -320,3 +322,19 @@ class BigQueryClient:
             dry_run=self.dry_run,
             **kwargs,
         )
+
+        # The approach we use works best if the tables are clustered by account_slug. This is
+        # because we only need to refresh the data for a subset of accounts, and clustering by
+        # account_slug allows BigQuery to only scan the data for the accounts that need to be
+        # refreshed. This is a good practice in general, but it's particularly important in this
+        # case.
+        # WIP
+        # if (
+        #     script is not None
+        #     and not script.table_ref.is_test
+        #     and script.table_ref.name.endswith("___audit")
+        #     and "account_slug" in {field.name for field in script.fields}
+        # ):
+        #     job_config.clustering_fields = ["account_slug"]
+
+        # return job_config

{lea_cli-0.7.2 → lea_cli-0.7.3}/lea/scripts.py

@@ -72,7 +72,11 @@ class SQLScript:
 
     @classmethod
     def from_path(
-        cls, scripts_dir: pathlib.Path, relative_path: pathlib.Path, sql_dialect: SQLDialect
+        cls,
+        scripts_dir: pathlib.Path,
+        relative_path: pathlib.Path,
+        sql_dialect: SQLDialect,
+        project_name: str,
     ) -> SQLScript:
         # Either the file is a Jinja template
         if relative_path.suffixes == [".sql", ".jinja"]:
@@ -85,7 +89,9 @@ class SQLScript:
             code = (scripts_dir / relative_path).read_text().rstrip().rstrip(";")
 
         return cls(
-            table_ref=TableRef.from_path(scripts_dir=scripts_dir, relative_path=relative_path),
+            table_ref=TableRef.from_path(
+                scripts_dir=scripts_dir, relative_path=relative_path, project_name=project_name
+            ),
             code=code,
             sql_dialect=sql_dialect,
             updated_at=dt.datetime.fromtimestamp(
@@ -107,8 +113,15 @@ class SQLScript:
 
     @functools.cached_property
     def dependencies(self) -> set[TableRef]:
+        def add_default_project(table_ref: TableRef) -> TableRef:
+            if table_ref.project is None:
+                return table_ref.replace_project(self.table_ref.project)
+            return table_ref
+
         return {
-            self.sql_dialect.parse_table_ref(table_ref=sqlglot.exp.table_name(table))
+            add_default_project(
+                self.sql_dialect.parse_table_ref(table_ref=sqlglot.exp.table_name(table))
+            )
             for scope in sqlglot.optimizer.scope.traverse_scope(self.ast)
             for table in scope.tables
             if (
@@ -136,6 +149,7 @@ class SQLScript:
                 dataset=self.table_ref.dataset,
                 schema=("tests",),
                 name=f"{'__'.join(self.table_ref.schema)}__{self.table_ref.name}__{field.name}___{tag.lower().lstrip('#')}",
+                project=self.table_ref.project,
             )
 
         def make_assertion_test(table_ref, field, tag):
@@ -181,14 +195,14 @@ class SQLScript:
     def __rich__(self):
         code = textwrap.dedent(self.code).strip()
         code_with_table_ref = f"""-- {self.table_ref}\n\n{code}\n"""
-        return rich.syntax.Syntax(code_with_table_ref, "sql", line_numbers=True, theme="ansi_dark")
+        return rich.syntax.Syntax(code_with_table_ref, "sql", line_numbers=False, theme="ansi_dark")
 
 
 Script = SQLScript
 
 
 def read_scripts(
-    scripts_dir: pathlib.Path, sql_dialect: SQLDialect, dataset_name: str
+    scripts_dir: pathlib.Path, sql_dialect: SQLDialect, dataset_name: str, project_name: str
 ) -> list[Script]:
     def read_script(path: pathlib.Path) -> Script:
         match tuple(path.suffixes):
@@ -197,6 +211,7 @@ def read_scripts(
                     scripts_dir=scripts_dir,
                     relative_path=path.relative_to(scripts_dir),
                     sql_dialect=sql_dialect,
+                    project_name=project_name,
                 )
             case _:
                 raise ValueError(f"Unsupported script type: {path}")

{lea_cli-0.7.2 → lea_cli-0.7.3}/lea/table_ref.py

@@ -12,33 +12,39 @@ class TableRef:
     dataset: str
     schema: tuple[str, ...]
     name: str
-    project: str | None = None
+    project: str | None
 
     def __str__(self):
         return ".".join(filter(None, [self.project, self.dataset, *self.schema, self.name]))
 
     @classmethod
-    def from_path(cls, scripts_dir: pathlib.Path, relative_path: pathlib.Path) -> TableRef:
+    def from_path(
+        cls, scripts_dir: pathlib.Path, relative_path: pathlib.Path, project_name: str
+    ) -> TableRef:
         parts = list(filter(None, relative_path.parts))
         *schema, filename = parts
         return cls(
             dataset=scripts_dir.name,
             schema=tuple(schema),
-            # Remove the ex
-            name=filename.split(".")[0],
+            name=filename.split(".")[0],  # remove the extension
+            project=project_name,
         )
 
     def replace_dataset(self, dataset: str) -> TableRef:
         return dataclasses.replace(self, dataset=dataset)
 
-    def replace_project(self, project: str | None) -> TableRef:
+    def replace_project(self, project: str) -> TableRef:
         return dataclasses.replace(self, project=project)
 
     def add_audit_suffix(self) -> TableRef:
+        if self.is_audit_table:
+            return self
         return dataclasses.replace(self, name=f"{self.name}{AUDIT_TABLE_SUFFIX}")
 
     def remove_audit_suffix(self) -> TableRef:
-        return dataclasses.replace(self, name=re.sub(rf"{AUDIT_TABLE_SUFFIX}$", "", self.name))
+        if self.is_audit_table:
+            return dataclasses.replace(self, name=re.sub(rf"{AUDIT_TABLE_SUFFIX}$", "", self.name))
+        return self
 
     @property
     def is_audit_table(self) -> bool:

{lea_cli-0.7.2 → lea_cli-0.7.3}/lea/test_big_query.py

@@ -19,7 +19,7 @@ def scripts() -> dict[TableRef, Script]:
         script.table_ref: script
         for script in [
             Script(
-                table_ref=TableRef("read", ("raw",), "users"),
+                table_ref=TableRef("read", ("raw",), "users", "test_project"),
                 code="""
                 SELECT * FROM UNNEST([
                     STRUCT(1 AS id, 'Alice' AS name, 30 AS age),
@@ -30,7 +30,7 @@ def scripts() -> dict[TableRef, Script]:
                 sql_dialect=BigQueryDialect(),
             ),
             Script(
-                table_ref=TableRef("read", ("core",), "users"),
+                table_ref=TableRef("read", ("core",), "users", "test_project"),
                 code="""
                 SELECT
                     id,
@@ -42,7 +42,7 @@ def scripts() -> dict[TableRef, Script]:
                 sql_dialect=BigQueryDialect(),
             ),
             Script(
-                table_ref=TableRef("read", ("analytics",), "n_users"),
+                table_ref=TableRef("read", ("analytics",), "n_users", "test_project"),
                 code="""
                 SELECT COUNT(*)
                 FROM read.core__users
@@ -70,7 +70,9 @@ def test_simple_run(scripts):
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("raw",), "users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("raw",), "users", "test_project")]
+        ).code,
         """
         SELECT * FROM UNNEST([
             STRUCT(1 AS id, 'Alice' AS name, 30 AS age),
@@ -80,10 +82,12 @@ def test_simple_run(scripts):
         """,
     )
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("analytics",), "n_users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("analytics",), "n_users", "test_project")]
+        ).code,
         """
         SELECT COUNT(*)
-        FROM write.core__users___audit
+        FROM `test_project`.write.core__users___audit
         """,
     )
 
@@ -101,29 +105,33 @@ def test_incremental_field(scripts):
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("core",), "users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("core",), "users", "test_project")]
+        ).code,
         """
         SELECT *
         FROM (
             SELECT id, name, age
-            FROM write.raw__users___audit
+            FROM `test_project`.write.raw__users___audit
         )
         WHERE name IN ('Alice')
         """,
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("analytics",), "n_users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("analytics",), "n_users", "test_project")]
+        ).code,
         """
         SELECT COUNT(*) FROM (
             SELECT *
-            FROM write.core__users___audit
+            FROM `test_project`.write.core__users___audit
             WHERE name IN ('Alice')
 
             UNION ALL
 
             SELECT *
-            FROM write.core__users
+            FROM `test_project`.write.core__users
             WHERE name NOT IN ('Alice')
         )
         """,
@@ -136,21 +144,23 @@ def test_incremental_field_but_no_incremental_table_selected(scripts):
         base_dataset="read",
         write_dataset="write",
         scripts=scripts,
-        selected_table_refs={TableRef("read", ("analytics",), "n_users")},
+        selected_table_refs={TableRef("read", ("analytics",), "n_users", "test_project")},
         existing_audit_tables={},
         incremental_field_name="name",
         incremental_field_values={"Alice"},
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("core",), "users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("core",), "users", "test_project")]
+        ).code,
         """
         SELECT
             id,
             -- #INCREMENTAL
             name,
             age
-        FROM write.raw__users
+        FROM `test_project`.write.raw__users
         """,
     )
 
@@ -161,19 +171,21 @@ def test_incremental_field_with_just_incremental_table_selected(scripts):
         base_dataset="read",
         write_dataset="write",
         scripts=scripts,
-        selected_table_refs={TableRef("read", ("core",), "users")},
+        selected_table_refs={TableRef("read", ("core",), "users", "test_project")},
         existing_audit_tables={},
         incremental_field_name="name",
         incremental_field_values={"Alice"},
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("core",), "users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("core",), "users", "test_project")]
+        ).code,
         """
         SELECT *
         FROM (
             SELECT id, name, age
-            FROM write.raw__users
+            FROM `test_project`.write.raw__users
         )
         WHERE name IN ('Alice')
         """,
@@ -188,19 +200,23 @@ def test_incremental_field_with_just_incremental_table_selected_and_materialized
         base_dataset="read",
         write_dataset="write",
         scripts=scripts,
-        selected_table_refs={TableRef("read", ("core",), "users")},
-        existing_audit_tables={TableRef("read", ("raw",), "users"): DUMMY_TABLE_STATS},
+        selected_table_refs={TableRef("read", ("core",), "users", "test_project")},
+        existing_audit_tables={
+            TableRef("read", ("raw",), "users", "test_project"): DUMMY_TABLE_STATS
+        },
         incremental_field_name="name",
         incremental_field_values={"Alice"},
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("core",), "users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("core",), "users", "test_project")]
+        ).code,
         """
         SELECT *
         FROM (
             SELECT id, name, age
-            FROM write.raw__users___audit
+            FROM `test_project`.write.raw__users___audit
         )
         WHERE name IN ('Alice')
         """,
@@ -215,27 +231,29 @@ def test_incremental_field_but_no_incremental_table_selected_and_yet_dependency_
         base_dataset="read",
         write_dataset="write",
         scripts=scripts,
-        selected_table_refs={TableRef("read", ("analytics",), "n_users")},
+        selected_table_refs={TableRef("read", ("analytics",), "n_users", "test_project")},
         existing_audit_tables={
-            TableRef("read", ("core",), "users"): DUMMY_TABLE_STATS,
+            TableRef("read", ("core",), "users", "test_project"): DUMMY_TABLE_STATS,
         },
         incremental_field_name="name",
         incremental_field_values={"Alice"},
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("analytics",), "n_users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("analytics",), "n_users", "test_project")]
+        ).code,
         """
         SELECT COUNT(*)
         FROM (
             SELECT *
-            FROM write.core__users___audit
+            FROM `test_project`.write.core__users___audit
             WHERE name IN ('Alice')
 
             UNION ALL
 
             SELECT *
-            FROM write.core__users
+            FROM `test_project`.write.core__users
             WHERE name NOT IN ('Alice')
         )
         """,
@@ -255,27 +273,29 @@ def test_incremental_field_but_no_incremental_table_selected_and_yet_dependency_
         base_dataset="read",
         write_dataset="write",
         scripts=scripts,
-        selected_table_refs={TableRef("read", ("analytics",), "n_users")},
+        selected_table_refs={TableRef("read", ("analytics",), "n_users", "test_project")},
         existing_audit_tables={
-            TableRef("read", ("core",), "users"): DUMMY_TABLE_STATS,
+            TableRef("read", ("core",), "users", "test_project"): DUMMY_TABLE_STATS,
         },
         incremental_field_name="name",
         incremental_field_values={"Alice"},
     )
 
     assert_queries_are_equal(
-        session.add_context_to_script(scripts[TableRef("read", ("analytics",), "n_users")]).code,
+        session.add_context_to_script(
+            scripts[TableRef("read", ("analytics",), "n_users", "test_project")]
+        ).code,
         """
         SELECT COUNT(*)
         FROM (
             SELECT *
-            FROM `write-project-id`.write.core__users___audit
+            FROM `test_project`.write.core__users___audit
             WHERE name IN ('Alice')
 
             UNION ALL
 
             SELECT *
-            FROM `write-project-id`.write.core__users
+            FROM `test_project`.write.core__users
             WHERE name NOT IN ('Alice')
         )
         """,

{lea_cli-0.7.2 → lea_cli-0.7.3}/lea/test_table_ref.py

@@ -12,11 +12,17 @@ from lea.table_ref import TableRef
     [
         pytest.param(table_ref, expected, id=str(table_ref))
         for table_ref, expected in [
-            (TableRef("my_dataset", ("my_schema",), "my_table"), "my_dataset.my_schema.my_table"),
-            (TableRef("my_dataset", (), "my_table"), "my_dataset.my_table"),
             (
-                TableRef("my_dataset", ("my_schema", "my_subschema"), "my_table"),
-                "my_dataset.my_schema.my_subschema.my_table",
+                TableRef("my_dataset", ("my_schema",), "my_table", "my_project"),
+                "my_project.my_dataset.my_schema.my_table",
+            ),
+            (
+                TableRef("my_dataset", (), "my_table", "my_project"),
+                "my_project.my_dataset.my_table",
+            ),
+            (
+                TableRef("my_dataset", ("my_schema", "my_subschema"), "my_table", "my_project"),
+                "my_project.my_dataset.my_schema.my_subschema.my_table",
             ),
         ]
     ],
@@ -31,16 +37,16 @@ def test_str(table_ref, expected):
         pytest.param(table_ref, expected, id=str(table_ref))
         for table_ref, expected in [
             (
-                TableRef("my_dataset", ("my_schema",), "my_table"),
+                TableRef("my_dataset", ("my_schema",), "my_table", None),
                 "TableRef(dataset='my_dataset', schema=('my_schema',), name='my_table', project=None)",
             ),
             (
-                TableRef("my_dataset", (), "my_table"),
+                TableRef("my_dataset", (), "my_table", None),
                 "TableRef(dataset='my_dataset', schema=(), name='my_table', project=None)",
             ),
             (
-                TableRef("my_dataset", ("my_schema", "my_subschema"), "my_table"),
-                "TableRef(dataset='my_dataset', schema=('my_schema', 'my_subschema'), name='my_table', project=None)",
+                TableRef("my_dataset", ("my_schema", "my_subschema"), "my_table", "my_project"),
+                "TableRef(dataset='my_dataset', schema=('my_schema', 'my_subschema'), name='my_table', project='my_project')",
             ),
         ]
     ],
@@ -52,5 +58,5 @@ def test_repr(table_ref, expected):
 def test_from_path():
     scripts_dir = pathlib.Path("my_dataset")
     relative_path = pathlib.Path("my_schema/my_table.sql")
-    table_ref = TableRef.from_path(scripts_dir, relative_path)
-    assert table_ref == TableRef("my_dataset", ("my_schema",), "my_table")
+    table_ref = TableRef.from_path(scripts_dir, relative_path, "my_project")
+    assert table_ref == TableRef("my_dataset", ("my_schema",), "my_table", "my_project")

{lea_cli-0.7.2 → lea_cli-0.7.3}/pyproject.toml

@@ -5,7 +5,7 @@ name = "lea-cli"
 packages = [
   {include = "lea", from = "."},
 ]
-version = "0.7.2"
+version = "0.7.3"
 
 [tool.poetry.dependencies]
 click = "^8.1.7"