ogrpgtablelayer.cpp

来自「支持各种栅格图像和矢量图像读取的库」· C++ 代码 · 共 1,798 行 · 第 1/4 页

    hResult = PQexec(hPGConn, "BEGIN");
    OGRPGClearResult( hResult );

    osCommand.Printf( "ALTER TABLE %s ADD COLUMN \"%s\" %s",
                      pszSqlTableName, oField.GetNameRef(), szFieldType );
    hResult = PQexec(hPGConn, osCommand);
    if( PQresultStatus(hResult) != PGRES_COMMAND_OK )
    {
        CPLError( CE_Failure, CPLE_AppDefined,
                  "%s\n%s", 
                  osCommand.c_str(), 
                  PQerrorMessage(hPGConn) );

        OGRPGClearResult( hResult );

        hResult = PQexec( hPGConn, "ROLLBACK" );
        OGRPGClearResult( hResult );

        return OGRERR_FAILURE;
    }

    OGRPGClearResult( hResult );

    hResult = PQexec(hPGConn, "COMMIT");
    OGRPGClearResult( hResult );

    poFeatureDefn->AddFieldDefn( &oField );

    return OGRERR_NONE;
}

/************************************************************************/
/*                             GetFeature()                             */
/************************************************************************/

OGRFeature *OGRPGTableLayer::GetFeature( long nFeatureId )

{
    if( pszFIDColumn == NULL )
        return OGRLayer::GetFeature( nFeatureId );

/* -------------------------------------------------------------------- */
/*      Discard any existing resultset.                                 */
/* -------------------------------------------------------------------- */
    ResetReading();

/* -------------------------------------------------------------------- */
/*      Issue query for a single record.                                */
/* -------------------------------------------------------------------- */
    OGRFeature  *poFeature = NULL;
    PGresult    *hResult = NULL;
    PGconn      *hPGConn = poDS->GetPGConn();
    char        *pszFieldList = BuildFields();
    char        *pszCommand = (char *) CPLMalloc(strlen(pszFieldList)+2000);

    poDS->FlushSoftTransaction();
    poDS->SoftStartTransaction();

    sprintf( pszCommand,
             "DECLARE getfeaturecursor CURSOR for "
             "SELECT %s FROM %s WHERE %s = %ld",
             pszFieldList, pszSqlTableName, pszFIDColumn,
             nFeatureId );
    CPLFree( pszFieldList );

    hResult = PQexec(hPGConn, pszCommand );
    CPLFree( pszCommand );

    if( hResult && PQresultStatus(hResult) == PGRES_COMMAND_OK )
    {
        OGRPGClearResult( hResult );

        hResult = PQexec(hPGConn, "FETCH ALL in getfeaturecursor" );

        if( hResult && PQresultStatus(hResult) == PGRES_TUPLES_OK )
        {
            hCursorResult = hResult;
            poFeature = RecordToFeature( 0 );
            hCursorResult = NULL;
        }
    }

/* -------------------------------------------------------------------- */
/*      Cleanup                                                         */
/* -------------------------------------------------------------------- */
    OGRPGClearResult( hResult );

    hResult = PQexec(hPGConn, "CLOSE getfeaturecursor");
    OGRPGClearResult( hResult );

    poDS->FlushSoftTransaction();

    return poFeature;
}

/************************************************************************/
/*                          GetFeatureCount()                           */
/*                                                                      */
/*      If a spatial filter is in effect, we turn control over to       */
/*      the generic counter.  Otherwise we return the total count.      */
/*      Eventually we should consider implementing a more efficient     */
/*      way of counting features matching a spatial query.              */
/************************************************************************/

int OGRPGTableLayer::GetFeatureCount( int bForce )

