Skip to content

Commit

Permalink
chore(data-warehouse): fix column typing and remove unnecessary UI (#…
Browse files Browse the repository at this point in the history
…23130)

* fix column typing and remove unneessary UI

* typing
  • Loading branch information
EDsCODE authored and timgl committed Jun 27, 2024
1 parent 675c50c commit eedbb81
Show file tree
Hide file tree
Showing 4 changed files with 133 additions and 166 deletions.
93 changes: 4 additions & 89 deletions frontend/src/scenes/data-warehouse/external/TableData.tsx
Original file line number Diff line number Diff line change
Expand Up @@ -3,13 +3,11 @@ import { LemonButton, LemonModal } from '@posthog/lemon-ui'
import { useActions, useValues } from 'kea'
import { capitalizeFirstLetter } from 'kea-forms'
import { humanFriendlyDetailedTime } from 'lib/utils'
import { Dispatch, SetStateAction, useEffect, useState } from 'react'
import { Dispatch, SetStateAction } from 'react'
import { DatabaseTable } from 'scenes/data-management/database/DatabaseTable'

import { HogQLQueryEditor } from '~/queries/nodes/HogQLQuery/HogQLQueryEditor'
import { DatabaseSchemaTable, HogQLQuery, NodeKind } from '~/queries/schema'
import { DatabaseSchemaTable } from '~/queries/schema'

import { ViewLinkModal } from '../ViewLinkModal'
import { dataWarehouseSceneLogic } from './dataWarehouseSceneLogic'

export function TableData(): JSX.Element {
Expand All @@ -18,30 +16,13 @@ export function TableData(): JSX.Element {
isEditingSavedQuery,
inEditSchemaMode,
editSchemaIsLoading,
databaseLoading,
} = useValues(dataWarehouseSceneLogic)
const {
deleteDataWarehouseSavedQuery,
deleteDataWarehouseTable,
setIsEditingSavedQuery,
updateDataWarehouseSavedQuery,
toggleEditSchemaMode,
updateSelectedSchema,
saveSchema,
cancelEditSchema,
} = useActions(dataWarehouseSceneLogic)
const [localQuery, setLocalQuery] = useState<HogQLQuery>()
const [isDeleteModalOpen, setIsDeleteModalOpen] = useState(false)
const { setIsEditingSavedQuery, toggleEditSchemaMode, updateSelectedSchema, saveSchema, cancelEditSchema } =
useActions(dataWarehouseSceneLogic)

const isExternalTable = table?.type === 'data_warehouse'
const isManuallyLinkedTable = isExternalTable && !table.source

useEffect(() => {
if (table && table.type === 'view') {
setLocalQuery(table.query)
}
}, [table])

return (
<div className="border rounded p-3 bg-bg-light">
{table ? (
Expand Down Expand Up @@ -91,11 +72,6 @@ export function TableData(): JSX.Element {
Edit schema
</LemonButton>
)}
{table.type === 'view' && (
<LemonButton type="primary" onClick={() => setIsEditingSavedQuery(true)}>
Edit
</LemonButton>
)}
</div>
)}
</div>
Expand Down Expand Up @@ -139,71 +115,10 @@ export function TableData(): JSX.Element {
/>
</div>
)}

{table.type === 'view' && isEditingSavedQuery && (
<div className="mt-2">
<span className="card-secondary">Update View Definition</span>
<HogQLQueryEditor
query={{
kind: NodeKind.HogQLQuery,
// TODO: Use `hogql` tag?
query: `${localQuery && localQuery.query}`,
}}
onChange={(queryInput) => {
setLocalQuery({
kind: NodeKind.HogQLQuery,
query: queryInput,
})
}}
editorFooter={(hasErrors, error, isValidView) => (
<LemonButton
className="ml-2"
onClick={() => {
localQuery &&
updateDataWarehouseSavedQuery({
...table,
query: localQuery,
})
}}
loading={databaseLoading}
type="primary"
center
disabledReason={
hasErrors
? error ?? 'Query has errors'
: !isValidView
? 'All fields must have an alias'
: ''
}
data-attr="hogql-query-editor-save-as-view"
>
Save as view
</LemonButton>
)}
/>
</div>
)}
</>
) : (
<div className="px-4 py-3 h-100 col-span-2 flex justify-center items-center" />
)}
{table && (
<DeleteTableModal
table={table}
isOpen={isDeleteModalOpen}
setIsOpen={setIsDeleteModalOpen}
onDelete={() => {
if (table) {
if (table.type === 'view') {
deleteDataWarehouseSavedQuery(table.id)
} else {
deleteDataWarehouseTable(table.id)
}
}
}}
/>
)}
<ViewLinkModal />
</div>
)
}
Expand Down
85 changes: 66 additions & 19 deletions posthog/warehouse/models/datawarehouse_saved_query.py
Original file line number Diff line number Diff line change
@@ -1,14 +1,20 @@
import re
from sentry_sdk import capture_exception
from django.core.exceptions import ValidationError
from django.db import models
from typing import Optional, Any

