From 514d28d6fc5e14c71f67f4c4f93da8e524b4a204 Mon Sep 17 00:00:00 2001
From: unytics <paul.marcombes@unytics.io>
Date: Fri, 10 May 2024 19:48:01 +0000
Subject: [PATCH] fix for aggregate function

---
 bigfun/bigfunctions.py          | 6 +++++-
 bigfun/cli.py                   | 2 +-
 bigfun/templates/bigfunction.md | 2 +-
 bigfun/utils.py                 | 4 ++--
 bigfunctions/README.md          | 1 +
 5 files changed, 10 insertions(+), 5 deletions(-)

diff --git a/bigfun/bigfunctions.py b/bigfun/bigfunctions.py
index 01a1eafe0..8ce9598c7 100644
--- a/bigfun/bigfunctions.py
+++ b/bigfun/bigfunctions.py
@@ -91,6 +91,10 @@ def dataset(self):
             self._dataset = self.bigquery.get_dataset(f'{self.project}.{self.dataset_name}')
         return self._dataset
 
+    @property
+    def location(self):
+        return self.dataset.location
+
     def test(self):
         # WARNING: TO CHANGE THIS AND DEPLOY A PYTHON FUNCTION HERE WE NEED TO HAVE A REMOTE CONNECTION PER DATASET AS users between dataset and remote connection are identical
         if self.config['type'] == 'function_py':
@@ -135,7 +139,7 @@ def deploy(self):
         template = jinja2.Template(open(template_file, encoding='utf-8').read())
         query = template.render(**self.config)
         print_info('Creating function in dataset')
-        self.bigquery.query(query)
+        self.bigquery.query(query, location=self.location)
         print_success(f'successfully created {self.project}.{self.dataset_name}.{self.name}')
 
     @property
diff --git a/bigfun/cli.py b/bigfun/cli.py
index 80e090d98..60c331704 100644
--- a/bigfun/cli.py
+++ b/bigfun/cli.py
@@ -154,7 +154,7 @@ def deploy(bigfunction, project, dataset):
         if len(datasets) > 1:
             with multiprocessing.Pool(processes=8) as pool:
                 pool.map(
-                    BigFunction.deploy,
+                    bf.BigFunction.deploy,
                     [
                         bf.BigFunction(name, project=project, dataset=dataset)
                         for dataset in datasets[1:]
diff --git a/bigfun/templates/bigfunction.md b/bigfun/templates/bigfunction.md
index 6d4e039b7..a98aca1db 100644
--- a/bigfun/templates/bigfunction.md
+++ b/bigfun/templates/bigfunction.md
@@ -73,7 +73,7 @@ description: "BigFunction {{ name }}: {{ description.split('\n')[0] }}"
     ```sql
     {% if example.with_clause is defined %}
     with sample_data as (
-      {{ example.with_clause | indent(2) }}
+      {{ example.with_clause | indent(6) }}
     )
     {% endif %}
     {% if type == 'procedure' %}call{% elif type == 'table_function' %}select * from{% else %}select{% endif %} {{ project }}.{{ dataset }}.{{ name }}({% for argument in example.arguments %}{{ argument | replace('{BIGFUNCTIONS_DATASET}', dataset) | replace('\n', '\n      ') }}{% if not loop.last %}, {% endif %}{% endfor %}){% if type == 'procedure' %};{% elif 'output' in bigfunction and type != 'table_function' %} as {{ output.name }}{% endif %}
diff --git a/bigfun/utils.py b/bigfun/utils.py
index 6e3940823..b559a8743 100644
--- a/bigfun/utils.py
+++ b/bigfun/utils.py
@@ -137,9 +137,9 @@ def get_dataset(self, dataset):
         ]
         return dataset
 
-    def query(self, query):
+    def query(self, query, **kwargs):
         try:
-            return self.client.query(query).result()
+            return self.client.query(query, **kwargs).result()
         except google.api_core.exceptions.Forbidden as e:
             handle_error("Access Denied", e.message)
         except (
diff --git a/bigfunctions/README.md b/bigfunctions/README.md
index 736fe0812..1cf9afb53 100644
--- a/bigfunctions/README.md
+++ b/bigfunctions/README.md
@@ -124,6 +124,7 @@ BigFunctions are open-source BigQuery routines that give you **SQL-superpowers**
 - [<code>format_percentage(first_number, second_number, nb_decimals)</code>](format_percentage.md): Return `first_number / second_number` as a formatted percentage
 - [<code>quantize_into_fixed_width_bins(value, min_bound, max_bound, nb_bins)</code>](quantize_into_fixed_width_bins.md): Get the `bin_range` in which belongs `value`
 - [<code>quantize_into_bins(value, bin_bounds)</code>](quantize_into_bins.md): Get the `bin_range` in which belongs `value`
+- [<code>weighted_average(element, weight)</code>](weighted_average.md): Returns the weigthed average elements.
 
 
 ## ✨ Transform string