{
/* -------------------------------------------------------------------- */
/*      Use a more brute force mechanism if we have a spatial query     */
/*      in play.                                                        */
/* -------------------------------------------------------------------- */
    if( m_poFilterGeom != NULL && !bHasPostGISGeometry )
        return OGRPGLayer::GetFeatureCount( bForce );

/* -------------------------------------------------------------------- */
/*      In theory it might be wise to cache this result, but it         */
/*      won't be trivial to work out the lifetime of the value.         */
/*      After all someone else could be adding records from another     */
/*      application when working against a database.                    */
/* -------------------------------------------------------------------- */
    PGconn              *hPGConn = poDS->GetPGConn();
    PGresult            *hResult = NULL;
    CPLString           osCommand;
    int                 nCount = 0;

    poDS->FlushSoftTransaction();
    hResult = PQexec(hPGConn, "BEGIN");
    OGRPGClearResult( hResult );

    osCommand.Printf(
        "DECLARE countCursor CURSOR for "
        "SELECT count(*) FROM %s "
        "%s",
        pszSqlTableName, osWHERE.c_str() );

    CPLDebug( "OGR_PG", "PQexec(%s)\n",
              osCommand.c_str() );

    hResult = PQexec(hPGConn, osCommand);
    OGRPGClearResult( hResult );

    hResult = PQexec(hPGConn, "FETCH ALL in countCursor");
    if( hResult != NULL && PQresultStatus(hResult) == PGRES_TUPLES_OK )
        nCount = atoi(PQgetvalue(hResult,0,0));
    else
        CPLDebug( "OGR_PG", "%s; failed.", osCommand.c_str() );
    OGRPGClearResult( hResult );

    hResult = PQexec(hPGConn, "CLOSE countCursor");
    OGRPGClearResult( hResult );

    hResult = PQexec(hPGConn, "COMMIT");
    OGRPGClearResult( hResult );

    return nCount;
}

/************************************************************************/
/*                           GetSpatialRef()                            */
/*                                                                      */
/*      We override this to try and fetch the table SRID from the       */
/*      geometry_columns table if the srsid is -2 (meaning we           */
/*      haven't yet even looked for it).                                */
/************************************************************************/

OGRSpatialReference *OGRPGTableLayer::GetSpatialRef()

{
    if( nSRSId == -2 )
    {
        PGconn      *hPGConn = poDS->GetPGConn();
        PGresult    *hResult = NULL;
        char        szCommand[1024];

        nSRSId = -1;

        poDS->SoftStartTransaction();

        sprintf( szCommand,
                 "SELECT srid FROM geometry_columns "
                 "WHERE f_table_name = '%s' AND f_table_schema = '%s'",
                 pszTableName, pszSchemaName );
        hResult = PQexec(hPGConn, szCommand );

        if( hResult
            && PQresultStatus(hResult) == PGRES_TUPLES_OK
            && PQntuples(hResult) == 1 )
        {
            nSRSId = atoi(PQgetvalue(hResult,0,0));
        }
        else // I think perhaps an older version used f_schema_name.
        {
            OGRPGClearResult( hResult );

            poDS->SoftCommit();
            poDS->SoftStartTransaction();

            sprintf( szCommand,
                     "SELECT srid FROM geometry_columns "
                     "WHERE f_table_name = '%s' AND f_schema_name = '%s'",
                     pszTableName, pszSchemaName );
            hResult = PQexec(hPGConn, szCommand );
            if( hResult
                && PQresultStatus(hResult) == PGRES_TUPLES_OK
                && PQntuples(hResult) == 1 )
            {
                nSRSId = atoi(PQgetvalue(hResult,0,0));
            }
        }

        OGRPGClearResult( hResult );

        poDS->SoftCommit();
    }

    return OGRPGLayer::GetSpatialRef();
}

/************************************************************************/
/*                             GetExtent()                              */
/*                                                                      */
/*      For PostGIS use internal Extend(geometry) function              */
/*      in other cases we use standard OGRLayer::GetExtent()            */
/************************************************************************/