from posthog.hogql.database.database import Database
from posthog.hogql.database.models import SavedQuery
from posthog.hogql.database.models import SavedQuery, FieldOrTable
from posthog.hogql import ast
from posthog.models.team import Team
from posthog.models.utils import CreatedMetaFields, DeletedMetaFields, UUIDModel
from posthog.warehouse.models.util import remove_named_tuples
from posthog.warehouse.models.util import (
remove_named_tuples,
CLICKHOUSE_HOGQL_MAPPING,
clean_type,
STR_TO_HOGQL_MAPPING,
)
from posthog.schema import HogQLQueryModifiers


def validate_saved_query_name(value):
Expand Down Expand Up @@ -52,13 +58,37 @@ class Meta:
)
]

def get_columns(self) -> dict[str, str]:
def get_columns(self) -> dict[str, dict[str, Any]]:
from posthog.api.services.query import process_query_dict
from posthog.hogql_queries.query_runner import ExecutionMode

response = process_query_dict(self.team, self.query, execution_mode=ExecutionMode.CALCULATE_BLOCKING_ALWAYS)
types = getattr(response, "types", {})
return dict(types)
result = getattr(response, "types", [])

if result is None or isinstance(result, int):
raise Exception("No columns types provided by clickhouse in get_columns")

columns = {
str(item[0]): {
"hogql": CLICKHOUSE_HOGQL_MAPPING[clean_type(str(item[1]))].__name__,
"clickhouse": item[1],
"valid": True,
}
for item in result
}

return columns

def get_clickhouse_column_type(self, column_name: str) -> Optional[str]:
clickhouse_type = self.columns.get(column_name, None)

if isinstance(clickhouse_type, dict) and self.columns[column_name].get("clickhouse"):
clickhouse_type = self.columns[column_name].get("clickhouse")

if clickhouse_type.startswith("Nullable("):
clickhouse_type = clickhouse_type.replace("Nullable(", "")[:-1]

return clickhouse_type

@property
def s3_tables(self):
Expand All @@ -83,26 +113,43 @@ def s3_tables(self):

return list(table_collector.tables)

def hogql_definition(self) -> SavedQuery:
def hogql_definition(self, modifiers: Optional[HogQLQueryModifiers] = None) -> SavedQuery:
from posthog.warehouse.models.table import CLICKHOUSE_HOGQL_MAPPING

columns = self.columns or {}

fields = {}
fields: dict[str, FieldOrTable] = {}
structure = []
for column, type in columns.items():
if type.startswith("Nullable("):
type = type.replace("Nullable(", "")[:-1]
# Support for 'old' style columns
if isinstance(type, str):
clickhouse_type = type
else:
clickhouse_type = type["clickhouse"]

if clickhouse_type.startswith("Nullable("):
clickhouse_type = clickhouse_type.replace("Nullable(", "")[:-1]

# TODO: remove when addressed https://github.com/ClickHouse/ClickHouse/issues/37594
if type.startswith("Array("):
type = remove_named_tuples(type)

type = type.partition("(")[0]
try:
type = CLICKHOUSE_HOGQL_MAPPING[type]
fields[column] = type(name=column)
except KeyError as e:
capture_exception(e)
if clickhouse_type.startswith("Array("):
clickhouse_type = remove_named_tuples(clickhouse_type)

if isinstance(type, dict):
column_invalid = not type.get("valid", True)
else:
column_invalid = False

if not column_invalid or (modifiers is not None and modifiers.s3TableUseInvalidColumns):
structure.append(f"`{column}` {clickhouse_type}")

