init. improve PK as primaryKey, required and unique (when count eq. 1)

datacontract/imports/odcs_helper.py

@@ -8,6 +8,7 @@
     DataQuality,
     EnumValue,
     OpenDataContractStandard,
+    Relationship,
     Role,
     SchemaObject,
     SchemaProperty,
@@ -96,6 +97,7 @@ def create_property(
     map_value: "SchemaProperty" = None,
     dimensions: int = None,
     element_type: str = None,
+    relationships: List[Relationship] = None,
 ) -> SchemaProperty:
     """Create a SchemaProperty (equivalent to DCS Field).
 
@@ -152,6 +154,8 @@ def create_property(
         logical_type_options["elementType"] = element_type
     if logical_type_options:
         prop.logicalTypeOptions = logical_type_options
+    if relationships:
+        prop.relationships = relationships
 
     # precision/scale are forbidden in logicalTypeOptions for number types per ODCS v3.1.0,
     # so carry them in customProperties instead.

datacontract/imports/sql_importer.py

@@ -2,9 +2,10 @@
 import os
 import re
 from enum import Enum
+from typing import List
 
 import sqlglot
-from open_data_contract_standard.model import OpenDataContractStandard, SchemaProperty
+from open_data_contract_standard.model import OpenDataContractStandard, Relationship, SchemaProperty
 from sqlglot.dialects.dialect import Dialects
 
 from datacontract.imports.importer import Importer
@@ -80,16 +81,22 @@ def import_sql(source: str, import_args: dict = None) -> OpenDataContractStandar
         table_name = table.this.name
         properties = []
 
+        columns = list(create.find_all(sqlglot.exp.ColumnDef))
+        primary_key_columns = [column for column in columns if get_primary_key(column, create)]
+        has_single_primary_key = len(primary_key_columns) == 1
+
         primary_key_position = 1
-        for column in create.find_all(sqlglot.exp.ColumnDef):
+        for column in columns:
             col_name = column.this.name
             col_type = to_col_type(column, dialect)
             logical_type, format = map_type_from_sql(col_type)
             col_description = get_description(column)
             max_length = get_max_length(column)
             precision, scale = get_precision_scale(column)
-            is_primary_key = get_primary_key(column)
-            is_required = column.find(sqlglot.exp.NotNullColumnConstraint) is not None or None
+            is_primary_key = get_primary_key(column, create)
+            is_required = column.find(sqlglot.exp.NotNullColumnConstraint) is not None or is_primary_key or None
+            is_unique = True if is_primary_key and has_single_primary_key else None
+            col_relationship = get_relationship(column, create)
             tags = get_tags(column)
 
             map_key, map_value = map_key_value_from_type(col_type) if logical_type == "map" else (None, None)
@@ -107,11 +114,13 @@ def import_sql(source: str, import_args: dict = None) -> OpenDataContractStandar
                 primary_key=is_primary_key,
                 primary_key_position=primary_key_position if is_primary_key else None,
                 required=is_required if is_required else None,
+                unique=is_unique,
                 tags=tags,
                 map_key=map_key,
                 map_value=map_value,
                 dimensions=dimensions,
                 element_type=element_type,
+                relationships=col_relationship,
             )
 
             if is_primary_key:
@@ -144,14 +153,37 @@ def import_sql(source: str, import_args: dict = None) -> OpenDataContractStandar
     return odcs
 
 
-def get_primary_key(column) -> bool | None:
+def get_primary_key(column, table) -> bool | None:
     if column.find(sqlglot.exp.PrimaryKeyColumnConstraint) is not None:
         return True
     if column.find(sqlglot.exp.PrimaryKey) is not None:
         return True
+    if table.find(sqlglot.exp.PrimaryKey) is not None and column.name in [
+        c.name for c in table.find(sqlglot.exp.PrimaryKey).expressions
+    ]:
+        return True
     return None
 
 
+def get_relationship(column, table) -> List[Relationship] | None:
+    reference = column.find(sqlglot.exp.Reference)
+    if reference is None:
+        for foreign_key in table.find_all(sqlglot.exp.ForeignKey):
+            if column.name in [c.name for c in foreign_key.expressions]:
+                reference = foreign_key.args.get("reference")
+                break
+    if reference is None:
+        return None
+
+    referenced_table = reference.this.find(sqlglot.exp.Table)
+    referenced_columns = reference.this.expressions
+    if referenced_table is None or not referenced_columns:
+        return None
+
+    to = f"schema/{referenced_table.this.name}/properties/{referenced_columns[0].name}"
+    return [Relationship(to=to)]
+
+
 def to_dialect(import_args: dict) -> Dialects | None:
     if import_args is None:
         return None

tests/fixtures/snowflake/import/datacontract.yaml

@@ -28,6 +28,20 @@ schema:
       - property: scale
         value: 0
     required: true
+    primaryKey: true
+    primaryKeyPosition: 1
+    unique: true
+  - name: field_parent_primary_key
+    physicalType: DECIMAL(38, 0)
+    description: Parent primary key
+    logicalType: number
+    customProperties:
+      - property: precision
+        value: 38
+      - property: scale
+        value: 0
+    relationships:
+      - to: schema/my_table/properties/field_primary_key
   - name: field_not_null
     physicalType: INT
     description: Not null

tests/fixtures/snowflake/import/ddl.sql

@@ -1,6 +1,7 @@
 CREATE TABLE IF NOT EXISTS  ${database_name}.PUBLIC.my_table (
   -- https://docs.snowflake.com/en/sql-reference/intro-summary-data-types
   field_primary_key      NUMBER(38,0) NOT NULL autoincrement start 1 increment 1 COMMENT 'Primary key',
+  field_parent_primary_key   NUMBER(38,0) COMMENT 'Parent primary key',
   field_not_null         INT NOT NULL     COMMENT 'Not null',
   field_char             CHAR(10)         COMMENT 'Fixed-length string',
   field_character        CHARACTER(10)    COMMENT 'Fixed-length string',
@@ -38,5 +39,6 @@ CREATE TABLE IF NOT EXISTS  ${database_name}.PUBLIC.my_table (
   field_variant          VARIANT          COMMENT 'VARIANT data',
   field_json             OBJECT           COMMENT 'JSON (Stored as text)',
   UNIQUE(field_not_null),
-  PRIMARY KEY (field_primary_key)
+  PRIMARY KEY (field_primary_key),
+  FOREIGN KEY (field_parent_primary_key) REFERENCES ${database_name}.PUBLIC.my_table(field_primary_key)
 ) COMMENT = 'My Comment'

tests/test_import_sql_oracle.py

@@ -51,6 +51,8 @@ def test_import_sql_oracle():
         physicalType: INT
         primaryKey: true
         primaryKeyPosition: 1
+        required: true
+        unique: true
         description: Primary key
       - name: field_not_null
         logicalType: integer
@@ -177,7 +179,10 @@ def test_import_sql_constraints():
       - name: id
         logicalType: number
         physicalType: DECIMAL
+        primaryKey: true
+        primaryKeyPosition: 1
         required: true
+        unique: true
       - name: created_by
         logicalType: string
         logicalTypeOptions:

tests/test_import_sql_postgres.py

@@ -49,6 +49,8 @@ def test_import_sql_postgres():
         physicalType: VARCHAR(10)
         primaryKey: true
         primaryKeyPosition: 1
+        required: true
+        unique: true
       - name: field_two
         logicalType: integer
         physicalType: INT
@@ -87,7 +89,10 @@ def test_import_sql_constraints():
       - name: id
         logicalType: number
         physicalType: DECIMAL
+        primaryKey: true
+        primaryKeyPosition: 1
         required: true
+        unique: true
       - name: created_by
         logicalType: string
         logicalTypeOptions:

tests/test_import_sql_snowflake.py

@@ -39,6 +39,20 @@ def test_import_sql_snowflake():
       value: 0
     logicalType: number
     required: true
+    primaryKey: true
+    primaryKeyPosition: 1
+    unique: true
+  - name: field_parent_primary_key
+    physicalType: DECIMAL(38, 0)
+    description: Parent primary key
+    customProperties:
+    - property: precision
+      value: 38
+    - property: scale
+      value: 0
+    logicalType: number
+    relationships:
+    - to: schema/my_table/properties/field_primary_key
   - name: field_not_null
     physicalType: INT
     description: Not null

tests/test_import_sql_sqlserver.py

@@ -33,6 +33,8 @@ def test_import_sql_sqlserver():
         physicalType: INTEGER
         primaryKey: true
         primaryKeyPosition: 1
+        required: true
+        unique: true
         description: Primary key
       - name: field_not_null
         logicalType: integer