OGRErr OGRPGTableLayer::GetExtent( OGREnvelope *psExtent, int bForce )
{
    if ( psExtent == NULL )
        return OGRERR_FAILURE;

    if ( bHasPostGISGeometry )
    {
        PGconn      *hPGConn = poDS->GetPGConn();
        PGresult    *hResult = NULL;
        CPLString   osCommand;

        osCommand.Printf( "SELECT Extent(\"%s\") FROM %s", 
                          pszGeomColumn, pszSqlTableName );

        hResult = PQexec( hPGConn, osCommand );
        if( ! hResult || PQresultStatus(hResult) != PGRES_TUPLES_OK || PQgetisnull(hResult,0,0) )
        {
            OGRPGClearResult( hResult );
            CPLDebug("OGR_PG","Unable to get extent by PostGIS. Using standard OGRLayer method.");
            return OGRPGLayer::GetExtent( psExtent, bForce );
        }

        char * pszBox = PQgetvalue(hResult,0,0);
        char * ptr = pszBox;
        char szVals[64*6+6];

        while ( *ptr != '(' && ptr ) ptr++; ptr++;

        strncpy(szVals,ptr,strstr(ptr,")") - ptr);
        szVals[strstr(ptr,")") - ptr] = '\0';

        char ** papszTokens = CSLTokenizeString2(szVals," ,",CSLT_HONOURSTRINGS);
        int nTokenCnt = poDS->sPostGISVersion.nMajor >= 1 ? 4 : 6;

        if ( CSLCount(papszTokens) != nTokenCnt )
        {
            CPLError( CE_Failure, CPLE_IllegalArg,
                      "Bad extent representation: '%s'", pszBox);
            CSLDestroy(papszTokens);

            OGRPGClearResult( hResult );
            return OGRERR_FAILURE;
        }

        // Take X,Y coords
        // For PostGis ver >= 1.0.0 -> Tokens: X1 Y1 X2 Y2 (nTokenCnt = 4)
        // For PostGIS ver < 1.0.0 -> Tokens: X1 Y1 Z1 X2 Y2 Z2 (nTokenCnt = 6)
        // =>   X2 index calculated as nTokenCnt/2
        //      Y2 index caluclated as nTokenCnt/2+1
        
        psExtent->MinX = CPLScanDouble(papszTokens[0],strlen(papszTokens[0]),"C");
        psExtent->MinY = CPLScanDouble(papszTokens[1],strlen(papszTokens[1]),"C");
        psExtent->MaxX = CPLScanDouble(papszTokens[nTokenCnt/2],strlen(papszTokens[nTokenCnt/2]),"C");
        psExtent->MaxY = CPLScanDouble(papszTokens[nTokenCnt/2+1],strlen(papszTokens[nTokenCnt/2+1]),"C");

        CSLDestroy(papszTokens);
        OGRPGClearResult( hResult );

        return OGRERR_NONE;
    }

    return OGRLayer::GetExtent( psExtent, bForce );
}

/************************************************************************/
/*                             StartCopy()                              */
/************************************************************************/

OGRErr OGRPGTableLayer::StartCopy()

{
    OGRErr result = OGRERR_NONE;

    /* Tell the datasource we are now planning to copy data */
    poDS->StartCopy( this ); 

    char *pszFields = BuildCopyFields();

    int size = strlen(pszFields) +  strlen(pszSqlTableName) + 100;
    char *pszCommand = (char *) CPLMalloc(size);

    sprintf( pszCommand,
             "COPY %s (%s) FROM STDIN;",
             pszSqlTableName, pszFields );

    CPLFree( pszFields );

    PGconn *hPGConn = poDS->GetPGConn();
    CPLDebug( "OGR_PG", "%s", pszCommand );
    PGresult *hResult = PQexec(hPGConn, pszCommand);

    if ( !hResult || (PQresultStatus(hResult) != PGRES_COPY_IN))
    {
        CPLError( CE_Failure, CPLE_AppDefined,
                  "%s", PQerrorMessage(hPGConn) );
        result = OGRERR_FAILURE;
    }
    else
        bCopyActive = TRUE;

    OGRPGClearResult( hResult );
    CPLFree( pszCommand );

    return OGRERR_NONE;
}