# Support for 'old' style columns
if isinstance(type, str):
hogql_type_str = clickhouse_type.partition("(")[0]
hogql_type = CLICKHOUSE_HOGQL_MAPPING[hogql_type_str]
else:
hogql_type = STR_TO_HOGQL_MAPPING[type["hogql"]]

fields[column] = hogql_type(name=column)

return SavedQuery(
name=self.name,
Expand Down
59 changes: 1 addition & 58 deletions posthog/warehouse/models/table.py
Original file line number Diff line number Diff line change
@@ -1,20 +1,11 @@
import re
from typing import Optional, TypeAlias
from django.db import models

from posthog.client import sync_execute
from posthog.errors import wrap_query_error
from posthog.hogql import ast
from posthog.hogql.database.models import (
BooleanDatabaseField,
DateDatabaseField,
DateTimeDatabaseField,
FieldOrTable,
FloatDatabaseField,
IntegerDatabaseField,
StringArrayDatabaseField,
StringDatabaseField,
StringJSONDatabaseField,
)
from posthog.hogql.database.s3_table import S3Table
from posthog.models.team import Team
Expand All @@ -32,6 +23,7 @@
from uuid import UUID
from sentry_sdk import capture_exception
from posthog.warehouse.util import database_sync_to_async
from posthog.warehouse.models.util import CLICKHOUSE_HOGQL_MAPPING, clean_type, STR_TO_HOGQL_MAPPING
from .external_table_definitions import external_tables

SERIALIZED_FIELD_TO_CLICKHOUSE_MAPPING: dict[DatabaseSerializedFieldType, str] = {
Expand All @@ -45,44 +37,6 @@
DatabaseSerializedFieldType.JSON: "Map",
}

CLICKHOUSE_HOGQL_MAPPING = {
"UUID": StringDatabaseField,
"String": StringDatabaseField,
"DateTime64": DateTimeDatabaseField,
"DateTime32": DateTimeDatabaseField,
"DateTime": DateTimeDatabaseField,
"Date": DateDatabaseField,
"Date32": DateDatabaseField,
"UInt8": IntegerDatabaseField,
"UInt16": IntegerDatabaseField,
"UInt32": IntegerDatabaseField,
"UInt64": IntegerDatabaseField,
"Float8": FloatDatabaseField,
"Float16": FloatDatabaseField,
"Float32": FloatDatabaseField,
"Float64": FloatDatabaseField,
"Int8": IntegerDatabaseField,
"Int16": IntegerDatabaseField,
"Int32": IntegerDatabaseField,
"Int64": IntegerDatabaseField,
"Tuple": StringJSONDatabaseField,
"Array": StringArrayDatabaseField,
"Map": StringJSONDatabaseField,
"Bool": BooleanDatabaseField,
"Decimal": FloatDatabaseField,
}

STR_TO_HOGQL_MAPPING = {
"BooleanDatabaseField": BooleanDatabaseField,
"DateDatabaseField": DateDatabaseField,
"DateTimeDatabaseField": DateTimeDatabaseField,
"IntegerDatabaseField": IntegerDatabaseField,
"FloatDatabaseField": FloatDatabaseField,
"StringArrayDatabaseField": StringArrayDatabaseField,
"StringDatabaseField": StringDatabaseField,
"StringJSONDatabaseField": StringJSONDatabaseField,
}

ExtractErrors = {
"The AWS Access Key Id you provided does not exist": "The Access Key you provided does not exist",
"Access Denied: while reading key:": "Access was denied when reading the provided file",
Expand Down Expand Up @@ -183,17 +137,6 @@ def get_columns(self, safe_expose_ch_error=True) -> DataWarehouseTableColumns:
if result is None or isinstance(result, int):
raise Exception("No columns types provided by clickhouse in get_columns")

def clean_type(column_type: str) -> str:
if column_type.startswith("Nullable("):
column_type = column_type.replace("Nullable(", "")[:-1]

if column_type.startswith("Array("):
column_type = remove_named_tuples(column_type)

column_type = re.sub(r"\(.+\)+", "", column_type)

return column_type

columns = {
str(item[0]): {
"hogql": CLICKHOUSE_HOGQL_MAPPING[clean_type(str(item[1]))].__name__,
Expand Down
Loading

0 comments on commit eedbb81

Please sign in to comment.