make telem key build-time

bruin-data · Dec 23, 2024 · 6f9a356 · 6f9a356
1 parent d9e5527
commit 6f9a356
Show file tree

Hide file tree

Showing 6 changed files with 34 additions and 39 deletions.
diff --git a/Makefile b/Makefile
@@ -8,6 +8,7 @@ ERROR_COLOR=\033[31;01m
 WARN_COLOR=\033[33;01m
 TELEMETRY_OPTOUT=1
 CURRENT_DIR=$(pwd)
+TELEMETRY_KEY=""
 FILES := $(wildcard *.yml *.txt *.py)
 
 .PHONY: all clean test build tools format pre-commit tools-update
@@ -19,8 +20,7 @@ deps: tools
 
 build: deps
 	@echo "$(OK_COLOR)==> Building the application...$(NO_COLOR)"
-	@CGO_ENABLED=1 go build -v -tags="no_duckdb_arrow" -ldflags="-s -w -X main.Version=$(or $(tag), dev-$(shell git describe --tags --abbrev=0))" -o "$(BUILD_DIR)/$(NAME)" "$(BUILD_SRC)"
-
+	@CGO_ENABLED=1 go build -v -tags="no_duckdb_arrow" -ldflags="-s -w -X main.Version=$(or $(tag), dev-$(shell git describe --tags --abbrev=0)) -X main.TelemetryKey=$(TELEMETRY_KEY)" -o "$(BUILD_DIR)/$(NAME)" "$(BUILD_SRC)"
 
 integration-test: build
 	@touch integration-tests/.git

diff --git a/main.go b/main.go
@@ -12,23 +12,22 @@ import (
 )
 
 var (
-	version = "dev"
-	commit  = ""
+	version      = "dev"
+	commit       = ""
+	TelemetryKey string
 )
 
 func main() {
 	isDebug := false
 	color.NoColor = false
-
 	var optOut bool
 	if os.Getenv("TELEMETRY_OPTOUT") != "" {
 		optOut = true
 	}
 
 	v.Version = version
 	v.Commit = commit
-
-	telemetry.TelemetryKey = os.Getenv("TELEMETRY_KEY")
+	telemetry.TelemetryKey = TelemetryKey
 	telemetry.OptOut = optOut
 	telemetry.AppVersion = v.Version
 	client := telemetry.Init()

diff --git a/pkg/telemetry/telemetry.go b/pkg/telemetry/telemetry.go
@@ -21,13 +21,13 @@ const (
 
 type contextKey string
 
+var TelemetryKey string
 var (
-	TelemetryKey = ""
-	OptOut       = false
-	AppVersion   = ""
-	RunID        = ""
-	client       analytics.Client
-	lock         sync.Mutex
+	OptOut     = false
+	AppVersion = ""
+	RunID      = ""
+	client     analytics.Client
+	lock       sync.Mutex
 )
 
 func Init() io.Closer {

diff --git a/pythonsrc/main.py b/pythonsrc/main.py
@@ -3,7 +3,6 @@
 import logging
 import os
 from parser.main import get_column_lineage, get_tables
-from datetime import datetime
 
 from pathlib import Path
 

diff --git a/pythonsrc/parser/main.py b/pythonsrc/parser/main.py
@@ -116,11 +116,13 @@ def get_column_lineage(query: str, schema: dict, dialect: str):
         nested_schema = schema_dict_to_schema_object(schema)
         try:
             optimized = optimize(parsed, nested_schema, dialect=dialect)
-        except Exception as e:
+        except Exception:
             # try again without dialect, this solves some issues, e.g. https://github.com/tobymao/sqlglot/issues/4538
             optimized = optimize(parsed, nested_schema)
     except Exception as e:
-        logging.error(f"Error optimizing query: {e}, query and schema: { json.dumps({'query': query, 'schema': schema}) }")
+        logging.error(
+            f"Error optimizing query: {e}, query and schema: { json.dumps({'query': query, 'schema': schema}) }"
+        )
         return {"columns": [], "error": str(e)}
 
     result = []

diff --git a/pythonsrc/parser/main_test.py b/pythonsrc/parser/main_test.py
@@ -1424,7 +1424,7 @@
             "raw.Teams": {
                 "Id": "INTEGER",
                 "Name": "STRING",
-            }
+            },
         },
         "expected": [
             {
@@ -1440,12 +1440,13 @@
         ],
         "expected_non_selected": [
             {
-                'name': 'id',
-                'upstream': [{'column': 'id', 'table': 'raw.Organizations'}]},
+                "name": "id",
+                "upstream": [{"column": "id", "table": "raw.Organizations"}],
+            },
             {
-                'name': 'organizationid',
-                'upstream': [{'column': 'organizationid', 'table': 'raw.Users'}]
-            }
+                "name": "organizationid",
+                "upstream": [{"column": "organizationid", "table": "raw.Users"}],
+            },
         ],
     },
     {
@@ -1466,7 +1467,7 @@
                 "CreatedDate": "TIMESTAMP",
                 "CreatedById": "INTEGER",
                 "UpdatedDate": "TIMESTAMP",
-                "UpdatedById": "INTEGER"
+                "UpdatedById": "INTEGER",
             },
             "raw.teams": {
                 "Id": "INTEGER",
@@ -1482,7 +1483,7 @@
                 "CreatedById": "INTEGER",
                 "UpdatedDate": "TIMESTAMP",
                 "UpdatedById": "INTEGER",
-                "IsGuest": "BOOLEAN"
+                "IsGuest": "BOOLEAN",
             },
         },
         "expected": [
@@ -1498,18 +1499,12 @@
             },
         ],
         "expected_non_selected": [
+            {"name": "id", "upstream": [{"column": "id", "table": "raw.teams"}]},
+            {"name": "name", "upstream": [{"column": "name", "table": "raw.teams"}]},
             {
-                'name': 'id',
-                'upstream': [{'column': 'id', 'table': 'raw.teams'}]
-            },
-            {
-                'name': 'name',
-                'upstream': [{'column': 'name', 'table': 'raw.teams'}]
+                "name": "teamid",
+                "upstream": [{"column": "teamid", "table": "raw.teammemberships"}],
             },
-            {
-                'name': 'teamid',
-                'upstream': [{'column': 'teamid', 'table': 'raw.teammemberships'}]
-            }
         ],
     },
 ]
@@ -1519,11 +1514,11 @@
     "query,schema,expected,expected_non_selected,dialect",
     [
         (
-                tc["query"],
-                tc["schema"],
-                tc["expected"],
-                tc["expected_non_selected"],
-                tc["dialect"],
+            tc["query"],
+            tc["schema"],
+            tc["expected"],
+            tc["expected_non_selected"],
+            tc["dialect"],
         )
         for tc in test_cases
     ],