/************************************************************************/
/*                              EndCopy()                               */
/************************************************************************/

OGRErr OGRPGTableLayer::EndCopy()

{
    if( !bCopyActive )
        return OGRERR_NONE;

    /* This method is called from the datasource when
       a COPY operation is ended */
    OGRErr result = OGRERR_NONE;

    PGconn *hPGConn = poDS->GetPGConn();
    CPLDebug( "OGR_PG", "PQputCopyEnd()" );

    bCopyActive = FALSE;

    /* This is for postgresql 7.4 and higher */
#if !defined(PG_PRE74)
    int copyResult = PQputCopyEnd(hPGConn, NULL);

    switch (copyResult)
    {
      case 0:
        CPLError( CE_Failure, CPLE_AppDefined, "Writing COPY data blocked.");
        result = OGRERR_FAILURE;
        break;
      case -1:
        CPLError( CE_Failure, CPLE_AppDefined, "%s", PQerrorMessage(hPGConn) );
        result = OGRERR_FAILURE;
        break;
    }

#else /* defined(PG_PRE74) */
    PQputline(hPGConn, "\\.\n");
    int copyResult = PQendcopy(hPGConn);

    if (copyResult != 0)
    {
      CPLError( CE_Failure, CPLE_AppDefined, "%s", PQerrorMessage(hPGConn) );
      result = OGRERR_FAILURE;
    }
#endif /* defined(PG_PRE74) */

    /* Now check the results of the copy */
    PGresult * hResult = PQgetResult( hPGConn );

    if( hResult && PQresultStatus(hResult) != PGRES_COMMAND_OK )
    {
        CPLError( CE_Failure, CPLE_AppDefined,
                  "COPY statement failed.\n%s",
                  PQerrorMessage(hPGConn) );

        result = OGRERR_FAILURE;
    }

    OGRPGClearResult( hResult );

    bUseCopy = USE_COPY_UNSET;

    return result;
}

/************************************************************************/
/*                          BuildCopyFields()                           */
/************************************************************************/

char *OGRPGTableLayer::BuildCopyFields()
{
    int     i = 0;
    int     nSize = 0;
    char    *pszFieldList;
        
    nSize = 25;
    if( pszGeomColumn )
        nSize += strlen(pszGeomColumn);

    if( bHasFid && poFeatureDefn->GetFieldIndex( pszFIDColumn ) != -1 )
        nSize += strlen(pszFIDColumn);

    for( i = 0; i < poFeatureDefn->GetFieldCount(); i++ )
        nSize += strlen(poFeatureDefn->GetFieldDefn(i)->GetNameRef()) + 4;

    pszFieldList = (char *) CPLMalloc(nSize);
    pszFieldList[0] = '\0';

    if( bHasFid && poFeatureDefn->GetFieldIndex( pszFIDColumn ) != -1 )
        sprintf( pszFieldList, "\"%s\"", pszFIDColumn );

    if( pszGeomColumn )
    {
        if( strlen(pszFieldList) > 0 )
            strcat( pszFieldList, ", " );

        sprintf( pszFieldList+strlen(pszFieldList),
                 "\"%s\"", pszGeomColumn );
    }

    for( i = 0; i < poFeatureDefn->GetFieldCount(); i++ )
    {
        const char *pszName = poFeatureDefn->GetFieldDefn(i)->GetNameRef();

        if( strlen(pszFieldList) > 0 )
            strcat( pszFieldList, ", " );

        strcat( pszFieldList, "\"" );
        strcat( pszFieldList, pszName );
        strcat( pszFieldList, "\"" );
    }

    CPLAssert( (int) strlen(pszFieldList) < nSize );

    return pszFieldList;
}